futhark 0.22.4 → 0.22.5

raw patch · 54 files changed

+3013/−878 lines, 54 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• docs/c-api.rst +5/−0
• docs/error-index.rst +120/−26
• docs/installation.rst +17/−13
• futhark.cabal +2/−1
• rts/c/cuda.h +4/−0
• rts/c/opencl.h +6/−0
• rts/c/server.h +3/−3
• rts/python/opencl.py +2/−0
• src/Futhark/AD/Rev/Hist.hs +889/−0
• src/Futhark/AD/Rev/Reduce.hs +139/−0
• src/Futhark/AD/Rev/SOAC.hs +140/−8
• src/Futhark/AD/Rev/Scan.hs +275/−80
• src/Futhark/Actions.hs +28/−8
• src/Futhark/Analysis/Interference.hs +1/−1
• src/Futhark/Analysis/LastUse.hs +6/−6
• src/Futhark/Analysis/PrimExp.hs +6/−4
• src/Futhark/CLI/Dev.hs +36/−1
• src/Futhark/CLI/Test.hs +7/−1
• src/Futhark/CodeGen/Backends/CCUDA.hs +24/−11
• src/Futhark/CodeGen/Backends/COpenCL.hs +23/−8
• src/Futhark/CodeGen/Backends/GenericC/Pretty.hs +5/−1
• src/Futhark/CodeGen/Backends/MulticoreISPC.hs +22/−8
• src/Futhark/CodeGen/Backends/PyOpenCL.hs +14/−9
• src/Futhark/CodeGen/ImpCode/GPU.hs +19/−5
• src/Futhark/CodeGen/ImpCode/OpenCL.hs +4/−1
• src/Futhark/CodeGen/ImpGen/GPU/Base.hs +14/−1
• src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs +26/−14
• src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs +39/−26
• src/Futhark/CodeGen/ImpGen/GPU/Transpose.hs +1/−1
• src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs +144/−106
• src/Futhark/IR/Mem/IxFun.hs +6/−3
• src/Futhark/IR/Pretty.hs +1/−1
• src/Futhark/IR/Prop/Names.hs +3/−0
• src/Futhark/Internalise/Defunctionalise.hs +30/−38
• src/Futhark/Internalise/Exps.hs +9/−9
• src/Futhark/Optimise/ArrayShortCircuiting.hs +97/−55
• src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs +400/−205
• src/Futhark/Optimise/ArrayShortCircuiting/DataStructs.hs +13/−36
• src/Futhark/Optimise/ArrayShortCircuiting/LastUse.hs +205/−94
• src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs +1/−1
• src/Futhark/Optimise/ArrayShortCircuiting/TopdownAnalysis.hs +28/−6
• src/Futhark/Pass/LiftAllocations.hs +40/−10
• src/Futhark/Pass/LowerAllocations.hs +40/−10
• src/Futhark/Passes.hs +10/−0
• src/Futhark/Test/Spec.hs +4/−0
• src/Language/Futhark/Interpreter/Values.hs +4/−4
• src/Language/Futhark/Primitive.hs +14/−0
• src/Language/Futhark/Primitive/Parse.hs +1/−1
• src/Language/Futhark/Prop.hs +7/−4
• src/Language/Futhark/TypeChecker.hs +12/−3
• src/Language/Futhark/TypeChecker/Terms.hs +10/−9
• src/Language/Futhark/TypeChecker/Terms/DoLoop.hs +1/−0
• src/Language/Futhark/TypeChecker/Terms/Pat.hs +7/−1
• src/Language/Futhark/TypeChecker/Unify.hs +49/−45

@@ -8,6 +8,11 @@ ``futlib.h``.  The API provided in the ``.h`` file is documented in the following. +The ``.h`` file can be included by a C++ source file to access the+functions (``extern "C"`` is added automatically), but the ``.c`` file+must be compiled with a proper C compiler and the resulting object+file linked with the rest of the program.+ Using the API requires creating a *configuration object*, which is then used to obtain a *context object*, which is then used to perform most other operations, such as calling Futhark functions.

@@ -582,32 +582,6 @@ Module errors ------------- -.. _nested-entry:--"Entry points may not be declared inside modules."-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~--This occurs when the program uses the ``entry`` keyword inside a-module:--.. code-block:: futhark--  module m = {-    entry f x = x + 1-  }--Entry points can only be declared at the top level of a file.  When we-wish to make a function from inside a module available as an entry-point, we must define a wrapper function:--.. code-block:: futhark--  module m = {-    def f x = x + 1-  }--  entry f = m.f- .. _module-is-parametric:  "Module *x* is a parametric module@@ -767,3 +741,123 @@ .. code-block:: futhark    def f (r : {x:i32}) = r with x = 0++Entry points+------------++.. _nested-entry:++"Entry points may not be declared inside modules."+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~++This occurs when the program uses the ``entry`` keyword inside a+module:++.. code-block:: futhark++  module m = {+    entry f x = x + 1+  }++Entry points can only be declared at the top level of a file.  When we+wish to make a function from inside a module available as an entry+point, we must define a wrapper function:++.. code-block:: futhark++  module m = {+    def f x = x + 1+  }++  entry f = m.f++.. _polymorphic-entry:++"Entry point functions may not be polymorphic."+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~++Entry points are Futhark functions that can be called from other+languages, and are therefore limited how advanced their types can be.+In this case, the problem is that an entry point may not have a+polymorphic type, for example:++.. code-block:: futhark++   entry dup 't (x: t) : (t,t) = x++This is an invalid entry point because it uses a type parameter+``'t``.  This error occurs frequently when we want to test a+polymorphic function.  In such cases, the solution is to define one or+more *monomorphic* entry points, each for a distinct type.  For+example, to we can define a variety of monomorphic entry points that+call the built-in function ``scan``:++.. code-block:: futhark++   entry scan_i32 (xs: []i32) = scan (+) 0 xs++   entry scan_f32 (xs: []i32) = scan (*) 1 xs++.. _higher-order-entry:++"Entry point functions may not be higher-order."+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~++Entry points are Futhark functions that can be called from other+languages, and are therefore limited how advanced their types can be.+In this case, the problem is that an entry point may use functions as+input or output.  For example:++.. code-block:: futhark++   entry apply (f: i32 -> i32) (x: i32) = f x++There is no simple workaround for such cases.  One option is to+manually `defunctionalise+<https://en.wikipedia.org/wiki/Defunctionalization>`_ to use a+non-functional encoding of the functional values, but this can quickly+get very elaborate.  Following up on the example above, if we know+that the only functions that would ever be passed are ``(+y)`` or+``(*y)`` for some ``y``, we could do something like the following:++.. code-block:: futhark++   type function = #add i32 | #mul i32++   entry apply (f: function) (x: i32) =+     match f+     case #add y -> x + y+     case #mul y -> x + y++Although in many cases, the best solution is simply to define a+handful of simpler entry points instead of a single complicated one.++.. _size-polymorphic-entry:++"Entry point functions must not be size-polymorphic in their return type."+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~++This somewhat rare error occurs when an entry point returns an array+that can have an arbitrary size chosen by its caller.  Contrived example:++.. code-block:: futhark++   -- Entry point taking no parameters.+   entry f [n] : [0][n]i32 = []++The size ``n`` is chosen by the caller.  Note that the ``n`` might be+inferred and invisible, as in this example:++.. code-block:: futhark++   entry g : [0][]i32 = []++When calling functions within a Futhark program, size parameters are+handled by type inference, but entry points are called from the+outside world, which is not subject to type inference.  If you really+must have entry points like this, turn the size parameter into an+ordinary parameter:++.. code-block:: futhark++   entry f (n: i64) : [0][n]i32 = []

@@ -38,7 +38,7 @@  To compile Futhark you must first install an appropriate version of GHC, either with `ghcup <https://www.haskell.org/ghcup>`_ or a package-manager.  Any version since GHC 8.10 should work.  You also need the+manager.  Any version since GHC 9.0 should work.  You also need the ``cabal`` command line program, which ghcup will install for you as well. @@ -96,6 +96,12 @@ **Linux (x86_64)**   `futhark-nightly-linux-x86_64.tar.xz <https://futhark-lang.org/releases/futhark-nightly-linux-x86_64.tar.xz>`_ +**macOS (x86_64)**+  `futhark-nightly-macos-x86_64.zip <https://futhark-lang.org/releases/futhark-nightly-macos-x86_64.zip>`_++**Windows (x86_64)**+  `futhark-nightly-windows-x86_64.zip <https://futhark-lang.org/releases/futhark-nightly-windows-x86_64.zip>`_+   You will still likely need to make a C compiler (such as GCC) available on your own.  .. _linux-installation:@@ -103,16 +109,13 @@ Installing Futhark on Linux --------------------------- -* `Linuxbrew`_ is a distribution-agnostic package manager that-  contains a formula for Futhark.  If Linuxbrew is installed (which-  does not require ``root`` access), installation is as easy as::+* `Homebrew`_ is a distribution-agnostic package manager for macOS and+  Linux that contains a formula for Futhark. If Homebrew is already+  installed (which does not require ``root`` access), installation is+  as easy as::      $ brew install futhark -  Note that as of this writing, Linuxbrew is hampered by limited-  compute resources for building packages, so the Futhark version may-  be a bit behind.- * Arch Linux users can use a `futhark-nightly package   <https://aur.archlinux.org/packages/futhark-nightly/>`_ or a   `regular futhark package@@ -179,13 +182,14 @@ Setting up Futhark on Windows ----------------------------- -Due to limited maintenance and testing resources, Futhark is not-directly supported on Windows.  Install `WSL+Due to limited maintenance and testing resources, Futhark is only+partially supported on Windows.  A precompiled nightly snapshot is+available above.++In most cases, it is better to install `WSL <https://docs.microsoft.com/en-us/windows/wsl/install>`_ and follow the Linux instructions above.  The C code generated by the Futhark-compiler should work on Windows.--In the future, we may support Windows directly again.+compiler should work on Windows, except for the ``multicore`` backend.  Futhark with Nix ----------------

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.22.4+version:        0.22.5 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -150,6 +150,7 @@       Futhark.AD.Fwd       Futhark.AD.Rev       Futhark.AD.Rev.Loop+      Futhark.AD.Rev.Hist       Futhark.AD.Rev.Map       Futhark.AD.Rev.Monad       Futhark.AD.Rev.Reduce

@@ -349,6 +349,9 @@   } else {     opts[i++] = strdup("--disable-warnings");   }+  opts[i++] = msgprintf("-D%s=%d",+                        "max_group_size",+                        (int)ctx->max_block_size);   for (size_t j = 0; j < ctx->cfg.num_sizes; j++) {     opts[i++] = msgprintf("-D%s=%zu", ctx->cfg.size_vars[j],                           ctx->cfg.size_values[j]);@@ -427,6 +430,7 @@   }   cache_hash(h, src, strlen(src));   size_t ptxsize;+  errno = 0;   if (cache_restore(cache_fname, h, (unsigned char**)ptx, &ptxsize) != 0) {     if (ctx->cfg.logging) {       fprintf(stderr, "Failed to restore cache (errno: %s)\n", strerror(errno));

@@ -517,6 +517,11 @@                    "-DLOCKSTEP_WIDTH=%d ",                    (int)ctx->lockstep_width); +  w += snprintf(compile_opts+w, compile_opts_size-w,+                "-D%s=%d ",+                "max_group_size",+                (int)ctx->max_group_size);+   for (int i = 0; i < ctx->cfg.num_sizes; i++) {     w += snprintf(compile_opts+w, compile_opts_size-w,                   "-D%s=%d ",@@ -742,6 +747,7 @@        unsigned char *buf;       size_t bufsize;+      errno = 0;       if (cache_restore(cache_fname, &h, &buf, &bufsize) != 0) {         if (ctx->cfg.logging) {           fprintf(stderr, "Failed to restore cache (errno: %s)\n", strerror(errno));

@@ -388,6 +388,7 @@         return;       } +      errno = 0;       if (t->restore(t->aux, f, s->ctx, value_ptr(&v->value)) != 0) {         failure();         printf("Failed to restore variable %s.\n"@@ -852,11 +853,10 @@   if (arr == NULL) {     return 1;   }-  assert(futhark_context_sync(ctx) == 0);-+  int err = futhark_context_sync(ctx);   *(void**)p = arr;   free(data);-  return 0;+  return err; }  void store_array(const struct array_aux *aux, FILE *f,

@@ -217,6 +217,8 @@     if (len(program_src) >= 0):         build_options += ["-DLOCKSTEP_WIDTH={}".format(lockstep_width)] +        build_options += ["-D{}={}".format('max_group_size', max_group_size)]+         build_options += ["-D{}={}".format(s.                                            replace('z', 'zz').                                            replace('.', 'zi').

@@ -0,0 +1,889 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeFamilies #-}++module Futhark.AD.Rev.Hist+  ( diffMinMaxHist,+    diffMulHist,+    diffAddHist,+    diffHist,+  )+where++import Control.Monad+import Futhark.AD.Rev.Monad+import Futhark.Analysis.PrimExp.Convert+import Futhark.Builder+import Futhark.IR.SOACS+import Futhark.Tools+import Futhark.Transform.Rename++getBinOpPlus :: PrimType -> BinOp+getBinOpPlus (IntType x) = Add x OverflowUndef+getBinOpPlus (FloatType f) = FAdd f+getBinOpPlus _ = error "In getBinOpMul, Hist.hs: input not supported"++getBinOpDiv :: PrimType -> BinOp+getBinOpDiv (IntType t) = SDiv t Unsafe+getBinOpDiv (FloatType t) = FDiv t+getBinOpDiv _ = error "In getBinOpDiv, Hist.hs: input not supported"++withinBounds :: [(SubExp, VName)] -> TPrimExp Bool VName+withinBounds [] = TPrimExp $ ValueExp (BoolValue True)+withinBounds [(q, i)] = (le64 i .<. pe64 q) .&&. (pe64 (intConst Int64 (-1)) .<. le64 i)+withinBounds (qi : qis) = withinBounds [qi] .&&. withinBounds qis++elseIf :: PrimType -> [(ADM (Exp SOACS), ADM (Exp SOACS))] -> [ADM (Body SOACS)] -> ADM (Exp SOACS)+elseIf t [(c1, c2)] [bt, bf] =+  eIf+    (eCmpOp (CmpEq t) c1 c2)+    bt+    bf+elseIf t ((c1, c2) : cs) (bt : bs) =+  eIf+    (eCmpOp (CmpEq t) c1 c2)+    bt+    $ eBody+    $ pure+    $ elseIf t cs bs+elseIf _ _ _ = error "In elseIf, Hist.hs: input not supported"++bindSubExpRes :: String -> [SubExpRes] -> ADM [VName]+bindSubExpRes s =+  traverse+    ( \(SubExpRes cs se) -> do+        bn <- newVName s+        certifying cs $ letBindNames [bn] $ BasicOp $ SubExp se+        pure bn+    )++nestedmap :: [SubExp] -> [PrimType] -> Lambda SOACS -> ADM (Lambda SOACS)+nestedmap [] _ lam = pure lam+nestedmap s@(h : r) pt lam = do+  params <- traverse (\tp -> newParam "x" $ Array tp (Shape s) NoUniqueness) pt+  body <- nestedmap r pt lam+  mkLambda params $+    fmap varsRes . letTupExp "res" . Op $+      Screma h (map paramName params) (mapSOAC body)++-- \ds hs -> map2 lam ds hs+mkF' :: Lambda SOACS -> [Type] -> SubExp -> ADM ([VName], [VName], Lambda SOACS)+mkF' lam tps n = do+  lam' <- renameLambda lam++  ds_params <- traverse (newParam "ds_param") tps+  hs_params <- traverse (newParam "hs_param") tps+  let ds_pars = fmap paramName ds_params+  let hs_pars = fmap paramName hs_params+  lam_map <-+    mkLambda (ds_params <> hs_params) $+      fmap varsRes . letTupExp "map_f'" . Op $+        Screma n (ds_pars <> hs_pars) (mapSOAC lam')++  pure (ds_pars, hs_pars, lam_map)++-- \ls as rs -> map3 (\li ai ri -> li `lam` ai `lam` ri) ls as rs+mkF :: Lambda SOACS -> [Type] -> SubExp -> ADM ([VName], Lambda SOACS)+mkF lam tps n = do+  lam_l <- renameLambda lam+  lam_r <- renameLambda lam+  let q = length $ lambdaReturnType lam+      (lps, aps) = splitAt q $ lambdaParams lam_l+      (ips, rps) = splitAt q $ lambdaParams lam_r+  lam' <- mkLambda (lps <> aps <> rps) $ do+    lam_l_res <- bodyBind $ lambdaBody lam_l+    forM_ (zip ips lam_l_res) $ \(ip, SubExpRes cs se) ->+      certifying cs $ letBindNames [paramName ip] $ BasicOp $ SubExp se+    bodyBind $ lambdaBody lam_r++  ls_params <- traverse (newParam "ls_param") tps+  as_params <- traverse (newParam "as_param") tps+  rs_params <- traverse (newParam "rs_param") tps+  let map_params = ls_params <> as_params <> rs_params+  lam_map <-+    mkLambda map_params $+      fmap varsRes . letTupExp "map_f" $+        Op $+          Screma n (map paramName map_params) $+            mapSOAC lam'++  pure (map paramName as_params, lam_map)++mapout :: VName -> SubExp -> SubExp -> ADM VName+mapout is n w = do+  par_is <- newParam "is" $ Prim int64+  is'_lam <-+    mkLambda [par_is] $+      fmap varsRes . letTupExp "is'"+        =<< eIf+          (toExp $ withinBounds $ pure (w, paramName par_is))+          (eBody $ pure $ eParam par_is)+          (eBody $ pure $ eSubExp w)++  letExp "is'" $ Op $ Screma n (pure is) $ mapSOAC is'_lam++multiScatter :: SubExp -> [VName] -> VName -> [VName] -> ADM [VName]+multiScatter n dst is vs = do+  tps <- traverse lookupType vs+  par_i <- newParam "i" $ Prim int64+  scatter_params <- traverse (newParam "scatter_param" . rowType) tps+  scatter_lam <-+    mkLambda (par_i : scatter_params) $+      fmap subExpsRes . mapM (letSubExp "scatter_map_res") =<< do+        p1 <- replicateM (length scatter_params) $ eParam par_i+        p2 <- traverse eParam scatter_params+        pure $ p1 <> p2++  letTupExp "scatter_res" . Op $+    Scatter n (is : vs) scatter_lam $+      zipWith (\t -> (,,) (Shape $ pure $ arraySize 0 t) 1) tps dst++multiIndex :: [VName] -> [DimIndex SubExp] -> ADM [VName]+multiIndex vs s = do+  traverse+    ( \x -> do+        t <- lookupType x+        letExp "sorted" $ BasicOp $ Index x (fullSlice t s)+    )+    vs++--+-- special case of histogram with min/max as operator.+-- Original, assuming `is: [n]i64` and `dst: [w]btp`+--     let x = reduce_by_index dst minmax ne is vs+-- Forward sweep:+--     need to copy dst: reverse sweep might use it 7+--       (see ex. in reducebyindexminmax6.fut where the first map requires the original dst to be differentiated).+--     let dst_cpy = copy dst+--     let (x, x_inds) = zip vs (iota n)+--                       |> reduce_by_index (dst_cpy,-1s) argminmax (ne,-1) is+--+-- Reverse sweep:+--     dst_bar += map2 (\i b -> if i == -1+--                              then b+--                              else 0+--                     ) x_inds x_bar++--     vs_ctrbs = map2 (\i b -> if i == -1+--                              then 0+--                              else vs_bar[i] + b+--                     ) x_inds x_bar+--     vs_bar <- scatter vs_bar x_inds vs_ctrbs+diffMinMaxHist ::+  VjpOps -> VName -> StmAux () -> SubExp -> BinOp -> SubExp -> VName -> VName -> SubExp -> SubExp -> VName -> ADM () -> ADM ()+diffMinMaxHist _ops x aux n minmax ne is vs w rf dst m = do+  let t = binOpType minmax+  vs_type <- lookupType vs+  let vs_elm_type = elemType vs_type+  let vs_dims = arrayDims vs_type+  let inner_dims = tail vs_dims+  let nr_dims = length vs_dims+  dst_type <- lookupType dst+  let dst_dims = arrayDims dst_type++  dst_cpy <- letExp (baseString dst <> "_copy") $ BasicOp $ Copy dst++  acc_v_p <- newParam "acc_v" $ Prim t+  acc_i_p <- newParam "acc_i" $ Prim int64+  v_p <- newParam "v" $ Prim t+  i_p <- newParam "i" $ Prim int64+  hist_lam_inner <-+    mkLambda [acc_v_p, acc_i_p, v_p, i_p] $+      fmap varsRes . letTupExp "idx_res"+        =<< eIf+          (eCmpOp (CmpEq t) (eParam acc_v_p) (eParam v_p))+          ( eBody+              [ eParam acc_v_p,+                eBinOp (SMin Int64) (eParam acc_i_p) (eParam i_p)+              ]+          )+          ( eBody+              [ eIf+                  ( eCmpOp+                      (CmpEq t)+                      (eParam acc_v_p)+                      (eBinOp minmax (eParam acc_v_p) (eParam v_p))+                  )+                  (eBody [eParam acc_v_p, eParam acc_i_p])+                  (eBody [eParam v_p, eParam i_p])+              ]+          )+  hist_lam <- nestedmap inner_dims [vs_elm_type, int64, vs_elm_type, int64] hist_lam_inner++  dst_minus_ones <-+    letExp "minus_ones" . BasicOp $+      Replicate (Shape dst_dims) (intConst Int64 (-1))+  ne_minus_ones <-+    letSubExp "minus_ones" . BasicOp $+      Replicate (Shape inner_dims) (intConst Int64 (-1))+  iota_n <-+    letExp "red_iota" . BasicOp $+      Iota n (intConst Int64 0) (intConst Int64 1) Int64++  inp_iota <- do+    if nr_dims == 1+      then pure iota_n+      else do+        i <- newParam "i" $ Prim int64+        lam <-+          mkLambda [i] $+            fmap varsRes . letTupExp "res" =<< do+              pure $ BasicOp $ Replicate (Shape inner_dims) $ Var $ paramName i++        letExp "res" $ Op $ Screma n [iota_n] $ mapSOAC lam++  let hist_op = HistOp (Shape [w]) rf [dst_cpy, dst_minus_ones] [ne, if nr_dims == 1 then intConst Int64 (-1) else ne_minus_ones] hist_lam+  f' <- mkIdentityLambda [Prim int64, rowType vs_type, rowType $ Array int64 (Shape vs_dims) NoUniqueness]+  x_inds <- newVName (baseString x <> "_inds")+  auxing aux $+    letBindNames [x, x_inds] $+      Op $+        Hist n [is, vs, inp_iota] [hist_op] f'++  m++  x_bar <- lookupAdjVal x++  x_ind_dst <- newParam (baseString x <> "_ind_param") $ Prim int64+  x_bar_dst <- newParam (baseString x <> "_bar_param") $ Prim t+  dst_lam_inner <-+    mkLambda [x_ind_dst, x_bar_dst] $+      fmap varsRes . letTupExp "dst_bar"+        =<< eIf+          (toExp $ le64 (paramName x_ind_dst) .==. -1)+          (eBody $ pure $ eParam x_bar_dst)+          (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)+  dst_lam <- nestedmap inner_dims [int64, vs_elm_type] dst_lam_inner++  dst_bar <-+    letExp (baseString dst <> "_bar") . Op $+      Screma w [x_inds, x_bar] (mapSOAC dst_lam)++  updateAdj dst dst_bar++  vs_bar <- lookupAdjVal vs++  inds' <- traverse (letExp "inds" . BasicOp . Replicate (Shape [w]) . Var) =<< mk_indices inner_dims []+  let inds = x_inds : inds'++  par_x_ind_vs <- replicateM nr_dims $ newParam (baseString x <> "_ind_param") $ Prim int64+  par_x_bar_vs <- newParam (baseString x <> "_bar_param") $ Prim t+  vs_lam_inner <-+    mkLambda (par_x_bar_vs : par_x_ind_vs) $+      fmap varsRes . letTupExp "res"+        =<< eIf+          (toExp $ le64 (paramName $ head par_x_ind_vs) .==. -1)+          (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)+          ( eBody $+              pure $ do+                vs_bar_i <-+                  letSubExp (baseString vs_bar <> "_el") . BasicOp $+                    Index vs_bar . Slice $+                      fmap (DimFix . Var . paramName) par_x_ind_vs+                eBinOp (getBinOpPlus t) (eParam par_x_bar_vs) (eSubExp vs_bar_i)+          )+  vs_lam <- nestedmap inner_dims (vs_elm_type : replicate nr_dims int64) vs_lam_inner++  vs_bar_p <-+    letExp (baseString vs <> "_partial") . Op $+      Screma w (x_bar : inds) (mapSOAC vs_lam)++  q <-+    letSubExp "q"+      =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) dst_dims++  scatter_inps <- do+    -- traverse (letExp "flat" . BasicOp . Reshape [DimNew q]) $ inds ++ [vs_bar_p]+    -- ToDo: Cosmin asks: is the below the correct translation of the line above?+    traverse (letExp "flat" . BasicOp . Reshape ReshapeArbitrary (Shape [q])) $+      inds ++ [vs_bar_p]++  f'' <- mkIdentityLambda $ replicate nr_dims (Prim int64) ++ [Prim t]+  vs_bar' <-+    letExp (baseString vs <> "_bar") . Op $+      Scatter q scatter_inps f'' [(Shape vs_dims, 1, vs_bar)]+  insAdj vs vs_bar'+  where+    mk_indices :: [SubExp] -> [SubExp] -> ADM [VName]+    mk_indices [] _ = pure []+    mk_indices [d] iotas = do+      reps <- traverse (letExp "rep" . BasicOp . Replicate (Shape [d])) iotas+      iota_d <-+        letExp "red_iota" . BasicOp $+          Iota d (intConst Int64 0) (intConst Int64 1) Int64+      pure $ reps ++ [iota_d]+    mk_indices (d : dims) iotas = do+      iota_d <-+        letExp "red_iota" . BasicOp $+          Iota d (intConst Int64 0) (intConst Int64 1) Int64++      i_param <- newParam "i" $ Prim int64+      lam <-+        mkLambda [i_param] $+          fmap varsRes $+            mk_indices dims $+              iotas ++ [Var $ paramName i_param]++      letTupExp "res" $ Op $ Screma d [iota_d] $ mapSOAC lam++--+-- special case of histogram with multiplication as operator.+-- Original, assuming `is: [n]i64` and `dst: [w]btp`+--     let x = reduce_by_index dst (*) ne is vs+-- Forward sweep:+--     dst does not need to be copied: dst is not overwritten+--     let (ps, zs) = map (\v -> if v == 0 then (1,1) else (v,0)) vs+--     let non_zero_prod = reduce_by_index nes (*) ne is ps+--     let zero_count = reduce_by_index 0s (+) 0 is zs+--     let h_part = map2 (\p c -> if c == 0 then p else 0+--                       ) non_zero_prod zero_count+--     let x = map2 (*) dst h_part+--+-- Reverse sweep:+--     dst_bar += map2 (*) h_part x_bar++--     let part_bar = map2 (*) dst x_bar+--     vs_bar += map2 (\i v -> let zr_cts = zero_count[i]+--                             let pr_bar = part_bar[i]+--                             let nz_prd = non_zero_prod[i]+--                             in if zr_cts == 0+--                             then pr_bar * (nz_prd / v)+--                             else if zr_cts == 1 and v == 0+--                             then nz_prd * pr_bar+--                             else 0+--                    ) is vs+diffMulHist ::+  VjpOps -> VName -> StmAux () -> SubExp -> BinOp -> SubExp -> VName -> VName -> SubExp -> SubExp -> VName -> ADM () -> ADM ()+diffMulHist _ops x aux n mul ne is vs w rf dst m = do+  let t = binOpType mul+  vs_type <- lookupType vs+  let vs_dims = arrayDims vs_type+  let vs_elm_type = elemType vs_type+  dst_type <- lookupType dst+  let dst_dims = arrayDims dst_type+  let inner_dims = tail vs_dims++  v_param <- newParam "v" $ Prim t+  lam_ps_zs_inner <-+    mkLambda [v_param] $+      fmap varsRes . letTupExp "map_res"+        =<< eIf+          (eCmpOp (CmpEq t) (eParam v_param) (eSubExp $ Constant $ blankPrimValue t))+          (eBody $ fmap eSubExp [Constant $ onePrimValue t, intConst Int64 1])+          (eBody [eParam v_param, eSubExp $ intConst Int64 0])+  lam_ps_zs <- nestedmap vs_dims [vs_elm_type] lam_ps_zs_inner+  ps_zs_res <- eLambda lam_ps_zs [eSubExp $ Var vs]+  ps_zs <- bindSubExpRes "ps_zs" ps_zs_res+  let [ps, zs] = ps_zs++  lam_mul_inner <- binOpLambda mul t+  lam_mul <- nestedmap inner_dims [vs_elm_type, vs_elm_type] lam_mul_inner+  nz_prods0 <- letExp "nz_prd" $ BasicOp $ Replicate (Shape [w]) ne+  let hist_nzp = HistOp (Shape [w]) rf [nz_prods0] [ne] lam_mul++  lam_add_inner <- binOpLambda (Add Int64 OverflowUndef) int64+  lam_add <- nestedmap inner_dims [int64, int64] lam_add_inner+  zr_counts0 <- letExp "zr_cts" $ BasicOp $ Replicate (Shape dst_dims) (intConst Int64 0)+  zrn_ne <- letSubExp "zr_ne" $ BasicOp $ Replicate (Shape inner_dims) (intConst Int64 0)+  let hist_zrn = HistOp (Shape [w]) rf [zr_counts0] [if length vs_dims == 1 then intConst Int64 0 else zrn_ne] lam_add++  f' <- mkIdentityLambda [Prim int64, Prim int64, rowType vs_type, rowType $ Array int64 (Shape vs_dims) NoUniqueness]+  nz_prods <- newVName "non_zero_prod"+  zr_counts <- newVName "zero_count"+  auxing aux $+    letBindNames [nz_prods, zr_counts] $+      Op $+        Hist n [is, is, ps, zs] [hist_nzp, hist_zrn] f'++  p_param <- newParam "prod" $ Prim t+  c_param <- newParam "count" $ Prim int64+  lam_h_part_inner <-+    mkLambda [p_param, c_param] $+      fmap varsRes . letTupExp "h_part"+        =<< eIf+          (toExp $ 0 .==. le64 (paramName c_param))+          (eBody $ pure $ eParam p_param)+          (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)+  lam_h_part <- nestedmap dst_dims [vs_elm_type, int64] lam_h_part_inner+  h_part_res <- eLambda lam_h_part $ map (eSubExp . Var) [nz_prods, zr_counts]+  h_part' <- bindSubExpRes "h_part" h_part_res+  let [h_part] = h_part'++  lam_mul_inner' <- binOpLambda mul t+  lam_mul' <- nestedmap dst_dims [vs_elm_type, vs_elm_type] lam_mul_inner'+  x_res <- eLambda lam_mul' $ map (eSubExp . Var) [dst, h_part]+  x' <- bindSubExpRes "x" x_res+  auxing aux $ letBindNames [x] $ BasicOp $ SubExp $ Var $ head x'++  m++  x_bar <- lookupAdjVal x++  lam_mul'' <- renameLambda lam_mul'+  dst_bar_res <- eLambda lam_mul'' $ map (eSubExp . Var) [h_part, x_bar]+  dst_bar <- bindSubExpRes (baseString dst <> "_bar") dst_bar_res+  updateAdj dst $ head dst_bar++  lam_mul''' <- renameLambda lam_mul'+  part_bar_res <- eLambda lam_mul''' $ map (eSubExp . Var) [dst, x_bar]+  part_bar' <- bindSubExpRes "part_bar" part_bar_res+  let [part_bar] = part_bar'++  inner_params <- zipWithM newParam ["zr_cts", "pr_bar", "nz_prd", "a"] $ map Prim [int64, t, t, t]+  let [zr_cts, pr_bar, nz_prd, a_param] = inner_params+  lam_vsbar_inner <-+    mkLambda inner_params $+      fmap varsRes . letTupExp "vs_bar" =<< do+        eIf+          (eCmpOp (CmpEq int64) (eSubExp $ intConst Int64 0) (eParam zr_cts))+          (eBody $ pure $ eBinOp mul (eParam pr_bar) $ eBinOp (getBinOpDiv t) (eParam nz_prd) $ eParam a_param)+          ( eBody $+              pure $+                eIf+                  ( eBinOp+                      LogAnd+                      (eCmpOp (CmpEq int64) (eSubExp $ intConst Int64 1) (eParam zr_cts))+                      (eCmpOp (CmpEq t) (eSubExp $ Constant $ blankPrimValue t) $ eParam a_param)+                  )+                  (eBody $ pure $ eBinOp mul (eParam nz_prd) (eParam pr_bar))+                  (eBody $ pure $ eSubExp $ Constant $ blankPrimValue t)+          )++  lam_vsbar_middle <- nestedmap inner_dims [int64, t, t, t] lam_vsbar_inner++  i_param <- newParam "i" $ Prim int64+  a_param' <- newParam "a" $ rowType vs_type+  lam_vsbar <-+    mkLambda [i_param, a_param'] $+      fmap varsRes . letTupExp "vs_bar"+        =<< eIf+          (toExp $ withinBounds $ pure (w, paramName i_param))+          ( buildBody_ $ do+              let i = fullSlice vs_type [DimFix $ Var $ paramName i_param]+              names <- traverse newVName ["zr_cts", "pr_bar", "nz_prd"]+              zipWithM_ (\name -> letBindNames [name] . BasicOp . flip Index i) names [zr_counts, part_bar, nz_prods]+              eLambda lam_vsbar_middle $ map (eSubExp . Var) names <> [eParam a_param']+          )+          (eBody $ pure $ pure $ zeroExp $ rowType dst_type)++  vs_bar <-+    letExp (baseString vs <> "_bar") $ Op $ Screma n [is, vs] $ mapSOAC lam_vsbar++  updateAdj vs vs_bar++--+-- special case of histogram with add as operator.+-- Original, assuming `is: [n]i64` and `dst: [w]btp`+--     let x = reduce_by_index dst (+) ne is vs+-- Forward sweep:+--     need to copy dst: reverse sweep might use it 7+--       (see ex. in reducebyindexminmax6.fut where the first map requires the original dst to be differentiated).+--     let dst_cpy = copy dst+--     let x = reduce_by_index dst_cpy (+) ne is vs+--+-- Reverse sweep:+--     dst_bar += x_bar+--+--     vs_bar += map (\i -> x_bar[i]) is+diffAddHist ::+  VjpOps -> VName -> StmAux () -> SubExp -> Lambda SOACS -> SubExp -> VName -> VName -> SubExp -> SubExp -> VName -> ADM () -> ADM ()+diffAddHist _ops x aux n add ne is vs w rf dst m = do+  let t = paramDec $ head $ lambdaParams add++  dst_cpy <- letExp (baseString dst <> "_copy") $ BasicOp $ Copy dst++  f <- mkIdentityLambda [Prim int64, t]+  auxing aux . letBindNames [x] . Op $+    Hist n [is, vs] [HistOp (Shape [w]) rf [dst_cpy] [ne] add] f++  m++  x_bar <- lookupAdjVal x++  updateAdj dst x_bar++  x_type <- lookupType x+  i_param <- newParam (baseString vs <> "_i") $ Prim int64+  let i = paramName i_param+  lam_vsbar <-+    mkLambda [i_param] $+      fmap varsRes . letTupExp "vs_bar"+        =<< eIf+          (toExp $ withinBounds $ pure (w, i))+          (eBody $ pure $ pure $ BasicOp $ Index x_bar $ fullSlice x_type [DimFix $ Var i])+          (eBody $ pure $ eSubExp ne)++  vs_bar <- letExp (baseString vs <> "_bar") $ Op $ Screma n [is] $ mapSOAC lam_vsbar+  updateAdj vs vs_bar++--+-- a step in the radix sort implementation+-- it assumes the key we are sorting+-- after is [n]i64 and it is the first VName+--+-- local def radix_sort_step [n] 't (xs: [n]t) (get_bit: i32 -> t -> i32)+--                                  (digit_n: i32): [n]t =+--   let num x = get_bit (digit_n+1) x * 2 + get_bit digit_n x+--   let pairwise op (a1,b1,c1,d1) (a2,b2,c2,d2) =+--     (a1 `op` a2, b1 `op` b2, c1 `op` c2, d1 `op` d2)+--   let bins = xs |> map num+--   let flags = bins |> map (\x -> if x == 0 then (1,0,0,0)+--                                  else if x == 1 then (0,1,0,0)+--                                  else if x == 2 then (0,0,1,0)+--                                  else (0,0,0,1))+--   let offsets = scan (pairwise (+)) (0,0,0,0) flags+--   let (na,nb,nc,_nd) = last offsets+--   let f bin (a,b,c,d) = match bin+--                         case 0 -> a-1+--                         case 1 -> na+b-1+--                         case 2 -> na+nb+c-1+--                         case _ -> na+nb+nc+d-1+--   let is = map2 f bins offsets+--   in scatter scratch is xs+radixSortStep :: [VName] -> [Type] -> SubExp -> SubExp -> SubExp -> ADM [VName]+radixSortStep xs tps bit n w = do+  -- let is = head xs+  is <- mapout (head xs) n w++  num_param <- newParam "num" $ Prim int64+  num_lam <-+    mkLambda [num_param] $+      fmap varsRes . letTupExp "num_res"+        =<< eBinOp+          (Add Int64 OverflowUndef)+          ( eBinOp+              (And Int64)+              (eBinOp (AShr Int64) (eParam num_param) (eSubExp bit))+              (iConst 1)+          )+          ( eBinOp+              (Mul Int64 OverflowUndef)+              (iConst 2)+              ( eBinOp+                  (And Int64)+                  (eBinOp (AShr Int64) (eParam num_param) (eBinOp (Add Int64 OverflowUndef) (eSubExp bit) (iConst 1)))+                  (iConst 1)+              )+          )++  bins <- letExp "bins" $ Op $ Screma n [is] $ mapSOAC num_lam+  flag_param <- newParam "flag" $ Prim int64+  flag_lam <-+    mkLambda [flag_param] $+      fmap varsRes . letTupExp "flag_res"+        =<< elseIf+          int64+          (map ((,) (eParam flag_param) . iConst) [0 .. 2])+          (map (eBody . fmap iConst . (\i -> map (\j -> if i == j then 1 else 0) [0 .. 3])) ([0 .. 3] :: [Integer]))++  flags <- letTupExp "flags" $ Op $ Screma n [bins] $ mapSOAC flag_lam++  scan_params <- traverse (flip newParam $ Prim int64) ["a1", "b1", "c1", "d1", "a2", "b2", "c2", "d2"]+  scan_lam <-+    mkLambda scan_params $+      fmap subExpsRes . mapM (letSubExp "scan_res") =<< do+        uncurry (zipWithM (eBinOp $ Add Int64 OverflowUndef)) $ splitAt 4 $ map eParam scan_params++  scan <- scanSOAC $ pure $ Scan scan_lam $ map (intConst Int64) [0, 0, 0, 0]+  offsets <- letTupExp "offsets" $ Op $ Screma n flags scan++  ind <- letSubExp "ind_last" =<< eBinOp (Sub Int64 OverflowUndef) (eSubExp n) (iConst 1)+  let i = Slice [DimFix ind]+  nabcd <- traverse newVName ["na", "nb", "nc", "nd"]+  zipWithM_ (\abcd -> letBindNames [abcd] . BasicOp . flip Index i) nabcd offsets++  let vars = map Var nabcd+  map_params <- traverse (flip newParam $ Prim int64) ["bin", "a", "b", "c", "d"]+  map_lam <-+    mkLambda map_params $+      fmap varsRes . letTupExp "map_res"+        =<< elseIf+          int64+          (map ((,) (eParam $ head map_params) . iConst) [0 .. 2])+          ( zipWith+              ( \j p ->+                  eBody $+                    pure $ do+                      t <- letSubExp "t" =<< eBinOp (Sub Int64 OverflowUndef) (eParam p) (iConst 1)+                      foldBinOp (Add Int64 OverflowUndef) (intConst Int64 0) (t : take j vars)+              )+              [0 .. 3]+              (tail map_params)+          )++  nis <- letExp "nis" $ Op $ Screma n (bins : offsets) $ mapSOAC map_lam++  scatter_dst <- traverse (\t -> letExp "scatter_dst" $ BasicOp $ Scratch (elemType t) (arrayDims t)) tps+  multiScatter n scatter_dst nis xs+  where+    iConst c = eSubExp $ intConst Int64 c++--+-- the radix sort implementation+-- def radix_sort [n] 't (xs: [n]i64) =+--   let iters = if n == 0 then 0 else 32+--   in loop xs for i < iters do radix_sort_step xs i64.get_bit (i*2)+radixSort :: [VName] -> SubExp -> SubExp -> ADM [VName]+radixSort xs n w = do+  logw <- log2 =<< letSubExp "w1" =<< toExp (pe64 w + 1)+  -- ceil logw by (logw + 1) / 2+  iters <- letSubExp "iters" =<< toExp (untyped (pe64 logw + 1) ~/~ untyped (pe64 (intConst Int64 2)))++  types <- traverse lookupType xs+  params <- zipWithM (\x -> newParam (baseString x) . flip toDecl Nonunique) xs types+  i <- newVName "i"+  loopbody <- buildBody_ . localScope (scopeOfFParams params) $+    fmap varsRes $ do+      bit <- letSubExp "bit" =<< toExp (le64 i * 2)+      radixSortStep (map paramName params) types bit n w++  letTupExp "sorted" $+    DoLoop+      (zip params $ map Var xs)+      (ForLoop i Int64 iters [])+      loopbody+  where+    log2 :: SubExp -> ADM SubExp+    log2 m = do+      params <- zipWithM newParam ["cond", "r", "i"] $ map Prim [Bool, int64, int64]+      let [cond, r, i] = params++      body <- buildBody_ . localScope (scopeOfFParams params) $ do+        r' <- letSubExp "r'" =<< toExp (le64 (paramName r) .>>. 1)+        cond' <- letSubExp "cond'" =<< toExp (bNot $ pe64 r' .==. 0)+        i' <- letSubExp "i'" =<< toExp (le64 (paramName i) + 1)+        pure $ subExpsRes [cond', r', i']++      cond_init <- letSubExp "test" =<< toExp (bNot $ pe64 m .==. 0)++      l <-+        letTupExp' "log2res" $+          DoLoop+            (zip params [cond_init, m, Constant $ blankPrimValue int64])+            (WhileLoop $ paramName cond)+            body++      let [_, _, res] = l+      pure res++radixSort' :: [VName] -> SubExp -> SubExp -> ADM [VName]+radixSort' xs n w = do+  iota_n <-+    letExp "red_iota" . BasicOp $+      Iota n (intConst Int64 0) (intConst Int64 1) Int64++  radres <- radixSort [head xs, iota_n] n w+  let [is', iota'] = radres++  i_param <- newParam "i" $ Prim int64+  let slice = [DimFix $ Var $ paramName i_param]+  map_lam <- mkLambda [i_param] $ varsRes <$> multiIndex (tail xs) slice++  sorted <- letTupExp "sorted" $ Op $ Screma n [iota'] $ mapSOAC map_lam+  pure $ iota' : is' : sorted++--+-- generic case of histogram.+-- Original, assuming `is: [n]i64` and `dst: [w]btp`+--   let xs = reduce_by_index dst odot ne is as+-- Forward sweep:+-- let h_part = reduce_by_index (replicate w ne) odot ne is as+-- let xs = map2 odot dst h_part+-- Reverse sweep:+-- h_part_bar += f'' dst h_part+-- dst_bar += f' dst h_part++-- let flag = map (\i -> i == 0 || sis[i] != sis[i-1]) (iota n)+-- let flag_rev = map (\i -> i==0 || flag[n-i]) (iota n)+-- let ls = seg_scan_exc odot ne flag sas+-- let rs = reverse sas |>+--          seg_scan_exc odot ne flag_rev |> reverse+-- let f_bar = map (\i -> if i < w && -1 < w+--                        then h_part_bar[i]+--                        else 0s+--                 ) sis+-- let sas_bar = f f_dst ls sas rs+-- as_bar += scatter (Scratch alpha n) siota sas_bar+-- Where:+--  siota: 'iota n' sorted wrt 'is'+--  sis: 'is' sorted wrt 'is'+--  sas: 'as' sorted wrt 'is'+--  f'' = vjpLambda xs_bar h_part (map2 odot)+--  f' = vjpLambda xs_bar dst (map2 odot)+--  f  = vjpLambda f_bar sas (map4 (\di li ai ri -> di odot li odot ai odot ri))+--  0s is an alpha-dimensional array with 0 (possibly 0-dim)+diffHist :: VjpOps -> [VName] -> StmAux () -> SubExp -> Lambda SOACS -> [SubExp] -> [VName] -> [SubExp] -> SubExp -> [VName] -> ADM () -> ADM ()+diffHist ops xs aux n lam0 ne as w rf dst m = do+  as_type <- traverse lookupType $ tail as+  dst_type <- traverse lookupType dst++  nes <- traverse (letExp "new_dst" . BasicOp . Replicate (Shape $ pure $ head w)) ne++  h_map <- mkIdentityLambda $ Prim int64 : map rowType as_type+  h_part <- traverse (newVName . flip (<>) "_h_part" . baseString) xs+  auxing aux . letBindNames h_part . Op $+    Hist n as [HistOp (Shape w) rf nes ne lam0] h_map++  lam0' <- renameLambda lam0+  auxing aux . letBindNames xs . Op $+    Screma (head w) (dst <> h_part) (mapSOAC lam0')++  m++  xs_bar <- traverse lookupAdjVal xs++  (dst_params, hp_params, f') <- mkF' lam0 dst_type $ head w+  f'_adj_dst <- vjpLambda ops (map adjFromVar xs_bar) dst_params f'+  f'_adj_hp <- vjpLambda ops (map adjFromVar xs_bar) hp_params f'++  dst_bar' <- eLambda f'_adj_dst $ map (eSubExp . Var) $ dst <> h_part+  dst_bar <- bindSubExpRes "dst_bar" dst_bar'+  zipWithM_ updateAdj dst dst_bar++  h_part_bar' <- eLambda f'_adj_hp $ map (eSubExp . Var) $ dst <> h_part+  h_part_bar <- bindSubExpRes "h_part_bar" h_part_bar'++  lam <- renameLambda lam0+  lam' <- renameLambda lam0++  -- is' <- mapout (head as) n (head w)+  -- sorted <- radixSort' (is' : tail as) n $ head w+  sorted <- radixSort' as n $ head w+  let siota = head sorted+  let sis = head $ tail sorted+  let sas = drop 2 sorted++  iota_n <-+    letExp "iota" $ BasicOp $ Iota n (intConst Int64 0) (intConst Int64 1) Int64++  par_i <- newParam "i" $ Prim int64+  flag_lam <- mkFlagLam par_i sis+  flag <- letExp "flag" $ Op $ Screma n [iota_n] $ mapSOAC flag_lam++  -- map (\i -> (if flag[i] then (true,ne) else (false,vs[i-1]), if i==0 || flag[n-i] then (true,ne) else (false,vs[n-i]))) (iota n)+  par_i' <- newParam "i" $ Prim int64+  let i' = paramName par_i'+  g_lam <-+    mkLambda [par_i'] $+      fmap subExpsRes . mapM (letSubExp "scan_inps") =<< do+        im1 <- letSubExp "i_1" =<< toExp (le64 i' - 1)+        nmi <- letSubExp "n_i" =<< toExp (pe64 n - le64 i')+        let s1 = [DimFix im1]+        let s2 = [DimFix nmi]++        -- flag array for left scan+        f1 <- letSubExp "f1" $ BasicOp $ Index flag $ Slice [DimFix $ Var i']++        -- array for left scan+        r1 <-+          letTupExp' "r1"+            =<< eIf+              (eSubExp f1)+              (eBody $ fmap eSubExp ne)+              (eBody . fmap (eSubExp . Var) =<< multiIndex sas s1)++        -- array for right scan inc flag+        r2 <-+          letTupExp' "r2"+            =<< eIf+              (toExp $ le64 i' .==. 0)+              (eBody $ fmap eSubExp $ Constant (onePrimValue Bool) : ne)+              ( eBody $+                  pure $ do+                    eIf+                      (pure $ BasicOp $ Index flag $ Slice s2)+                      (eBody $ fmap eSubExp $ Constant (onePrimValue Bool) : ne)+                      ( eBody . fmap eSubExp . (Constant (blankPrimValue Bool) :) . fmap Var+                          =<< multiIndex sas s2+                      )+              )++        traverse eSubExp $ f1 : r1 ++ r2++  -- scan (\(f1,v1) (f2,v2) ->+  --   let f = f1 || f2+  --   let v = if f2 then v2 else g v1 v2+  --   in (f,v) ) (false,ne) (zip flags vals)+  scan_lams <-+    traverse+      ( \l -> do+          f1 <- newParam "f1" $ Prim Bool+          f2 <- newParam "f2" $ Prim Bool+          ps <- lambdaParams <$> renameLambda lam0+          let (p1, p2) = splitAt (length ne) ps++          mkLambda (f1 : p1 ++ f2 : p2) $+            fmap varsRes . letTupExp "scan_res" =<< do+              let f = eBinOp LogOr (eParam f1) (eParam f2)+              eIf+                (eParam f2)+                (eBody $ f : fmap eParam p2)+                ( eBody . (f :) . fmap (eSubExp . Var)+                    =<< bindSubExpRes "gres"+                    =<< eLambda l (fmap eParam ps)+                )+      )+      [lam, lam']++  let ne' = Constant (BoolValue False) : ne++  scansres <-+    letTupExp "adj_ctrb_scan" . Op $+      Screma n [iota_n] (scanomapSOAC (map (`Scan` ne') scan_lams) g_lam)++  let (_ : ls_arr, _ : rs_arr_rev) = splitAt (length ne + 1) scansres++  -- map (\i -> if i < w && -1 < w then (xs_bar[i], dst[i]) else (0,ne)) sis+  par_i'' <- newParam "i" $ Prim int64+  let i'' = paramName par_i''+  map_lam <-+    mkLambda [par_i''] $+      fmap varsRes . letTupExp "scan_res"+        =<< eIf+          (toExp $ withinBounds $ pure (head w, i''))+          (eBody . fmap (eSubExp . Var) =<< multiIndex h_part_bar [DimFix $ Var i''])+          ( eBody $ do+              map (\t -> pure $ BasicOp $ Replicate (Shape $ tail $ arrayDims t) (Constant $ blankPrimValue $ elemType t)) as_type+          )++  f_bar <- letTupExp "f_bar" $ Op $ Screma n [sis] $ mapSOAC map_lam++  (as_params, f) <- mkF lam0 as_type n+  f_adj <- vjpLambda ops (map adjFromVar f_bar) as_params f++  -- map (\i -> rs_arr_rev[n-i-1]) (iota n)+  par_i''' <- newParam "i" $ Prim int64+  let i''' = paramName par_i'''+  rev_lam <- mkLambda [par_i'''] $ do+    nmim1 <- letSubExp "n_i_1" =<< toExp (pe64 n - le64 i''' - 1)+    varsRes <$> multiIndex rs_arr_rev [DimFix nmim1]++  rs_arr <- letTupExp "rs_arr" $ Op $ Screma n [iota_n] $ mapSOAC rev_lam++  sas_bar <-+    bindSubExpRes "sas_bar"+      =<< eLambda f_adj (map (eSubExp . Var) $ ls_arr <> sas <> rs_arr)++  scatter_dst <- traverse (\t -> letExp "scatter_dst" $ BasicOp $ Scratch (elemType t) (arrayDims t)) as_type+  as_bar <- multiScatter n scatter_dst siota sas_bar++  zipWithM_ updateAdj (tail as) as_bar+  where+    -- map (\i -> if i == 0 then true else is[i] != is[i-1]) (iota n)+    mkFlagLam :: LParam SOACS -> VName -> ADM (Lambda SOACS)+    mkFlagLam par_i sis =+      mkLambda [par_i] $+        fmap varsRes . letTupExp "flag" =<< do+          let i = paramName par_i+          eIf+            (toExp (le64 i .==. 0))+            (eBody $ pure $ eSubExp $ Constant $ onePrimValue Bool)+            ( eBody $+                pure $ do+                  i_p <- letExp "i_p" =<< toExp (le64 i - 1)+                  vs <- traverse (letExp "vs" . BasicOp . Index sis . Slice . pure . DimFix . Var) [i, i_p]+                  let [vs_i, vs_p] = vs+                  toExp $ bNot $ le64 vs_i .==. le64 vs_p+            )

@@ -3,6 +3,8 @@ module Futhark.AD.Rev.Reduce   ( diffReduce,     diffMinMaxReduce,+    diffVecReduce,+    diffMulReduce,   ) where @@ -201,3 +203,140 @@     letSubExp "minmax_in_bounds" . BasicOp $       CmpOp (CmpSlt Int64) (intConst Int64 0) w   updateAdjIndex as (CheckBounds (Just in_bounds), Var x_ind) (Var x_adj)++--+-- Special case of vectorised reduce:+--    let x = reduce (map2 op) nes as+-- Idea:+--    rewrite to+--      let x = map2 (\as ne -> reduce op ne as) (transpose as) nes+--    and diff+diffVecReduce ::+  VjpOps -> Pat Type -> StmAux () -> SubExp -> Commutativity -> Lambda SOACS -> VName -> VName -> ADM () -> ADM ()+diffVecReduce ops x aux w iscomm lam ne as m = do+  stms <- collectStms_ $ do+    rank <- arrayRank <$> lookupType as+    let rear = [1, 0] ++ drop 2 [0 .. rank - 1]++    tran_as <- letExp "tran_as" $ BasicOp $ Rearrange rear as+    ts <- lookupType tran_as+    t_ne <- lookupType ne++    as_param <- newParam "as_param" $ rowType ts+    ne_param <- newParam "ne_param" $ rowType t_ne++    reduce_form <- reduceSOAC [Reduce iscomm lam [Var $ paramName ne_param]]++    map_lam <-+      mkLambda [as_param, ne_param] $+        fmap varsRes . letTupExp "idx_res" $+          Op $+            Screma w [paramName as_param] reduce_form+    addStm $ Let x aux $ Op $ Screma (arraySize 0 ts) [tran_as, ne] $ mapSOAC map_lam++  foldr (vjpStm ops) m stms++--+-- Special case of reduce with mul:+--    let x = reduce (*) ne as+-- Forward trace (assuming w = length as):+--    let (p, z) = map (\a -> if a == 0 then (1, 1) else (a, 0)) as+--    non_zero_prod = reduce (*) ne p+--    zr_count = reduce (+) 0 z+--    let x =+--      if 0 == zr_count+--      then non_zero_prod+--      else 0+-- Reverse trace:+--    as_bar = map2+--      (\a a_bar ->+--        if zr_count == 0+--        then a_bar + non_zero_prod/a * x_bar+--        else if zr_count == 1+--        then a_bar + (if a == 0 then non_zero_prod * x_bar else 0)+--        else as_bar+--      ) as as_bar+diffMulReduce ::+  VjpOps -> VName -> StmAux () -> SubExp -> BinOp -> SubExp -> VName -> ADM () -> ADM ()+diffMulReduce _ops x aux w mul ne as m = do+  let t = binOpType mul+  let const_zero = eSubExp $ Constant $ blankPrimValue t++  a_param <- newParam "a" $ Prim t+  map_lam <-+    mkLambda [a_param] $+      fmap varsRes . letTupExp "map_res"+        =<< eIf+          (eCmpOp (CmpEq t) (eParam a_param) const_zero)+          (eBody $ fmap eSubExp [Constant $ onePrimValue t, intConst Int64 1])+          (eBody [eParam a_param, eSubExp $ intConst Int64 0])++  ps <- newVName "ps"+  zs <- newVName "zs"+  auxing aux $+    letBindNames [ps, zs] $+      Op $+        Screma w [as] $+          mapSOAC map_lam++  red_lam_mul <- binOpLambda mul t+  red_lam_add <- binOpLambda (Add Int64 OverflowUndef) int64++  red_form_mul <- reduceSOAC $ pure $ Reduce Commutative red_lam_mul $ pure ne+  red_form_add <- reduceSOAC $ pure $ Reduce Commutative red_lam_add $ pure $ intConst Int64 0++  nz_prods <- newVName "non_zero_prod"+  zr_count <- newVName "zero_count"+  auxing aux $ letBindNames [nz_prods] $ Op $ Screma w [ps] red_form_mul+  auxing aux $ letBindNames [zr_count] $ Op $ Screma w [zs] red_form_add++  auxing aux $+    letBindNames [x]+      =<< eIf+        (toExp $ 0 .==. le64 zr_count)+        (eBody $ pure $ eSubExp $ Var nz_prods)+        (eBody $ pure const_zero)++  m++  x_adj <- lookupAdjVal x++  a_param_rev <- newParam "a" $ Prim t+  map_lam_rev <-+    mkLambda [a_param_rev] $+      fmap varsRes . letTupExp "adj_res"+        =<< eIf+          (toExp $ 0 .==. le64 zr_count)+          ( eBody $+              pure $+                eBinOp mul (eSubExp $ Var x_adj) $+                  eBinOp (getDiv t) (eSubExp $ Var nz_prods) $+                    eParam a_param_rev+          )+          ( eBody $+              pure $+                eIf+                  (toExp $ 1 .==. le64 zr_count)+                  ( eBody $+                      pure $+                        eIf+                          (eCmpOp (CmpEq t) (eParam a_param_rev) const_zero)+                          ( eBody $+                              pure $+                                eBinOp mul (eSubExp $ Var x_adj) $+                                  eSubExp $+                                    Var nz_prods+                          )+                          (eBody $ pure const_zero)+                  )+                  (eBody $ pure const_zero)+          )++  as_adjup <- letExp "adjs" $ Op $ Screma w [as] $ mapSOAC map_lam_rev++  updateAdj as as_adjup+  where+    getDiv :: PrimType -> BinOp+    getDiv (IntType t) = SDiv t Unsafe+    getDiv (FloatType t) = FDiv t+    getDiv _ = error "In getDiv, Reduce.hs: input not supported"

@@ -3,6 +3,7 @@ module Futhark.AD.Rev.SOAC (vjpSOAC) where  import Control.Monad+import Futhark.AD.Rev.Hist import Futhark.AD.Rev.Map import Futhark.AD.Rev.Monad import Futhark.AD.Rev.Reduce@@ -34,14 +35,47 @@       onOps _ _ _ = m   onOps ops pat_per_op as_per_op +-- We split multi-op histograms into multiple operations so we+-- can take advantage of special cases. Post-AD, the result may+-- be fused again.+splitHist :: VjpOps -> Pat Type -> StmAux () -> [HistOp SOACS] -> SubExp -> [VName] -> [VName] -> ADM () -> ADM ()+splitHist vjpops pat aux ops w is as m = do+  let ks = map (length . histNeutral) ops+      pat_per_op = map Pat $ chunks ks $ patElems pat+      as_per_op = chunks ks as+      onOps (op : ops') (op_pat : op_pats') (op_is : op_is') (op_as : op_as') = do+        f <- mkIdentityLambda . (Prim int64 :) =<< traverse lookupType op_as+        vjpSOAC vjpops op_pat aux (Hist w (op_is : op_as) [op] f) $+          onOps ops' op_pats' op_is' op_as'+      onOps _ _ _ _ = m+  onOps ops pat_per_op is as_per_op++-- unfusing a map-histogram construct into a map and a histogram.+histomapToMapAndHist :: Pat Type -> (SubExp, [HistOp SOACS], Lambda SOACS, [VName]) -> ADM (Stm SOACS, Stm SOACS)+histomapToMapAndHist (Pat pes) (w, histops, map_lam, as) = do+  map_pat <- traverse accMapPatElem $ lambdaReturnType map_lam+  let map_stm = mkLet map_pat $ Op $ Screma w as $ mapSOAC map_lam+  new_lam <- mkIdentityLambda $ lambdaReturnType map_lam+  let hist_stm = Let (Pat pes) (defAux ()) $ Op $ Hist w (map identName map_pat) histops new_lam+  pure (map_stm, hist_stm)+  where+    accMapPatElem =+      newIdent "hist_map_res" . (`arrayOfRow` w)+ commonSOAC :: Pat Type -> StmAux () -> SOAC SOACS -> ADM () -> ADM [Adj] commonSOAC pat aux soac m = do   addStm $ Let pat aux $ Op soac   m   returnSweepCode $ mapM lookupAdj $ patNames pat +-- Reverse-mode differentiation of SOACs vjpSOAC :: VjpOps -> Pat Type -> StmAux () -> SOAC SOACS -> ADM () -> ADM ()+-- Differentiating Reduces vjpSOAC ops pat aux soac@(Screma w as form) m+  | Just [Reduce iscomm lam [Var ne]] <- isReduceSOAC form,+    [a] <- as,+    Just op <- mapOp lam =+      diffVecReduce ops pat aux w iscomm op ne a m   | Just reds <- isReduceSOAC form,     length reds > 1 =       splitScanRed ops (reduceSOAC, redNeutral) (pat, aux, reds, w, as) m@@ -52,35 +86,133 @@     Just [(op, _, _, _)] <- lamIsBinOp $ redLambda red,     isMinMaxOp op =       diffMinMaxReduce ops x aux w op ne a m+  | Just [red] <- isReduceSOAC form,+    [x] <- patNames pat,+    [ne] <- redNeutral red,+    [a] <- as,+    Just [(op, _, _, _)] <- lamIsBinOp $ redLambda red,+    isMulOp op =+      diffMulReduce ops x aux w op ne a m   | Just red <- singleReduce <$> isReduceSOAC form = do       pat_adj <- mapM adjVal =<< commonSOAC pat aux soac m       diffReduce ops pat_adj w as red-  where-    isMinMaxOp (SMin _) = True-    isMinMaxOp (UMin _) = True-    isMinMaxOp (FMin _) = True-    isMinMaxOp (SMax _) = True-    isMinMaxOp (UMax _) = True-    isMinMaxOp (FMax _) = True-    isMinMaxOp _ = False++-- Differentiating Scans vjpSOAC ops pat aux soac@(Screma w as form) m+  | Just [Scan lam [ne]] <- isScanSOAC form,+    [x] <- patNames pat,+    [a] <- as,+    Just [(op, _, _, _)] <- lamIsBinOp lam,+    isAddOp op = do+      void $ commonSOAC pat aux soac m+      diffScanAdd ops x w lam ne a+  | Just [Scan lam ne] <- isScanSOAC form,+    Just op <- mapOp lam = do+      diffScanVec ops (patNames pat) aux w op ne as m   | Just scans <- isScanSOAC form,     length scans > 1 =       splitScanRed ops (scanSOAC, scanNeutral) (pat, aux, scans, w, as) m   | Just red <- singleScan <$> isScanSOAC form = do       void $ commonSOAC pat aux soac m       diffScan ops (patNames pat) w as red++-- Differentiating Maps vjpSOAC ops pat aux soac@(Screma w as form) m   | Just lam <- isMapSOAC form = do       pat_adj <- commonSOAC pat aux soac m       vjpMap ops pat_adj aux w lam as++-- Differentiating Redomaps vjpSOAC ops pat _aux (Screma w as form) m   | Just (reds, map_lam) <-       isRedomapSOAC form = do       (mapstm, redstm) <-         redomapToMapAndReduce pat (w, reds, map_lam, as)       vjpStm ops mapstm $ vjpStm ops redstm m++-- Differentiating Scatter vjpSOAC ops pat aux (Scatter w lam ass written_info) m =   vjpScatter ops pat aux (w, lam, ass, written_info) m+-- Differentiating Histograms+vjpSOAC ops pat aux (Hist n as histops f) m+  | isIdentityLambda f,+    length histops > 1 = do+      let (is, vs) = splitAt (length histops) as+      splitHist ops pat aux histops n is vs m+vjpSOAC ops pat aux (Hist n [is, vs] [histop] f) m+  | isIdentityLambda f,+    [x] <- patNames pat,+    HistOp (Shape [w]) rf [dst] [ne] lam <- histop,+    lam' <- nestedMapOp lam,+    Just [(op, _, _, _)] <- lamIsBinOp lam',+    isMinMaxOp op =+      diffMinMaxHist ops x aux n op ne is vs w rf dst m+  | isIdentityLambda f,+    [x] <- patNames pat,+    HistOp (Shape [w]) rf [dst] [ne] lam <- histop,+    lam' <- nestedMapOp lam,+    Just [(op, _, _, _)] <- lamIsBinOp lam',+    isMulOp op =+      diffMulHist ops x aux n op ne is vs w rf dst m+  | isIdentityLambda f,+    [x] <- patNames pat,+    HistOp (Shape [w]) rf [dst] [ne] lam <- histop,+    lam' <- nestedMapOp lam,+    Just [(op, _, _, _)] <- lamIsBinOp lam',+    isAddOp op =+      diffAddHist ops x aux n lam ne is vs w rf dst m+vjpSOAC ops pat aux (Hist n as [histop] f) m+  | isIdentityLambda f,+    HistOp (Shape w) rf dst ne lam <- histop = do+      diffHist ops (patNames pat) aux n lam ne as w rf dst m+vjpSOAC ops pat _aux (Hist n as histops f) m+  | not (isIdentityLambda f) = do+      (mapstm, redstm) <-+        histomapToMapAndHist pat (n, histops, f, as)+      vjpStm ops mapstm $ vjpStm ops redstm m vjpSOAC _ _ _ soac _ =   error $ "vjpSOAC unhandled:\n" ++ prettyString soac++---------------+--- Helpers ---+---------------++isMinMaxOp :: BinOp -> Bool+isMinMaxOp (SMin _) = True+isMinMaxOp (UMin _) = True+isMinMaxOp (FMin _) = True+isMinMaxOp (SMax _) = True+isMinMaxOp (UMax _) = True+isMinMaxOp (FMax _) = True+isMinMaxOp _ = False++isMulOp :: BinOp -> Bool+isMulOp (Mul _ _) = True+isMulOp (FMul _) = True+isMulOp _ = False++isAddOp :: BinOp -> Bool+isAddOp (Add _ _) = True+isAddOp (FAdd _) = True+isAddOp _ = False++-- Identifies vectorized operators (lambdas):+--   if the lambda argument is a map, then returns+--   just the map's lambda; otherwise nothing.+mapOp :: Lambda SOACS -> Maybe (Lambda SOACS)+mapOp (Lambda [pa1, pa2] lam_body _)+  | [SubExpRes cs r] <- bodyResult lam_body,+    cs == mempty,+    [map_stm] <- stmsToList (bodyStms lam_body),+    (Let (Pat [pe]) _ (Op scrm)) <- map_stm,+    (Screma _ [a1, a2] (ScremaForm [] [] map_lam)) <- scrm,+    (a1 == paramName pa1 && a2 == paramName pa2) || (a1 == paramName pa2 && a2 == paramName pa1),+    r == Var (patElemName pe) =+      Just map_lam+mapOp _ = Nothing++-- getting the innermost lambda of a perfect-map nest+--   (i.e., the first lambda that does not consists of exactly a map)+nestedMapOp :: Lambda SOACS -> Lambda SOACS+nestedMapOp lam =+  maybe lam nestedMapOp (mapOp lam)

@@ -1,39 +1,76 @@-module Futhark.AD.Rev.Scan (diffScan) where+module Futhark.AD.Rev.Scan (diffScan, diffScanVec, diffScanAdd) where  import Control.Monad+import Data.List (transpose) import Futhark.AD.Rev.Monad import Futhark.Analysis.PrimExp.Convert import Futhark.Builder import Futhark.IR.SOACS+import Futhark.IR.SOACS.Simplify (simplifyLambda) import Futhark.Tools import Futhark.Transform.Rename-import Futhark.Util (pairs, unpairs)+import Futhark.Util (chunk)  data FirstOrSecond = WrtFirst | WrtSecond +identityM :: Int -> Type -> ADM [[SubExp]]+identityM n t =+  traverse+    (traverse (letSubExp "id"))+    [[if i == j then oneExp t else zeroExp t | i <- [1 .. n]] | j <- [1 .. n]]++matrixMul :: [[PrimExp VName]] -> [[PrimExp VName]] -> PrimType -> [[PrimExp VName]]+matrixMul m1 m2 t =+  let zero = primExpFromSubExp t $ Constant $ blankPrimValue t+   in [[foldl (~+~) zero $ zipWith (~*~) r q | q <- transpose m2] | r <- m1]++matrixVecMul :: [[PrimExp VName]] -> [PrimExp VName] -> PrimType -> [PrimExp VName]+matrixVecMul m v t =+  let zero = primExpFromSubExp t $ Constant $ blankPrimValue t+   in [foldl (~+~) zero $ zipWith (~*~) v r | r <- m]++vectorAdd :: [PrimExp VName] -> [PrimExp VName] -> [PrimExp VName]+vectorAdd = zipWith (~+~)++orderArgs :: Special -> [a] -> [[a]]+orderArgs s lst =+  let d = div (length lst) $ specialScans s+   in chunk d lst+ -- computes `d(x op y)/dx` or d(x op y)/dy-mkScanAdjointLam :: VjpOps -> Lambda SOACS -> FirstOrSecond -> ADM (Lambda SOACS)-mkScanAdjointLam ops lam0 which = do+mkScanAdjointLam :: VjpOps -> Lambda SOACS -> FirstOrSecond -> [SubExp] -> ADM (Lambda SOACS)+mkScanAdjointLam ops lam0 which adjs = do   let len = length $ lambdaReturnType lam0   lam <- renameLambda lam0   let p2diff =         case which of           WrtFirst -> take len $ lambdaParams lam           WrtSecond -> drop len $ lambdaParams lam-  p_adjs <- mapM unitAdjOfType (lambdaReturnType lam)-  vjpLambda ops p_adjs (map paramName p2diff) lam+  vjpLambda ops (fmap AdjVal adjs) (map paramName p2diff) lam  -- Should generate something like: -- `\ j -> let i = n - 1 - j---         if i < n-1 then ( ys_adj[i], df2dx ys[i] xs[i+1]) else (0,1) )`+--         if i < n-1 then ( ys_adj[i], df2dx ys[i] xs[i+1]) else (ys_adj[i],1) )` -- where `ys` is  the result of scan --       `xs` is  the input  of scan --       `ys_adj` is the known adjoint of ys --       `j` draw values from `iota n`-mkScanFusedMapLam :: VjpOps -> SubExp -> Lambda SOACS -> [VName] -> [VName] -> [VName] -> ADM (Lambda SOACS)-mkScanFusedMapLam ops w scn_lam xs ys ys_adj = do-  lam <- mkScanAdjointLam ops scn_lam WrtFirst-  ys_ts <- mapM lookupType ys+mkScanFusedMapLam ::+  VjpOps ->+  SubExp ->+  Lambda SOACS ->+  [VName] ->+  [VName] ->+  [VName] ->+  Special ->+  Int ->+  ADM (Lambda SOACS)+mkScanFusedMapLam ops w scn_lam xs ys ys_adj s d = do+  let sc = specialCase s+  let k = specialSubSize s+  ys_ts <- traverse lookupType ys+  idmat <- identityM (length ys) $ rowType $ head ys_ts+  lams <- traverse (mkScanAdjointLam ops scn_lam WrtFirst) idmat   par_i <- newParam "i" $ Prim int64   let i = paramName par_i   mkLambda [par_i] $@@ -41,110 +78,268 @@       =<< eIf         (toExp $ le64 i .==. 0)         ( buildBody_ $ do-            zs <- mapM (letSubExp "ct_zero" . zeroExp . rowType) ys_ts-            os <- mapM (letSubExp "ct_one" . oneExp . rowType) ys_ts-            pure $ subExpsRes $ unpairs $ zip zs os+            j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))+            y_s <- forM ys_adj $ \y_ ->+              letSubExp (baseString y_ ++ "_j") =<< eIndex y_ (eSubExp j)+            let zso = orderArgs s y_s+            let ido = orderArgs s $ case_jac k sc idmat+            pure $ subExpsRes $ concat $ zipWith (++) zso $ fmap concat ido         )         ( buildBody_ $ do             j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))             j1 <- letSubExp "j1" =<< toExp (pe64 w - le64 i)-            let index idx arr t = BasicOp $ Index arr $ fullSlice t [DimFix idx]-            y_s <- forM (zip ys_adj ys_ts) $ \(y_, t) ->-              letSubExp (baseString y_ ++ "_j") $ index j y_ t-            lam_rs <--              eLambda lam . map pure $-                zipWith (index j) ys ys_ts ++ zipWith (index j1) xs ys_ts-            pure $ unpairs $ zip (subExpsRes y_s) lam_rs+            y_s <- forM ys_adj $ \y_ ->+              letSubExp (baseString y_ ++ "_j") =<< eIndex y_ (eSubExp j)++            let args =+                  map (`eIndex` eSubExp j) ys ++ map (`eIndex` eSubExp j1) xs+            lam_rs <- traverse (`eLambda` args) lams++            let yso = orderArgs s $ subExpsRes y_s+            let jaco = orderArgs s $ case_jac k sc $ transpose lam_rs++            pure $ concat $ zipWith (++) yso $ fmap concat jaco         )+  where+    case_jac :: Int -> SpecialCase -> [[a]] -> [[a]]+    case_jac _ Generic jac = jac+    case_jac k ZeroQuadrant jac =+      concat+        $ zipWith+          (\i -> map (take k . drop (i * k)))+          [0 .. d `div` k]+        $ chunk k jac+    case_jac k MatrixMul jac =+      take k <$> take k jac --- \(a1, b1) (a2, b2) -> (a2 + b2 * a1, b1 * b2)-mkScanLinFunO :: Type -> ADM (Scan SOACS)-mkScanLinFunO t = do+-- a1 a2 b -> a2 + b * a1+linFunT0 :: [PrimExp VName] -> [PrimExp VName] -> [[PrimExp VName]] -> Special -> PrimType -> [PrimExp VName]+linFunT0 a1 a2 b s pt =+  let t = case specialCase s of+        MatrixMul ->+          concatMap (\v -> matrixVecMul b v pt) $ chunk (specialSubSize s) a1+        _ -> matrixVecMul b a1 pt+   in a2 `vectorAdd` t++-- \(a1, b1) (a2, b2) -> (a2 + b2 * a1, b2 * b1)+mkScanLinFunO :: Type -> Special -> ADM (Scan SOACS)+mkScanLinFunO t s = do   let pt = elemType t-  zero <- letSubExp "zeros" $ zeroExp t-  one <- letSubExp "ones" $ oneExp t-  tmp <- mapM newVName ["a1", "b1", "a2", "b2"]-  let [a1, b1, a2, b2] = tmp-      pet = primExpFromSubExp pt . Var-  lam <- mkLambda (map (\v -> Param mempty v t) [a1, b1, a2, b2]) . fmap varsRes $-    tabNest (arrayRank t) [a1, b1, a2, b2] $ \_ [a1', b1', a2', b2'] -> do-      x <- letExp "x" <=< toExp $ pet a2' ~+~ pet b2' ~*~ pet a1'-      y <- letExp "y" <=< toExp $ pet b1' ~*~ pet b2'-      pure [x, y]-  pure $ Scan lam [zero, one]+  neu_elm <- mkNeutral $ specialNeutral s+  let (as, bs) = specialParams s+  (a1s, b1s, a2s, b2s) <- mkParams (as, bs)+  let pet = primExpFromSubExp pt . Var+  let (_, n) = specialNeutral s --- build the map following the scan with linear-function-composition:--- for each (ds,cs) length-n array results of the scan, combine them as:---    `let rs = map2 (\ d_i c_i -> d_i + c_i * y_adj[n-1]) d c |> reverse`--- but insert explicit indexing to reverse inside the map.-mkScan2ndMaps :: SubExp -> (Type, VName, (VName, VName)) -> ADM VName-mkScan2ndMaps w (arr_tp, y_adj, (ds, cs)) = do-  y_adj_last <- letExp (baseString y_adj <> "_last") =<< eLast y_adj+  lam <- mkLambda (map (\v -> Param mempty v (rowType t)) (a1s ++ b1s ++ a2s ++ b2s)) . fmap subExpsRes $ do+    let [a1s', b1s', a2s', b2s'] = (fmap . fmap) pet [a1s, b1s, a2s, b2s]+    let (b1sm, b2sm) = (chunk n b1s', chunk n b2s') -  par_i <- newParam "i" $ Prim int64-  lam <- mkLambda [par_i] $ do-    let i = paramName par_i-    j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))-    dj <- letExp (baseString ds ++ "_dj") $ BasicOp $ Index ds $ fullSlice arr_tp [DimFix j]-    cj <- letExp (baseString cs ++ "_cj") $ BasicOp $ Index cs $ fullSlice arr_tp [DimFix j]+    let t0 = linFunT0 a1s' a2s' b2sm s pt+    let t1 = concat $ matrixMul b2sm b1sm pt+    traverse (letSubExp "r" <=< toExp) $ t0 ++ t1 -    let pet = primExpFromSubExp (elemType arr_tp) . Var-    fmap varsRes . tabNest (arrayRank (rowType arr_tp)) [y_adj_last, dj, cj] $ \_ [y_adj_last', dj', cj'] ->-      letTupExp "res" <=< toExp $ pet dj' ~+~ pet cj' ~*~ pet y_adj_last'+  pure $ Scan lam neu_elm+  where+    mkNeutral (a, b) = do+      zeros <- replicateM a $ letSubExp "zeros" $ zeroExp $ rowType t+      idmat <- identityM b $ Prim $ elemType t+      pure $ zeros ++ concat idmat -  iota <- letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64-  letExp "after_scan" $ Op (Screma w [iota] (ScremaForm [] [] lam))+    mkParams (a, b) = do+      a1s <- replicateM a $ newVName "a1"+      b1s <- replicateM b $ newVName "b1"+      a2s <- replicateM a $ newVName "a2"+      b2s <- replicateM b $ newVName "b2"+      pure (a1s, b1s, a2s, b2s) --- perform the final map, which is fusable with the maps obtained from `mkScan2ndMaps`+-- perform the final map -- let xs_contribs = --    map3 (\ i a r -> if i==0 then r else (df2dy (ys[i-1]) a) \bar{*} r)---         (iota n) xs rs+--         (iota n) xs (reverse ds) mkScanFinalMap :: VjpOps -> SubExp -> Lambda SOACS -> [VName] -> [VName] -> [VName] -> ADM [VName]-mkScanFinalMap ops w scan_lam xs ys rs = do+mkScanFinalMap ops w scan_lam xs ys ds = do   let eltps = lambdaReturnType scan_lam-  lam <- mkScanAdjointLam ops scan_lam WrtSecond+   par_i <- newParam "i" $ Prim int64   let i = paramName par_i-  par_x <- mapM (\(x, t) -> newParam (baseString x ++ "_par_x") t) $ zip xs eltps-  par_r <- mapM (\(r, t) -> newParam (baseString r ++ "_par_r") t) $ zip rs eltps+  par_x <- zipWithM (\x -> newParam (baseString x ++ "_par_x")) xs eltps    map_lam <--    mkLambda (par_i : par_x ++ par_r) $+    mkLambda (par_i : par_x) $ do+      j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))++      dj <-+        traverse+          (\dd -> letExp (baseString dd ++ "_dj") =<< eIndex dd (eSubExp j))+          ds+       fmap varsRes . letTupExp "scan_contribs"         =<< eIf           (toExp $ le64 i .==. 0)-          (resultBodyM $ map (Var . paramName) par_r)+          (resultBodyM $ fmap Var dj)           ( buildBody_ $ do-              im1 <- letSubExp "im1" =<< toExp (le64 i - 1)-              ys_im1 <- forM ys $ \y ->-                letSubExp (baseString y <> "_im1") =<< eIndex y (eSubExp im1)--              lam_res <--                mapM (letExp "const" . BasicOp . SubExp . resSubExp)-                  =<< eLambda lam (map eSubExp $ ys_im1 ++ map (Var . paramName) par_x)+              lam <- mkScanAdjointLam ops scan_lam WrtSecond $ fmap Var dj -              fmap (varsRes . mconcat) . forM (zip3 lam_res (map paramName par_r) eltps) $-                \(lam_r, r, eltp) -> do-                  let pet = primExpFromSubExp (elemType eltp) . Var+              im1 <- letSubExp "im1" =<< toExp (le64 i - 1)+              ys_im1 <- forM ys $ \y -> do+                y_t <- lookupType y+                letSubExp (baseString y ++ "_last") $ BasicOp $ Index y $ fullSlice y_t [DimFix im1] -                  tabNest (arrayRank eltp) [lam_r, r] $ \_ [lam_r', r'] ->-                    letTupExp "res" <=< toExp $ pet lam_r' ~*~ pet r'+              let args = map eSubExp $ ys_im1 ++ map (Var . paramName) par_x+              eLambda lam args           )    iota <- letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64-  letTupExp "scan_contribs" $ Op (Screma w (iota : xs ++ rs) (ScremaForm [] [] map_lam))+  letTupExp "scan_contribs" $ Op $ Screma w (iota : xs) $ mapSOAC map_lam +data SpecialCase+  = Generic+  | ZeroQuadrant+  | MatrixMul+  deriving (Show)++data Special = Special+  { specialNeutral :: (Int, Int),+    specialParams :: (Int, Int),+    specialScans :: Int,+    specialSubSize :: Int,+    specialCase :: SpecialCase+  }+  deriving (Show)++subMats :: Int -> [[Exp SOACS]] -> Exp SOACS -> Maybe Int+subMats d mat zero =+  let sub_d = filter (\x -> d `mod` x == 0) [1 .. (d `div` 2)]+      poss = map (\m -> all (ok m) $ zip mat [0 .. d - 1]) sub_d+      tmp = filter fst (zip poss sub_d)+   in if null tmp then Nothing else Just $ snd $ head tmp+  where+    ok m (row, i) =+      all (\(v, j) -> v == zero || (i `div` m == j `div` m)) $+        zip row [0 .. d - 1]++cases :: Int -> Type -> [[Exp SOACS]] -> Special+cases d t mat+  | Just k <- subMats d mat $ zeroExp t =+      let nonZeros = zipWith (\i -> map (take k . drop (i * k))) [0 .. d `div` k] $ chunk k mat+       in if all (== head nonZeros) $ tail nonZeros+            then Special (d, k) (d, k * k) 1 k MatrixMul+            else Special (k, k) (k, k * k) (d `div` k) k ZeroQuadrant+cases d _ _ = Special (d, d) (d, d * d) 1 d Generic++identifyCase :: VjpOps -> Lambda SOACS -> ADM Special+identifyCase ops lam = do+  let t = lambdaReturnType lam+  let d = length t+  idmat <- identityM d $ head t+  lams <- traverse (mkScanAdjointLam ops lam WrtFirst) idmat++  par1 <- traverse (newParam "tmp1") t+  par2 <- traverse (newParam "tmp2") t+  jac_lam <- mkLambda (par1 ++ par2) $ do+    let args = fmap eParam $ par1 ++ par2+    lam_rs <- traverse (`eLambda` args) lams++    pure $ concat (transpose lam_rs)++  simp <- simplifyLambda jac_lam+  let jac = chunk d $ fmap (BasicOp . SubExp . resSubExp) $ bodyResult $ lambdaBody simp+  pure $ cases d (head t) jac+ diffScan :: VjpOps -> [VName] -> SubExp -> [VName] -> Scan SOACS -> ADM () diffScan ops ys w as scan = do+  sc <- identifyCase ops (scanLambda scan)+  let d = length as   ys_adj <- mapM lookupAdjVal ys   as_ts <- mapM lookupType as-  map1_lam <- mkScanFusedMapLam ops w (scanLambda scan) as ys ys_adj-  scans_lin_fun_o <- mapM mkScanLinFunO $ lambdaReturnType $ scanLambda scan+  map1_lam <- mkScanFusedMapLam ops w (scanLambda scan) as ys ys_adj sc d+  scans_lin_fun_o <- mkScanLinFunO (head as_ts) sc+  scan_lams <- mkScans (specialScans sc) scans_lin_fun_o   iota <-     letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64   r_scan <--    letTupExp "adj_ctrb_scan" $-      Op (Screma w [iota] (ScremaForm scans_lin_fun_o [] map1_lam))-  red_nms <- mapM (mkScan2ndMaps w) (zip3 as_ts ys_adj (pairs r_scan))-  as_contribs <- mkScanFinalMap ops w (scanLambda scan) as ys red_nms+    letTupExp "adj_ctrb_scan" . Op . Screma w [iota] $+      scanomapSOAC scan_lams map1_lam++  as_contribs <- mkScanFinalMap ops w (scanLambda scan) as ys (splitScanRes sc r_scan d)   zipWithM_ updateAdj as as_contribs+  where+    mkScans :: Int -> Scan SOACS -> ADM [Scan SOACS]+    mkScans d s =+      replicateM d $ do+        lam' <- renameLambda $ scanLambda s+        pure $ Scan lam' $ scanNeutral s++    splitScanRes sc res d =+      concat $ take (div d $ specialScans sc) <$> orderArgs sc res++diffScanVec ::+  VjpOps ->+  [VName] ->+  StmAux () ->+  SubExp ->+  Lambda SOACS ->+  [SubExp] ->+  [VName] ->+  ADM () ->+  ADM ()+diffScanVec ops ys aux w lam ne as m = do+  stmts <- collectStms_ $ do+    rank <- arrayRank <$> lookupType (head as)+    let rear = [1, 0] ++ drop 2 [0 .. rank - 1]++    transp_as <-+      traverse+        (\a -> letExp (baseString a ++ "_transp") $ BasicOp $ Rearrange rear a)+        as++    ts <- traverse lookupType transp_as+    let n = arraysSize 0 ts++    as_par <- traverse (newParam "as_par" . rowType) ts+    ne_par <- traverse (newParam "ne_par") $ lambdaReturnType lam++    scan_form <- scanSOAC [Scan lam (map (Var . paramName) ne_par)]++    map_lam <-+      mkLambda (as_par ++ ne_par) . fmap varsRes . letTupExp "map_res" . Op $+        Screma w (map paramName as_par) scan_form++    transp_ys <-+      letTupExp "trans_ys" . Op $+        Screma n (transp_as ++ subExpVars ne) (mapSOAC map_lam)++    zipWithM+      (\y x -> auxing aux $ letBindNames [y] $ BasicOp $ Rearrange rear x)+      ys+      transp_ys++  foldr (vjpStm ops) m stmts++diffScanAdd :: VjpOps -> VName -> SubExp -> Lambda SOACS -> SubExp -> VName -> ADM ()+diffScanAdd _ops ys n lam' ne as = do+  lam <- renameLambda lam'+  ys_bar <- lookupAdjVal ys++  map_scan <- rev_arr_lam ys_bar++  iota <-+    letExp "iota" $ BasicOp $ Iota n (intConst Int64 0) (intConst Int64 1) Int64++  scan_res <-+    letExp "res_rev" $ Op $ Screma n [iota] $ scanomapSOAC [Scan lam [ne]] map_scan++  rev_lam <- rev_arr_lam scan_res+  contrb <- letExp "contrb" $ Op $ Screma n [iota] $ mapSOAC rev_lam++  updateAdj as contrb+  where+    rev_arr_lam :: VName -> ADM (Lambda SOACS)+    rev_arr_lam arr = do+      par_i <- newParam "i" $ Prim int64+      mkLambda [par_i] $ do+        a <-+          letExp "ys_bar_rev"+            =<< eIndex arr (toExp (pe64 n - le64 (paramName par_i) - 1))+        pure [varRes a]

@@ -4,6 +4,7 @@   ( printAction,     printAliasesAction,     printLastUseGPU,+    printLastUseGPUSS,     printFusionGraph,     printInterferenceGPU,     printMemAliasGPU,@@ -26,7 +27,9 @@  import Control.Monad import Control.Monad.IO.Class+import Data.Bifunctor import Data.List (intercalate)+import Data.Map qualified as M import Data.Maybe (fromMaybe) import Data.Text qualified as T import Data.Text.IO qualified as T@@ -55,6 +58,7 @@ import Futhark.IR.Prop.Aliases import Futhark.IR.SOACS (SOACS) import Futhark.IR.SeqMem (SeqMem)+import Futhark.Optimise.ArrayShortCircuiting.LastUse qualified as LastUseSS import Futhark.Optimise.Fusion.GraphRep qualified import Futhark.Util (runProgramWithExitCode, unixEnvironment) import Futhark.Version (versionString)@@ -87,9 +91,24 @@   Action     { actionName = "print last use gpu",       actionDescription = "Print last use information on gpu.",-      actionProcedure = liftIO . print . LastUse.analyseGPUMem+      actionProcedure = liftIO . putStrLn . prettyString . M.toList . LastUse.analyseGPUMem     } +-- | Print last use information to stdout.+printLastUseGPUSS :: Action GPUMem+printLastUseGPUSS =+  Action+    { actionName = "print last use ss gpu",+      actionDescription = "Print last use ss information on gpu.",+      actionProcedure =+        liftIO+          . putStrLn+          . prettyString+          . bimap M.toList (M.toList . fmap M.toList)+          . LastUseSS.lastUseGPUMem+          . aliasAnalysis+    }+ -- | Print fusion graph to stdout. printFusionGraph :: Action SOACS printFusionGraph =@@ -225,7 +244,7 @@   ret_ispc <-     liftIO $       runProgramWithExitCode-        "ispc"+        cmdISPC         ( [ispcpath, "-o", ispcbase `addExtension` "o"]             ++ ["--addressing=64", "--pic"]             ++ cmdISPCFLAGS ispc_flags -- These flags are always needed@@ -245,24 +264,25 @@         mempty   case ret_ispc of     Left err ->-      externalErrorS $ "Failed to run " ++ cmdCC ++ ": " ++ show err-    Right (ExitFailure code, _, gccerr) -> throwError code gccerr+      externalErrorS $ "Failed to run " ++ cmdISPC ++ ": " ++ show err+    Right (ExitFailure code, _, ispcerr) -> throwError cmdISPC code ispcerr     Right (ExitSuccess, _, _) ->       case ret of         Left err ->           externalErrorS $ "Failed to run ispc: " ++ show err-        Right (ExitFailure code, _, gccerr) -> throwError code gccerr+        Right (ExitFailure code, _, gccerr) -> throwError cmdCC code gccerr         Right (ExitSuccess, _, _) ->           pure ()   where+    cmdISPC = "ispc"     ispcbase = outpath <> ispcextension-    throwError code gccerr =+    throwError prog code err =       externalErrorS $-        cmdCC+        prog           ++ " failed with code "           ++ show code           ++ ":\n"-          ++ gccerr+          ++ err  -- | The @futhark c@ action. compileCAction :: FutharkConfig -> CompilerMode -> FilePath -> Action SeqMem

@@ -214,7 +214,7 @@   applyAliases (MemAlias.analyzeGPUMem prog) $     onConsts (progConsts prog) <> foldMap onFun (progFuns prog)   where-    (lumap, _) = LastUse.analyseGPUMem prog+    lumap = LastUse.analyseGPUMem prog     onFun f =       runReader (analyseGPU lumap $ bodyStms $ funDefBody f) $ scopeOf f     onConsts stms =

@@ -41,10 +41,10 @@  type LastUseM rep = Reader (Env rep) (LastUse, Used) -analyseGPUMem :: Prog GPUMem -> (LastUseMap, Used)+analyseGPUMem :: Prog GPUMem -> LastUseMap analyseGPUMem = analyseProg analyseGPUOp -analyseSeqMem :: Prog SeqMem -> (LastUseMap, Used)+analyseSeqMem :: Prog SeqMem -> LastUseMap analyseSeqMem = analyseProg analyseSeqOp  analyseGPUOp :: LastUseOp GPUMem@@ -100,7 +100,7 @@ -- | Analyses a program to return a last-use map, mapping each simple statement -- in the program to the values that were last used within that statement, and -- the set of all `VName` that were used inside.-analyseProg :: (CanBeAliased (Op rep), Mem rep inner) => LastUseOp rep -> Prog rep -> (LastUseMap, Used)+analyseProg :: (CanBeAliased (Op rep), Mem rep inner) => LastUseOp rep -> Prog rep -> LastUseMap analyseProg onOp prog =   runReader helper (Env onOp)   where@@ -112,9 +112,9 @@           prog_alias = aliasAnalysis prog           consts = progConsts prog_alias           funs = progFuns prog_alias-      (consts_lu, consts_used) <- analyseStms mempty mempty consts-      (lus, used) <- mconcat <$> mapM (analyseFun mempty bound_in_consts) funs-      pure (flipMap $ consts_lu <> lus, consts_used <> used)+      (consts_lu, _) <- analyseStms mempty mempty consts+      (lus, _) <- mconcat <$> mapM (analyseFun mempty bound_in_consts) funs+      pure $ flipMap $ consts_lu <> lus  analyseFun :: (FreeIn (OpWithAliases (Op rep)), ASTRep rep) => LastUse -> Used -> FunDef (Aliases rep) -> LastUseM rep analyseFun lumap used fun = do

@@ -731,13 +731,13 @@   S.unions $ map leafExpTypes pes  -- | Multiplication of untyped 'PrimExp's, which must have the same--- type.+-- type.  Uses 'OverflowWrap' for integer operations. (~*~) :: PrimExp v -> PrimExp v -> PrimExp v x ~*~ y = BinOpExp op x y   where     t = primExpType x     op = case t of-      IntType it -> Mul it OverflowUndef+      IntType it -> Mul it OverflowWrap       FloatType ft -> FMul ft       Bool -> LogAnd       Unit -> LogAnd@@ -755,23 +755,25 @@       Unit -> LogAnd  -- | Addition of untyped 'PrimExp's, which must have the same type.+-- Uses 'OverflowWrap' for integer operations. (~+~) :: PrimExp v -> PrimExp v -> PrimExp v x ~+~ y = BinOpExp op x y   where     t = primExpType x     op = case t of-      IntType it -> Add it OverflowUndef+      IntType it -> Add it OverflowWrap       FloatType ft -> FAdd ft       Bool -> LogOr       Unit -> LogOr  -- | Subtraction of untyped 'PrimExp's, which must have the same type.+-- Uses 'OverflowWrap' for integer operations. (~-~) :: PrimExp v -> PrimExp v -> PrimExp v x ~-~ y = BinOpExp op x y   where     t = primExpType x     op = case t of-      IntType it -> Sub it OverflowUndef+      IntType it -> Sub it OverflowWrap       FloatType ft -> FSub ft       Bool -> LogOr       Unit -> LogOr

@@ -44,6 +44,7 @@ import Futhark.Pass.AD import Futhark.Pass.ExpandAllocations import Futhark.Pass.ExplicitAllocations.GPU qualified as GPU+import Futhark.Pass.ExplicitAllocations.MC qualified as MC import Futhark.Pass.ExplicitAllocations.Seq qualified as Seq import Futhark.Pass.ExtractKernels import Futhark.Pass.ExtractMulticore@@ -229,6 +230,13 @@   externalErrorS $     "Pass " ++ name ++ " expects SeqMem representation, but got " ++ representation rep +mcMemProg :: String -> UntypedPassState -> FutharkM (Prog MCMem.MCMem)+mcMemProg _ (MCMem prog) =+  pure prog+mcMemProg name rep =+  externalErrorS $+    "Pass " ++ name ++ " expects MCMem representation, but got " ++ representation rep+ typedPassOption ::   Checkable torep =>   (String -> UntypedPassState -> FutharkM (Prog fromrep)) ->@@ -263,6 +271,13 @@ seqMemPassOption =   typedPassOption seqMemProg SeqMem +mcMemPassOption ::+  Pass MCMem.MCMem MCMem.MCMem ->+  String ->+  FutharkOption+mcMemPassOption =+  typedPassOption mcMemProg MCMem+ kernelsMemPassOption ::   Pass GPUMem.GPUMem GPUMem.GPUMem ->   String ->@@ -302,6 +317,9 @@     perform (Seq prog) config =       SeqMem         <$> runPipeline (onePass Seq.explicitAllocations) config prog+    perform (MC prog) config =+      MCMem+        <$> runPipeline (onePass MC.explicitAllocations) config prog     perform s _ =       externalErrorS $         "Pass '" ++ passDescription pass ++ "' cannot operate on " ++ representation s@@ -492,6 +510,14 @@       "Print last use information.",     Option       []+      ["print-last-use-gpu-ss"]+      ( NoArg $+          Right $ \opts ->+            opts {futharkAction = GPUMemAction $ \_ _ _ -> printLastUseGPUSS}+      )+      "Print last use information ss.",+    Option+      []       ["print-interference-gpu"]       ( NoArg $           Right $ \opts ->@@ -599,6 +625,7 @@     seqMemPassOption LowerAllocations.lowerAllocationsSeqMem [],     kernelsMemPassOption LowerAllocations.lowerAllocationsGPUMem [],     seqMemPassOption ArrayShortCircuiting.optimiseSeqMem [],+    mcMemPassOption ArrayShortCircuiting.optimiseMCMem [],     kernelsMemPassOption ArrayShortCircuiting.optimiseGPUMem [],     cseOption [],     simplifyOption "e",@@ -641,9 +668,17 @@       ["seq-mem"],     pipelineOption       getSOACSProg+      "MC"+      MC+      "Run the multicore compilation pipeline"+      mcPipeline+      []+      ["mc"],+    pipelineOption+      getSOACSProg       "MCMem"       MCMem-      "Run the multicore compilation pipeline"+      "Run the multicore+memory compilation pipeline"       multicorePipeline       []       ["mc-mem"]

@@ -138,10 +138,14 @@       check ["-s"]     GpuPipeline ->       check ["--gpu"]+    MCPipeline ->+      check ["--mc"]     SeqMemPipeline ->       check ["--seq-mem"]     GpuMemPipeline ->       check ["--gpu-mem"]+    MCMemPipeline ->+      check ["--mc-mem"]     NoPipeline ->       check []   where@@ -165,9 +169,11 @@   where     maybePipeline :: StructurePipeline -> T.Text     maybePipeline SOACSPipeline = "(soacs) "-    maybePipeline GpuPipeline = "(kernels) "+    maybePipeline GpuPipeline = "(gpu) "+    maybePipeline MCPipeline = "(mc) "     maybePipeline SeqMemPipeline = "(seq-mem) "     maybePipeline GpuMemPipeline = "(gpu-mem) "+    maybePipeline MCMemPipeline = "(mc-mem) "     maybePipeline NoPipeline = ""      ok (AstMetrics metrics) (name, expected_occurences) =

@@ -279,17 +279,13 @@ cudaMemoryType space =   error $ "CUDA backend does not support '" ++ space ++ "' memory space." -callKernel :: GC.OpCompiler OpenCL ()-callKernel (GetSize v key) =-  GC.stm [C.cstm|$id:v = *ctx->tuning_params.$id:key;|]-callKernel (CmpSizeLe v key x) = do-  x' <- GC.compileExp x-  GC.stm [C.cstm|$id:v = *ctx->tuning_params.$id:key <= $exp:x';|]-  sizeLoggingCode v key x'-callKernel (GetSizeMax v size_class) =-  let field = "max_" ++ cudaSizeClass size_class-   in GC.stm [C.cstm|$id:v = ctx->cuda.$id:field;|]+kernelConstToExp :: KernelConst -> C.Exp+kernelConstToExp (SizeConst key) =+  [C.cexp|*ctx->tuning_params.$id:key|]+kernelConstToExp (SizeMaxConst size_class) =+  [C.cexp|ctx->cuda.$id:field|]   where+    field = "max_" <> cudaSizeClass size_class     cudaSizeClass SizeThreshold {} = "threshold"     cudaSizeClass SizeGroup = "block_size"     cudaSizeClass SizeNumGroups = "grid_size"@@ -297,6 +293,23 @@     cudaSizeClass SizeRegTile = "reg_tile_size"     cudaSizeClass SizeLocalMemory = "shared_memory"     cudaSizeClass (SizeBespoke x _) = prettyString x++compileGroupDim :: GroupDim -> GC.CompilerM op s C.Exp+compileGroupDim (Left e) = GC.compileExp e+compileGroupDim (Right kc) = pure $ kernelConstToExp kc++callKernel :: GC.OpCompiler OpenCL ()+callKernel (GetSize v key) = do+  let e = kernelConstToExp $ SizeConst key+  GC.stm [C.cstm|$id:v = $exp:e;|]+callKernel (CmpSizeLe v key x) = do+  let e = kernelConstToExp $ SizeConst key+  x' <- GC.compileExp x+  GC.stm [C.cstm|$id:v = $exp:e <= $exp:x';|]+  sizeLoggingCode v key x'+callKernel (GetSizeMax v size_class) = do+  let e = kernelConstToExp $ SizeMaxConst size_class+  GC.stm [C.cstm|$id:v = $exp:e;|] callKernel (LaunchKernel safety kernel_name args num_blocks block_size) = do   args_arr <- newVName "kernel_args"   time_start <- newVName "time_start"@@ -310,7 +323,7 @@     GC.decl [C.cdecl|unsigned int $id:arg = $exp:offset;|]    (grid_x, grid_y, grid_z) <- mkDims <$> mapM GC.compileExp num_blocks-  (block_x, block_y, block_z) <- mkDims <$> mapM GC.compileExp block_size+  (block_x, block_y, block_z) <- mkDims <$> mapM compileGroupDim block_size   let perm_args         | length num_blocks == 3 = [[C.cinit|&perm[0]|], [C.cinit|&perm[1]|], [C.cinit|&perm[2]|]]         | otherwise = []

@@ -326,16 +326,30 @@ staticOpenCLArray _ space _ _ =   error $ "OpenCL backend cannot create static array in memory space '" ++ space ++ "'" +kernelConstToExp :: KernelConst -> C.Exp+kernelConstToExp (SizeConst key) =+  [C.cexp|*ctx->tuning_params.$id:key|]+kernelConstToExp (SizeMaxConst size_class) =+  [C.cexp|ctx->opencl.$id:field|]+  where+    field = "max_" <> prettyString size_class++compileGroupDim :: GroupDim -> GC.CompilerM op s C.Exp+compileGroupDim (Left e) = GC.compileExp e+compileGroupDim (Right kc) = pure $ kernelConstToExp kc+ callKernel :: GC.OpCompiler OpenCL ()-callKernel (GetSize v key) =-  GC.stm [C.cstm|$id:v = *ctx->tuning_params.$id:key;|]+callKernel (GetSize v key) = do+  let e = kernelConstToExp $ SizeConst key+  GC.stm [C.cstm|$id:v = $exp:e;|] callKernel (CmpSizeLe v key x) = do+  let e = kernelConstToExp $ SizeConst key   x' <- GC.compileExp x-  GC.stm [C.cstm|$id:v = *ctx->tuning_params.$id:key <= $exp:x';|]+  GC.stm [C.cstm|$id:v = $exp:e <= $exp:x';|]   sizeLoggingCode v key x'-callKernel (GetSizeMax v size_class) =-  let field = "max_" ++ prettyString size_class-   in GC.stm [C.cstm|$id:v = ctx->opencl.$id:field;|]+callKernel (GetSizeMax v size_class) = do+  let e = kernelConstToExp $ SizeMaxConst size_class+  GC.stm [C.cstm|$id:v = $exp:e;|] callKernel (LaunchKernel safety name args num_workgroups workgroup_size) = do   -- The other failure args are set automatically when the kernel is   -- first created.@@ -349,7 +363,7 @@    zipWithM_ setKernelArg [numFailureParams safety ..] args   num_workgroups' <- mapM GC.compileExp num_workgroups-  workgroup_size' <- mapM GC.compileExp workgroup_size+  workgroup_size' <- mapM compileGroupDim workgroup_size   local_bytes <- foldM localBytes [C.cexp|0|] args    launchKernel name num_workgroups' workgroup_size' local_bytes@@ -416,10 +430,11 @@         fprintf(ctx->log, $string:debug_str, $args:debug_args);         $id:time_start = get_wall_time();       }+      typename cl_event *pevent = $exp:(profilingEvent kernel_name);       OPENCL_SUCCEED_OR_RETURN(         clEnqueueNDRangeKernel(ctx->opencl.queue, ctx->$id:kernel_name, $int:kernel_rank, NULL,                                $id:global_work_size, $id:local_work_size,-                               0, NULL, $exp:(profilingEvent kernel_name)));+                               0, NULL, pevent));       if (ctx->debugging) {         OPENCL_SUCCEED_FATAL(clFinish(ctx->opencl.queue));         $id:time_end = get_wall_time();

@@ -7,6 +7,7 @@     definitionsText,     typeText,     idText,+    funcText,     funcsText,   ) where@@ -29,5 +30,8 @@ idText :: C.Id -> T.Text idText = T.pack . MPP.pretty 8000 . MPP.ppr +funcText :: C.Func -> T.Text+funcText = T.pack . MPP.pretty 8000 . MPP.ppr+ funcsText :: [C.Func] -> T.Text-funcsText = T.unlines . map (T.pack . MPP.pretty 8000 . MPP.ppr)+funcsText = T.unlines . map funcText

@@ -598,16 +598,26 @@ compileCode (Free name space) = do   cached <- isJust <$> GC.cacheMem name   unless cached $ unRefMem name space-compileCode (For i bound body) = do-  let i' = C.toIdent i-      t = GC.primTypeToCType $ primExpType bound-  bound' <- compileExp bound-  body' <- GC.collect $ compileCode body-  quals <- getVariabilityQuals i-  GC.stm-    [C.cstm|for ($tyquals:quals $ty:t $id:i' = 0; $id:i' < $exp:bound'; $id:i'++) {+compileCode (For i bound body)+  -- The special-case here is to avoid certain pathological/contrived+  -- programs that construct statically known zero-element arrays.+  -- Due to the way we do constant-fold index functions, this produces+  -- code that looks like it has uniform/varying mismatches (i.e. race+  -- conditions) to ISPC, even though that code is never actually run.+  | isZero bound = pure ()+  | otherwise = do+      let i' = C.toIdent i+          t = GC.primTypeToCType $ primExpType bound+      bound' <- compileExp bound+      body' <- GC.collect $ compileCode body+      quals <- getVariabilityQuals i+      GC.stm+        [C.cstm|for ($tyquals:quals $ty:t $id:i' = 0; $id:i' < $exp:bound'; $id:i'++) {             $items:body'           }|]+  where+    isZero (ValueExp v) = zeroIsh v+    isZero _ = False compileCode (While cond body) = do   cond' <- compileExp $ untyped cond   body' <- GC.collect $ compileCode body@@ -937,6 +947,10 @@         $items:body'       }     }|]+compileOp (ExtractLane dest (ValueExp v) _) =+  -- extract() on constants is not allowed (type is uniform, not+  -- varying), so just turn them into an assignment.+  GC.stm [C.cstm|$id:dest = $exp:v;|] compileOp (ExtractLane dest tar lane) = do   tar' <- compileExp tar   lane' <- compileExp lane

@@ -199,27 +199,32 @@ asLong :: PyExp -> PyExp asLong x = Py.simpleCall "np.int64" [x] +kernelConstToExp :: Imp.KernelConst -> PyExp+kernelConstToExp (Imp.SizeConst key) =+  Index (Var "self.sizes") (IdxExp $ String $ prettyString key)+kernelConstToExp (Imp.SizeMaxConst size_class) =+  Var $ "self.max_" <> prettyString size_class++compileGroupDim :: Imp.GroupDim -> Py.CompilerM op s PyExp+compileGroupDim (Left e) = asLong <$> Py.compileExp e+compileGroupDim (Right kc) = pure $ kernelConstToExp kc+ callKernel :: Py.OpCompiler Imp.OpenCL () callKernel (Imp.GetSize v key) = do   v' <- Py.compileVar v-  Py.stm $-    Assign v' $-      Index (Var "self.sizes") (IdxExp $ String $ prettyString key)+  Py.stm $ Assign v' $ kernelConstToExp $ Imp.SizeConst key callKernel (Imp.CmpSizeLe v key x) = do   v' <- Py.compileVar v   x' <- Py.compileExp x   Py.stm $     Assign v' $-      BinOp "<=" (Index (Var "self.sizes") (IdxExp $ String $ prettyString key)) x'+      BinOp "<=" (kernelConstToExp (Imp.SizeConst key)) x' callKernel (Imp.GetSizeMax v size_class) = do   v' <- Py.compileVar v-  Py.stm $-    Assign v' $-      Var $-        "self.max_" ++ prettyString size_class+  Py.stm $ Assign v' $ kernelConstToExp $ Imp.SizeMaxConst size_class callKernel (Imp.LaunchKernel safety name args num_workgroups workgroup_size) = do   num_workgroups' <- mapM (fmap asLong . Py.compileExp) num_workgroups-  workgroup_size' <- mapM (fmap asLong . Py.compileExp) workgroup_size+  workgroup_size' <- mapM compileGroupDim workgroup_size   let kernel_size = zipWith mult_exp num_workgroups' workgroup_size'       total_elements = foldl mult_exp (Integer 1) kernel_size       cond = BinOp "!=" total_elements (Integer 0)

@@ -10,6 +10,7 @@     KernelOp (..),     Fence (..),     AtomicOp (..),+    GroupDim,     Kernel (..),     KernelUse (..),     module Futhark.CodeGen.ImpCode,@@ -32,7 +33,9 @@ type KernelCode = Code KernelOp  -- | A run-time constant related to kernels.-newtype KernelConst = SizeConst Name+data KernelConst+  = SizeConst Name+  | SizeMaxConst SizeClass   deriving (Eq, Ord, Show)  -- | An expression whose variables are kernel constants.@@ -46,13 +49,16 @@   | GetSizeMax VName SizeClass   deriving (Show) +-- | The size of one dimension of a group.+type GroupDim = Either Exp KernelConst+ -- | A generic kernel containing arbitrary kernel code. data Kernel = Kernel   { kernelBody :: Code KernelOp,     -- | The host variables referenced by the kernel.     kernelUses :: [KernelUse],     kernelNumGroups :: [Exp],-    kernelGroupSize :: [Exp],+    kernelGroupSize :: [GroupDim],     -- | A short descriptive and _unique_ name - should be     -- alphanumeric and without spaces.     kernelName :: Name,@@ -81,7 +87,12 @@  instance Pretty KernelConst where   pretty (SizeConst key) = "get_size" <> parens (pretty key)+  pretty (SizeMaxConst size_class) = "get_max_size" <> parens (pretty size_class) +instance FreeIn KernelConst where+  freeIn' (SizeConst _) = mempty+  freeIn' (SizeMaxConst _) = mempty+ instance Pretty KernelUse where   pretty (ScalarUse name t) =     oneLine $ "scalar_copy" <> parens (commasep [pretty name, pretty t])@@ -118,8 +129,11 @@  instance FreeIn Kernel where   freeIn' kernel =-    freeIn' (kernelBody kernel)-      <> freeIn' [kernelNumGroups kernel, kernelGroupSize kernel]+    freeIn'+      ( kernelBody kernel,+        kernelNumGroups kernel,+        kernelGroupSize kernel+      )  instance Pretty Kernel where   pretty kernel =@@ -128,7 +142,7 @@         ( "groups"             <+> brace (pretty $ kernelNumGroups kernel)             </> "group_size"-            <+> brace (pretty $ kernelGroupSize kernel)+            <+> brace (list $ map (either pretty pretty) $ kernelGroupSize kernel)             </> "uses"             <+> brace (commasep $ map pretty $ kernelUses kernel)             </> "failure_tolerant"

@@ -16,6 +16,8 @@     numFailureParams,     KernelTarget (..),     FailureMsg (..),+    GroupDim,+    KernelConst (..),     module Futhark.CodeGen.ImpCode,     module Futhark.IR.GPU.Sizes,   )@@ -24,6 +26,7 @@ import Data.Map qualified as M import Data.Text qualified as T import Futhark.CodeGen.ImpCode+import Futhark.CodeGen.ImpCode.GPU (GroupDim, KernelConst (..)) import Futhark.IR.GPU.Sizes import Futhark.Util.Pretty @@ -92,7 +95,7 @@  -- | Host-level OpenCL operation. data OpenCL-  = LaunchKernel KernelSafety KernelName [KernelArg] [Exp] [Exp]+  = LaunchKernel KernelSafety KernelName [KernelArg] [Exp] [GroupDim]   | GetSize VName Name   | CmpSizeLe VName Name Exp   | GetSizeMax VName SizeClass

@@ -874,6 +874,8 @@         constExp =<< hasExp =<< M.lookup name vtable       constExp (Op (Inner (SizeOp (GetSize key _)))) =         Just $ LeafExp (Imp.SizeConst $ keyWithEntryPoint fname key) int32+      constExp (Op (Inner (SizeOp (GetSizeMax size_class)))) =+        Just $ LeafExp (Imp.SizeMaxConst size_class) int32       constExp e = primExpFromExp lookupConstExp e   pure $ replaceInPrimExpM onLeaf size   where@@ -1141,16 +1143,27 @@   HostEnv atomics _ locks <- askEnv   body <- makeAllMemoryGlobal $ subImpM_ (KernelEnv atomics constants locks) ops m   uses <- computeKernelUses body mempty+  group_size <- onGroupSize $ kernelGroupSize constants   emit . Imp.Op . Imp.CallKernel $     Imp.Kernel       { Imp.kernelBody = body,         Imp.kernelUses = uses,         Imp.kernelNumGroups = [untyped $ kernelNumGroups constants],-        Imp.kernelGroupSize = [untyped $ kernelGroupSize constants],+        Imp.kernelGroupSize = [group_size],         Imp.kernelName = name,         Imp.kernelFailureTolerant = kAttrFailureTolerant attrs,         Imp.kernelCheckLocalMemory = kAttrCheckLocalMemory attrs       }+  where+    -- Figure out if this expression actually corresponds to a+    -- KernelConst.+    onGroupSize e = do+      vtable <- getVTable+      x <- isConstExp vtable $ untyped e+      pure $+        case x of+          Just (LeafExp kc _) -> Right kc+          _ -> Left $ untyped e  sKernelFailureTolerant ::   Bool ->

@@ -39,6 +39,7 @@  import Control.Monad.Except import Data.List (foldl', genericLength, zip5)+import Data.Map qualified as M import Data.Maybe import Futhark.CodeGen.ImpCode.GPU qualified as Imp import Futhark.CodeGen.ImpGen@@ -542,10 +543,7 @@           AtomicCAS f -> pure $ const $ pure f           AtomicLocking f -> pure $ \hist_H_chk -> do             let lock_shape =-                  Shape $-                    tvSize num_subhistos_per_group-                      : shapeDims (histOpShape op)-                      ++ [hist_H_chk]+                  Shape [tvSize num_subhistos_per_group, hist_H_chk]              let dims = map pe64 $ shapeDims lock_shape @@ -902,6 +900,19 @@    max_group_size <- dPrim "max_group_size" int32   sOp $ Imp.GetSizeMax (tvVar max_group_size) Imp.SizeGroup++  -- XXX: we need to record for later use that max_group_size is the+  -- result of GetSizeMax.  This is an ugly hack that reflects our+  -- inability to track which variables are actually constants.+  let withSizeMax vtable =+        case M.lookup (tvVar max_group_size) vtable of+          Just (ScalarVar _ se) ->+            M.insert+              (tvVar max_group_size)+              (ScalarVar (Just (Op (Inner (SizeOp (GetSizeMax SizeGroup))))) se)+              vtable+          _ -> vtable+   let group_size = Imp.Count $ Var $ tvVar max_group_size   num_groups <-     fmap (Imp.Count . tvSize) $@@ -1030,16 +1041,17 @@               Just $                 untyped $                   unCount groups_per_segment-        histKernelLocal-          hist_M-          groups_per_segment-          map_pes-          num_groups-          group_size-          space-          hist_S-          slugs-          kbody+        localVTable withSizeMax $+          histKernelLocal+            hist_M+            groups_per_segment+            map_pes+            num_groups+            group_size+            space+            hist_S+            slugs+            kbody    pure (pick_local, run) 

@@ -11,6 +11,7 @@ import Control.Monad.Identity import Control.Monad.Reader import Control.Monad.State+import Data.Bifunctor (second) import Data.Map.Strict qualified as M import Data.Maybe import Data.Set qualified as S@@ -64,13 +65,13 @@        (device_prototypes, device_defs) = unzip $ M.elems device_funs       kernels' = M.map fst kernels-      opencl_code = openClCode $ map snd $ M.elems kernels+      opencl_code = T.unlines $ map snd $ M.elems kernels        opencl_prelude =         T.unlines           [ genPrelude target used_types,             definitionsText device_prototypes,-            funcsText device_defs+            T.unlines device_defs           ]    in ImpOpenCL.Program         opencl_code@@ -125,8 +126,8 @@ errorLabel = ("error_" ++) . show . kernelNextSync  data ToOpenCL = ToOpenCL-  { clGPU :: M.Map KernelName (KernelSafety, C.Func),-    clDevFuns :: M.Map Name (C.Definition, C.Func),+  { clGPU :: M.Map KernelName (KernelSafety, T.Text),+    clDevFuns :: M.Map Name (C.Definition, T.Text),     clUsedTypes :: S.Set PrimType,     clSizes :: M.Map Name SizeClass,     clFailures :: [FailureMsg]@@ -237,7 +238,7 @@   modify $ \s ->     s       { clUsedTypes = typesInCode (functionBody device_func) <> clUsedTypes s,-        clDevFuns = M.insert fname func $ clDevFuns s,+        clDevFuns = M.insert fname (second funcText func) $ clDevFuns s,         clFailures = kernelFailures kstate       } @@ -284,6 +285,15 @@     toDevice :: HostOp -> KernelOp     toDevice _ = bad +isConst :: GroupDim -> Maybe T.Text+isConst (Left (ValueExp (IntValue x))) =+  Just $ prettyText $ intToInt64 x+isConst (Right (SizeConst v)) =+  Just $ T.pack $ zEncodeString $ nameToString v+isConst (Right (SizeMaxConst size_class)) =+  Just $ T.pack $ "max_" <> prettyString size_class+isConst _ = Nothing+ onKernel :: KernelTarget -> Kernel -> OnKernelM OpenCL onKernel target kernel = do   called <- ensureDeviceFuns $ kernelBody kernel@@ -389,18 +399,30 @@           ++ catMaybes local_memory_params           ++ use_params +      attribute =+        case (target, mapM isConst $ kernelGroupSize kernel) of+          (TargetOpenCL, Just [x, y, z]) ->+            "__attribute__((reqd_work_group_size" <> prettyText (x, y, z) <> "))\n"+          (TargetOpenCL, Just [x, y]) ->+            "__attribute__((reqd_work_group_size" <> prettyText (x, y, 1 :: Int) <> "))\n"+          (TargetOpenCL, Just [x]) ->+            "__attribute__((reqd_work_group_size" <> prettyText (x, 1 :: Int, 1 :: Int) <> "))\n"+          _ -> ""+       kernel_fun =-        [C.cfun|__kernel void $id:name ($params:params) {-                  $items:(mconcat unpack_params)-                  $items:const_defs-                  $items:block_dim_init-                  $items:local_memory_init-                  $items:error_init-                  $items:kernel_body+        attribute+          <> funcText+            [C.cfun|__kernel void $id:name ($params:params) {+                    $items:(mconcat unpack_params)+                    $items:const_defs+                    $items:block_dim_init+                    $items:local_memory_init+                    $items:error_init+                    $items:kernel_body -                  $id:(errorLabel kstate): return;+                    $id:(errorLabel kstate): return; -                  $items:const_undefs+                    $items:const_undefs                 }|]   modify $ \s ->     s@@ -411,9 +433,7 @@    -- The argument corresponding to the global_failure parameters is   -- added automatically later.-  let args =-        catMaybes local_memory_args-          ++ kernelArgs kernel+  let args = catMaybes local_memory_args ++ kernelArgs kernel    pure $ LaunchKernel safety name args num_groups group_size   where@@ -472,15 +492,6 @@     undef = "#undef " <> idText (C.toIdent v mempty) constDef _ = Nothing -openClCode :: [C.Func] -> T.Text-openClCode kernels =-  definitionsText [C.cunit|$edecls:funcs|]-  where-    funcs =-      [ [C.cedecl|$func:kernel_func|]-        | kernel_func <- kernels-      ]- genOpenClPrelude :: S.Set PrimType -> T.Text genOpenClPrelude ts =   [untrimming|@@ -637,6 +648,8 @@   where     compileKernelConst (SizeConst key) =       pure [C.cexp|$id:(zEncodeString (prettyString key))|]+    compileKernelConst (SizeMaxConst size_class) =+      pure [C.cexp|$id:("max_" <> prettyString size_class)|]  kernelArgs :: Kernel -> [KernelArg] kernelArgs = mapMaybe useToArg . kernelUses

@@ -327,7 +327,7 @@           <> mapTranspose block_dim args t kind,       kernelUses = uses,       kernelNumGroups = map untyped num_groups,-      kernelGroupSize = map untyped group_size,+      kernelGroupSize = map (Left . untyped) group_size,       kernelName = nameFromString name,       kernelFailureTolerant = True,       kernelCheckLocalMemory = False

@@ -36,12 +36,15 @@ lamBody = lambdaBody . segBinOpLambda  -- Arrays for storing worker results.-resultArrays :: String -> [SegBinOp MCMem] -> MulticoreGen [[VName]]-resultArrays s segops =+carryArrays :: String -> TV Int32 -> [SegBinOp MCMem] -> MulticoreGen [[VName]]+carryArrays s nsubtasks segops =   forM segops $ \(SegBinOp _ lam _ shape) ->     forM (lambdaReturnType lam) $ \t -> do       let pt = elemType t-          full_shape = shape <> arrayShape t+          full_shape =+            Shape [Var (tvVar nsubtasks)]+              <> shape+              <> arrayShape t       sAllocArray s pt full_shape DefaultSpace  nonsegmentedScan ::@@ -69,14 +72,17 @@           | vectorize && no_array_param = (scanStage1Nested, scanStage3Nested)           | otherwise = (scanStage1Fallback, scanStage3Fallback) +    emit $ Imp.DebugPrint "Scan stage 1" Nothing     scanStage1 pat space kbody scan_ops      let nsubtasks' = tvExp nsubtasks     sWhen (nsubtasks' .>. 1) $ do       scan_ops2 <- renameSegBinOp scan_ops-      scanStage2 pat nsubtasks space scan_ops2 kbody+      emit $ Imp.DebugPrint "Scan stage 2" Nothing+      carries <- scanStage2 pat nsubtasks space scan_ops2       scan_ops3 <- renameSegBinOp scan_ops-      scanStage3 pat space kbody scan_ops3+      emit $ Imp.DebugPrint "Scan stage 3" Nothing+      scanStage3 pat space scan_ops3 carries  -- Different ways to generate code for a scan loop data ScanLoopType@@ -98,7 +104,10 @@ getExtract _ = extractVectorLane  genBinOpParams :: [SegBinOp MCMem] -> MulticoreGen ()-genBinOpParams scan_ops = dScope Nothing $ scopeOfLParams $ concatMap (lambdaParams . segBinOpLambda) scan_ops+genBinOpParams scan_ops =+  dScope Nothing $+    scopeOfLParams $+      concatMap (lambdaParams . segBinOpLambda) scan_ops  genLocalAccsStage1 :: [SegBinOp MCMem] -> MulticoreGen [[VName]] genLocalAccsStage1 scan_ops = do@@ -127,7 +136,41 @@ getNestLoop ScanNested = sLoopNestVectorized getNestLoop _ = sLoopNest --- Generate a loop which performs a potentially vectorized scan.+applyScanOps ::+  ScanLoopType ->+  Pat LetDecMem ->+  SegSpace ->+  [SubExp] ->+  [SegBinOp MCMem] ->+  [[VName]] ->+  ImpM MCMem HostEnv Imp.Multicore ()+applyScanOps typ pat space all_scan_res scan_ops local_accs = do+  let per_scan_res = segBinOpChunks scan_ops all_scan_res+      per_scan_pes = segBinOpChunks scan_ops $ patElems pat+  let (is, _) = unzip $ unSegSpace space++  -- Potential vector load and then do sequential scan+  getScanLoop typ $ \j ->+    forM_ (zip4 per_scan_pes scan_ops per_scan_res local_accs) $ \(pes, scan_op, scan_res, acc) ->+      getNestLoop typ (segBinOpShape scan_op) $ \vec_is -> do+        sComment "Read accumulator" $+          forM_ (zip (xParams scan_op) acc) $ \(p, acc') -> do+            copyDWIMFix (paramName p) [] (Var acc') vec_is+        sComment "Read next values" $+          forM_ (zip (yParams scan_op) scan_res) $ \(p, se) ->+            getExtract typ j $+              collect $+                copyDWIMFix (paramName p) [] se vec_is+        -- Scan body+        sComment "Scan op body" $+          compileStms mempty (bodyStms $ lamBody scan_op) $+            forM_ (zip3 acc pes $ map resSubExp $ bodyResult $ lamBody scan_op) $+              \(acc', pe, se) -> do+                copyDWIMFix (patElemName pe) (map Imp.le64 is ++ vec_is) se []+                copyDWIMFix acc' vec_is se []++-- Generate a loop which performs a potentially vectorized scan on the+-- result of a kernel body. genScanLoop ::   ScanLoopType ->   Pat LetDecMem ->@@ -138,37 +181,18 @@   Imp.TExp Int64 ->   ImpM MCMem HostEnv Imp.Multicore () genScanLoop typ pat space kbody scan_ops local_accs i = do-  let (all_scan_res, map_res) = splitAt (segBinOpResults scan_ops) $ kernelBodyResult kbody-      per_scan_res = segBinOpChunks scan_ops all_scan_res-      per_scan_pes = segBinOpChunks scan_ops $ patElems pat+  let (all_scan_res, map_res) =+        splitAt (segBinOpResults scan_ops) $ kernelBodyResult kbody   let (is, ns) = unzip $ unSegSpace space       ns' = map pe64 ns    zipWithM_ dPrimV_ is $ unflattenIndex ns' i   compileStms mempty (kernelBodyStms kbody) $ do-    -- Potential vector load and then do sequential scan-    getScanLoop typ $ \j -> do-      sComment "write mapped values results to memory" $ do-        let map_arrs = drop (segBinOpResults scan_ops) $ patElems pat-        zipWithM_ (compileThreadResult space) map_arrs map_res-      forM_ (zip4 per_scan_pes scan_ops per_scan_res local_accs) $ \(pes, scan_op, scan_res, acc) ->-        getNestLoop typ (segBinOpShape scan_op) $ \vec_is -> do-          -- Read accum value-          forM_ (zip (xParams scan_op) acc) $ \(p, acc') -> do-            copyDWIMFix (paramName p) [] (Var acc') vec_is-          -- Read next value-          sComment "Read next values" $-            forM_ (zip (yParams scan_op) scan_res) $ \(p, se) ->-              getExtract typ j $-                collect $-                  copyDWIMFix (paramName p) [] (kernelResultSubExp se) vec_is-          -- Scan body-          sComment "Scan body" $-            compileStms mempty (bodyStms $ lamBody scan_op) $-              forM_ (zip3 acc pes $ map resSubExp $ bodyResult $ lamBody scan_op) $-                \(acc', pe, se) -> do-                  copyDWIMFix (patElemName pe) (map Imp.le64 is ++ vec_is) se []-                  copyDWIMFix acc' vec_is se []+    let map_arrs = drop (segBinOpResults scan_ops) $ patElems pat+    sComment "write mapped values results to memory" $+      zipWithM_ (compileThreadResult space) map_arrs map_res+    sComment "Apply scan op" $+      applyScanOps typ pat space (map kernelResultSubExp all_scan_res) scan_ops local_accs  scanStage1Scalar ::   Pat LetDecMem ->@@ -200,11 +224,10 @@     dPrim_ (segFlat space) int64     sOp $ Imp.GetTaskId (segFlat space) -    lparams <- collect $ genBinOpParams scan_ops     local_accs <- genLocalAccsStage1 scan_ops      inISPC $ do-      emit lparams+      genBinOpParams scan_ops       generateChunkLoop "SegScan" Scalar $ \i -> do         genScanLoop ScanNested pat space kbody scan_ops local_accs i @@ -237,14 +260,12 @@   TV Int32 ->   SegSpace ->   [SegBinOp MCMem] ->-  KernelBody MCMem ->-  MulticoreGen ()-scanStage2 pat nsubtasks space scan_ops kbody = do-  emit $ Imp.DebugPrint "nonsegmentedScan stage 2" Nothing+  MulticoreGen [[VName]]+scanStage2 pat nsubtasks space scan_ops = do   let (is, ns) = unzip $ unSegSpace space       ns_64 = map pe64 ns       per_scan_pes = segBinOpChunks scan_ops $ patElems pat-      nsubtasks' = tvExp nsubtasks+      nsubtasks' = sExt64 $ tvExp nsubtasks    dScope Nothing $ scopeOfLParams $ concatMap (lambdaParams . segBinOpLambda) scan_ops   offset <- dPrimV "offset" (0 :: Imp.TExp Int64)@@ -255,103 +276,106 @@   -- Parameters used to find the chunk sizes   -- Perhaps get this information from ``scheduling information``   -- instead of computing it manually here.-  let iter_pr_subtask = product ns_64 `quot` sExt64 nsubtasks'-      remainder = product ns_64 `rem` sExt64 nsubtasks'+  let iter_pr_subtask = product ns_64 `quot` nsubtasks'+      remainder = product ns_64 `rem` nsubtasks' -  accs <- resultArrays "scan_stage_2_accum" scan_ops-  forM_ (zip scan_ops accs) $ \(scan_op, acc) ->-    sLoopNest (segBinOpShape scan_op) $ \vec_is ->-      forM_ (zip acc $ segBinOpNeutral scan_op) $ \(acc', ne) ->-        copyDWIMFix acc' vec_is ne []+  carries <- carryArrays "scan_stage_2_carry" nsubtasks scan_ops+  sComment "carry-in for first chunk is neutral" $+    forM_ (zip scan_ops carries) $ \(scan_op, carry) ->+      sLoopNest (segBinOpShape scan_op) $ \vec_is ->+        forM_ (zip carry $ segBinOpNeutral scan_op) $ \(carry', ne) ->+          copyDWIMFix carry' (0 : vec_is) ne []    -- Perform sequential scan over the last element of each chunk-  sFor "i" (nsubtasks' - 1) $ \i -> do+  sComment "scan carries" $ sFor "i" (nsubtasks' - 1) $ \i -> do     offset <-- iter_pr_subtask     sWhen (sExt64 i .<. remainder) (offset <-- offset' + 1)     offset_index <-- offset_index' + offset'     zipWithM_ dPrimV_ is $ unflattenIndex ns_64 $ sExt64 offset_index' -    compileStms mempty (kernelBodyStms kbody) $-      forM_ (zip3 per_scan_pes scan_ops accs) $ \(pes, scan_op, acc) ->-        sLoopNest (segBinOpShape scan_op) $ \vec_is -> do-          sComment "Read carry in" $-            forM_ (zip (xParams scan_op) acc) $ \(p, acc') ->-              copyDWIMFix (paramName p) [] (Var acc') vec_is--          sComment "Read next values" $-            forM_ (zip (yParams scan_op) pes) $ \(p, pe) ->-              copyDWIMFix (paramName p) [] (Var $ patElemName pe) ((offset_index' - 1) : vec_is)+    forM_ (zip3 per_scan_pes scan_ops carries) $ \(pes, scan_op, carry) ->+      sLoopNest (segBinOpShape scan_op) $ \vec_is -> do+        sComment "Read carry" $+          forM_ (zip (xParams scan_op) carry) $ \(p, carry') ->+            copyDWIMFix (paramName p) [] (Var carry') (i : vec_is) -          compileStms mempty (bodyStms $ lamBody scan_op) $-            forM_ (zip3 acc pes $ map resSubExp $ bodyResult $ lamBody scan_op) $-              \(acc', pe, se) -> do-                copyDWIMFix (patElemName pe) ((offset_index' - 1) : vec_is) se []-                copyDWIMFix acc' vec_is se []+        sComment "Read next values" $+          forM_ (zip (yParams scan_op) pes) $ \(p, pe) ->+            copyDWIMFix (paramName p) [] (Var $ patElemName pe) ((offset_index' - 1) : vec_is) -genLocalAccsStage3 :: [SegBinOp MCMem] -> [[PatElem LetDecMem]] -> MulticoreGen [[VName]]-genLocalAccsStage3 scan_ops per_scan_pes =-  forM (zip scan_ops per_scan_pes) $ \(scan_op, pes) -> do-    let shape = segBinOpShape scan_op-        ts = lambdaReturnType $ segBinOpLambda scan_op-    forM (zip4 (xParams scan_op) pes ts $ segBinOpNeutral scan_op) $ \(p, pe, t, ne) -> do-      acc <--        case shapeDims shape of-          [] -> pure $ paramName p-          _ -> do-            let pt = elemType t-            sAllocArray "local_acc" pt (shape <> arrayShape t) DefaultSpace+        compileStms mempty (bodyStms $ lamBody scan_op) $+          forM_ (zip carry $ map resSubExp $ bodyResult $ lamBody scan_op) $ \(carry', se) -> do+            copyDWIMFix carry' ((i + 1) : vec_is) se [] -      -- Initialise the accumulator with neutral from previous chunk.-      -- or read neutral if first ``iter``-      (start, _end) <- getLoopBounds-      sLoopNest (segBinOpShape scan_op) $ \vec_is -> do-        let read_carry_in =-              copyDWIMFix acc vec_is (Var $ patElemName pe) (start - 1 : vec_is)-            read_neutral =-              copyDWIMFix acc vec_is ne []-        sIf (start .==. 0) read_neutral read_carry_in-      pure acc+  -- Return the array of carries for each chunk.+  pure carries  scanStage3Scalar ::   Pat LetDecMem ->   SegSpace ->-  KernelBody MCMem ->   [SegBinOp MCMem] ->+  [[VName]] ->   MulticoreGen ()-scanStage3Scalar pat space kbody scan_ops = do+scanStage3Scalar pat space scan_ops per_scan_carries = do   let per_scan_pes = segBinOpChunks scan_ops $ patElems pat+      (is, ns) = unzip $ unSegSpace space+      ns' = map pe64 ns+   body <- collect $ do     dPrim_ (segFlat space) int64-    sOp $ Imp.GetTaskId (segFlat space)+    sOp $ Imp.GetTaskId $ segFlat space -    genBinOpParams scan_ops-    local_accs <- genLocalAccsStage3 scan_ops per_scan_pes+    inISPC $ do+      genBinOpParams scan_ops+      sComment "load carry-in" $+        forM_ (zip per_scan_carries scan_ops) $ \(op_carries, scan_op) ->+          forM_ (zip (xParams scan_op) op_carries) $ \(p, carries) ->+            copyDWIMFix (paramName p) [] (Var carries) [le64 (segFlat space)]+      generateChunkLoop "SegScan" Vectorized $ \i -> do+        zipWithM_ dPrimV_ is $ unflattenIndex ns' i+        sComment "load partial result" $+          forM_ (zip per_scan_pes scan_ops) $ \(scan_pes, scan_op) ->+            forM_ (zip (yParams scan_op) scan_pes) $ \(p, pe) ->+              copyDWIMFix (paramName p) [] (Var (patElemName pe)) (map le64 is)+        sComment "combine carry with partial result" $+          forM_ (zip per_scan_pes scan_ops) $ \(scan_pes, scan_op) ->+            compileStms mempty (bodyStms $ lamBody scan_op) $+              forM_ (zip scan_pes $ map resSubExp $ bodyResult $ lamBody scan_op) $ \(pe, se) ->+                copyDWIMFix (patElemName pe) (map Imp.le64 is) se [] -    inISPC $-      generateChunkLoop "SegScan" Vectorized $-        genScanLoop ScanScalar pat space kbody scan_ops local_accs   free_params <- freeParams body   emit $ Imp.Op $ Imp.ParLoop "scan_stage_3" body free_params  scanStage3Nested ::   Pat LetDecMem ->   SegSpace ->-  KernelBody MCMem ->   [SegBinOp MCMem] ->+  [[VName]] ->   MulticoreGen ()-scanStage3Nested pat space kbody scan_ops = do+scanStage3Nested pat space scan_ops per_scan_carries = do   let per_scan_pes = segBinOpChunks scan_ops $ patElems pat+      (is, ns) = unzip $ unSegSpace space+      ns' = map pe64 ns   body <- collect $ do     dPrim_ (segFlat space) int64     sOp $ Imp.GetTaskId (segFlat space) -    lparams <- collect $ genBinOpParams scan_ops-    local_accs <- genLocalAccsStage3 scan_ops per_scan_pes+    generateChunkLoop "SegScan" Scalar $ \i -> do+      genBinOpParams scan_ops+      zipWithM_ dPrimV_ is $ unflattenIndex ns' i+      forM_ (zip3 per_scan_pes per_scan_carries scan_ops) $ \(scan_pes, op_carries, scan_op) -> do+        sLoopNest (segBinOpShape scan_op) $ \vec_is -> do+          sComment "load carry-in" $+            forM_ (zip (xParams scan_op) op_carries) $ \(p, carries) ->+              copyDWIMFix (paramName p) [] (Var carries) (le64 (segFlat space) : vec_is) -    inISPC $ do-      emit lparams-      generateChunkLoop "SegScan" Scalar $ \i -> do-        genScanLoop ScanNested pat space kbody scan_ops local_accs i+          sComment "load partial result" $+            forM_ (zip (yParams scan_op) scan_pes) $ \(p, pe) ->+              copyDWIMFix (paramName p) [] (Var (patElemName pe)) (map le64 is ++ vec_is)+          sComment "combine carry with partial result" $+            compileStms mempty (bodyStms $ lamBody scan_op) $+              forM_ (zip scan_pes $ map resSubExp $ bodyResult $ lamBody scan_op) $ \(pe, se) ->+                copyDWIMFix (patElemName pe) (map Imp.le64 is ++ vec_is) se []    free_params <- freeParams body   emit $ Imp.Op $ Imp.ParLoop "scan_stage_3" body free_params@@ -359,20 +383,34 @@ scanStage3Fallback ::   Pat LetDecMem ->   SegSpace ->-  KernelBody MCMem ->   [SegBinOp MCMem] ->+  [[VName]] ->   MulticoreGen ()-scanStage3Fallback pat space kbody scan_ops = do+scanStage3Fallback pat space scan_ops per_scan_carries = do   let per_scan_pes = segBinOpChunks scan_ops $ patElems pat+      (is, ns) = unzip $ unSegSpace space+      ns' = map pe64 ns   body <- collect $ do     dPrim_ (segFlat space) int64     sOp $ Imp.GetTaskId (segFlat space)      genBinOpParams scan_ops-    local_accs <- genLocalAccsStage3 scan_ops per_scan_pes -    generateChunkLoop "SegScan" Scalar $-      genScanLoop ScanSeq pat space kbody scan_ops local_accs+    generateChunkLoop "SegScan" Scalar $ \i -> do+      zipWithM_ dPrimV_ is $ unflattenIndex ns' i+      forM_ (zip3 per_scan_pes per_scan_carries scan_ops) $ \(scan_pes, op_carries, scan_op) -> do+        sLoopNest (segBinOpShape scan_op) $ \vec_is -> do+          sComment "load carry-in" $+            forM_ (zip (xParams scan_op) op_carries) $ \(p, carries) ->+              copyDWIMFix (paramName p) [] (Var carries) (le64 (segFlat space) : vec_is)++          sComment "load partial result" $+            forM_ (zip (yParams scan_op) scan_pes) $ \(p, pe) ->+              copyDWIMFix (paramName p) [] (Var (patElemName pe)) (map le64 is ++ vec_is)+          sComment "combine carry with partial result" $+            compileStms mempty (bodyStms $ lamBody scan_op) $+              forM_ (zip scan_pes $ map resSubExp $ bodyResult $ lamBody scan_op) $ \(pe, se) ->+                copyDWIMFix (patElemName pe) (map Imp.le64 is ++ vec_is) se []   free_params <- freeParams body   emit $ Imp.Op $ Imp.ParLoop "scan_stage_3" body free_params 

@@ -141,10 +141,13 @@ -- contiguous, i.e., if we instantiate all the points of the current index -- function, do we get a contiguous memory interval? ----- By definition, the LMAD denotes the set of points (simplified):+-- By definition, the LMAD \( \sigma + \{ (n_1, s_1), \ldots, (n_k, s_k) \} \),+-- where \(n\) and \(s\) denote the shape and stride of each dimension, denotes+-- the set of points: -----   \{ o + \Sigma_{j=0}^{k} ((i_j+r_j) `mod` n_j)*s_j,---      \forall i_j such that 0<=i_j<n_j, j=1..k \}+-- \[+--    \{ ~ \sigma + i_1 * s_1 + \ldots + i_m * s_m ~ | ~ 0 \leq i_1 < n_1, \ldots, 0 \leq i_m < n_m ~ \}+-- \] data LMAD num = LMAD   { lmadOffset :: num,     lmadDims :: [LMADDim num]

@@ -358,7 +358,7 @@   pretty (Lambda [] (Body _ stms []) []) | stms == mempty = "nilFn"   pretty (Lambda params body rettype) =     "\\"-      <+> ppTuple' (map pretty params)+      <+> braces (commastack $ map pretty params)       </> indent 2 (colon <+> ppTupleLines' (map pretty rettype) <+> "->")       </> indent 2 (pretty body) 

@@ -213,6 +213,9 @@ instance (FreeIn a, FreeIn b, FreeIn c, FreeIn d) => FreeIn (a, b, c, d) where   freeIn' (a, b, c, d) = freeIn' a <> freeIn' b <> freeIn' c <> freeIn' d +instance (FreeIn a, FreeIn b) => FreeIn (Either a b) where+  freeIn' = either freeIn' freeIn'+ instance FreeIn a => FreeIn [a] where   freeIn' = foldMap freeIn' 

@@ -1,7 +1,6 @@ -- | Defunctionalization of typed, monomorphic Futhark programs without modules. module Futhark.Internalise.Defunctionalise (transformProg) where -import Control.Arrow qualified as Arrow import Control.Monad.Identity import Control.Monad.Reader import Control.Monad.State@@ -15,21 +14,16 @@ import Data.Set qualified as S import Futhark.IR.Pretty () import Futhark.MonadFreshNames+import Futhark.Util (mapAccumLM) import Language.Futhark import Language.Futhark.Traversals --- | An expression or an extended 'Lambda' (with size parameters,--- which AST lambdas do not support).-data ExtExp-  = ExtLambda [Pat] Exp StructRetType SrcLoc-  | ExtExp Exp-  deriving (Show)- -- | A static value stores additional information about the result of -- defunctionalization of an expression, aside from the residual expression. data StaticVal   = Dynamic PatType-  | LambdaSV Pat StructRetType ExtExp Env+  | -- | The Env is the lexical closure of the lambda.+    LambdaSV Pat StructRetType Exp Env   | RecordSV [(Name, StaticVal)]   | -- | The constructor that is actually present, plus     -- the others that are not.@@ -56,7 +50,7 @@ type Env = M.Map VName Binding  localEnv :: Env -> DefM a -> DefM a-localEnv env = local $ Arrow.second (env <>)+localEnv env = local $ second (env <>)  -- Even when using a "new" environment (for evaluating closures) we -- still ram the global environment of DynamicFuns in there.@@ -68,7 +62,7 @@ askEnv = asks snd  areGlobal :: [VName] -> DefM a -> DefM a-areGlobal vs = local $ Arrow.first (S.fromList vs <>)+areGlobal vs = local $ first (S.fromList vs <>)  replaceTypeSizes ::   M.Map VName SizeSubst ->@@ -98,7 +92,7 @@        in LambdaSV             (onAST substs param)             (RetType t_dims (replaceTypeSizes substs t))-            (onExtExp substs e)+            (onExp substs e)             (onEnv orig_substs closure_env) -- intentional     Dynamic t ->       Dynamic $ replaceTypeSizes orig_substs t@@ -140,6 +134,13 @@       AppExp (Coerce (onExp substs e) te' loc) (Info (AppRes (replaceTypeSizes substs t) ext))       where         te' = onTypeExp substs te+    onExp substs (Lambda params e ret (Info (als, RetType t_dims t)) loc) =+      Lambda+        (map (onAST substs) params)+        (onExp substs e)+        ret+        (Info (als, RetType t_dims (replaceTypeSizes substs t)))+        loc     onExp substs e = onAST substs e      onTypeExpDim substs d@(SizeExpNamed v loc) =@@ -176,15 +177,6 @@     onTypeExp _ (TEVar v loc) =       TEVar v loc -    onExtExp substs (ExtExp e) =-      ExtExp $ onExp substs e-    onExtExp substs (ExtLambda params e (RetType t_dims t) loc) =-      ExtLambda-        (map (onAST substs) params)-        (onExp substs e)-        (RetType t_dims (replaceTypeSizes substs t))-        loc-     onEnv substs =       M.fromList         . map (second (onBinding substs))@@ -389,11 +381,11 @@   -- remaining ones (if there are any) into the body of the lambda.   let (pat, ret', e0') = case pats of         [] -> error "Received a lambda with no parameters."-        [pat'] -> (pat', ret, ExtExp e0)+        [pat'] -> (pat', ret, e0)         (pat' : pats') ->           ( pat',             RetType [] $ foldFunType (map (toStruct . patternType) pats') ret,-            ExtLambda pats' e0 ret loc+            Lambda pats' e0 Nothing (Info (mempty, ret)) loc           )    -- Construct a record literal that closes over the environment of@@ -686,11 +678,6 @@ defuncExp' :: Exp -> DefM Exp defuncExp' = fmap fst . defuncExp -defuncExtExp :: ExtExp -> DefM (Exp, StaticVal)-defuncExtExp (ExtExp e) = defuncExp e-defuncExtExp (ExtLambda pats e0 ret loc) =-  defuncFun [] pats e0 ret loc- defuncCase :: StaticVal -> Case -> DefM (Case, StaticVal) defuncCase sv (CasePat p e loc) = do   let p' = updatePat p sv@@ -722,15 +709,8 @@ etaExpand :: PatType -> Exp -> DefM ([Pat], Exp, StructRetType) etaExpand e_t e = do   let (ps, ret) = getType $ RetType [] e_t-  (pats, vars) <- fmap unzip . forM ps $ \(p, t) -> do-    let t' = fromStruct t-    x <- case p of-      Named x -> pure x-      Unnamed -> newNameFromString "x"-    pure-      ( Id x (Info t') mempty,-        Var (qualName x) (Info t') mempty-      )+  -- Some careful hackery to avoid duplicate names.+  (_, (pats, vars)) <- second unzip <$> mapAccumLM f [] ps   let e' =         foldl'           ( \e1 (e2, t2, argtypes) ->@@ -746,6 +726,18 @@       let (ps, r) = getType t2 in ((p, t1) : ps, r)     getType t = ([], t) +    f prev (p, t) = do+      let t' = fromStruct t+      x <- case p of+        Named x | x `notElem` prev -> pure x+        _ -> newNameFromString "x"+      pure+        ( x : prev,+          ( Id x (Info t') mempty,+            Var (qualName x) (Info t') mempty+          )+        )+ -- | Defunctionalize an indexing of a single array dimension. defuncDimIndex :: DimIndexBase Info VName -> DefM (DimIndexBase Info VName) defuncDimIndex (DimFix e1) = DimFix . fst <$> defuncExp e1@@ -835,7 +827,7 @@           dims = mempty       (e0', sv) <-         localNewEnv (env' <> closure_env) $-          defuncExtExp e0+          defuncExp e0        let closure_pat = buildEnvPat dims closure_env           pat' = updatePat pat sv2

@@ -492,10 +492,10 @@          (loop_initial_cond, init_loop_cond_stms) <- collectStms $ do           forM_ (zip shapepat shapeinit) $ \(p, se) ->-            letBindNames [paramName p] $ BasicOp $ SubExp se+            letBindNames [I.paramName p] $ BasicOp $ SubExp se           forM_ (zip mergepat' mergeinit) $ \(p, se) ->-            unless (se == I.Var (paramName p)) $-              letBindNames [paramName p] $+            unless (se == I.Var (I.paramName p)) $+              letBindNames [I.paramName p] $                 BasicOp $                   case se of                     I.Var v@@ -515,13 +515,13 @@           -- Careful not to clobber anything.           loop_end_cond_body <- renameBody <=< buildBody_ $ do             forM_ (zip shapepat shapeargs) $ \(p, se) ->-              unless (se == I.Var (paramName p)) $-                letBindNames [paramName p] $+              unless (se == I.Var (I.paramName p)) $+                letBindNames [I.paramName p] $                   BasicOp $                     SubExp se             forM_ (zip mergepat' ses) $ \(p, se) ->-              unless (se == I.Var (paramName p)) $-                letBindNames [paramName p] $+              unless (se == I.Var (I.paramName p)) $+                letBindNames [I.paramName p] $                   BasicOp $                     case se of                       I.Var v@@ -1167,7 +1167,7 @@   img_params <- mapM (newParam "img_p" . rowType) =<< mapM lookupType img'   let params = bucket_params ++ img_params       rettype = replicate dim (I.Prim int64) ++ ne_ts-      body = mkBody mempty $ varsRes $ map paramName params+      body = mkBody mempty $ varsRes $ map I.paramName params   lam' <-     mkLambda params $       ensureResultShape@@ -1205,7 +1205,7 @@     let w = arraysSize 0 bs_ts     fmap subExpsRes . letValExp' "acc_res" $       I.Op $-        I.Screma w (paramName acc_p : bs') (I.mapSOAC lam')+        I.Screma w (I.paramName acc_p : bs') (I.mapSOAC lam')    op' <-     case op of

@@ -3,7 +3,12 @@ {-# LANGUAGE TypeFamilies #-}  -- | Perform array short circuiting-module Futhark.Optimise.ArrayShortCircuiting (optimiseSeqMem, optimiseGPUMem) where+module Futhark.Optimise.ArrayShortCircuiting+  ( optimiseSeqMem,+    optimiseGPUMem,+    optimiseMCMem,+  )+where  import Control.Monad.Reader import Data.Function ((&))@@ -12,6 +17,7 @@ import Futhark.Analysis.Alias qualified as AnlAls import Futhark.IR.Aliases import Futhark.IR.GPUMem+import Futhark.IR.MCMem import Futhark.IR.Mem.IxFun (substituteInIxFun) import Futhark.IR.SeqMem import Futhark.Optimise.ArrayShortCircuiting.ArrayCoalescing@@ -20,16 +26,13 @@ import Futhark.Pass qualified as Pass import Futhark.Util --------------------------------------------------------------------- Printer/Tester Main Program------------------------------------------------------------------ data Env inner = Env-  { envCoalesceTab :: M.Map VName Coalesced,-    onInner :: inner -> ReplaceM inner inner+  { envCoalesceTab :: CoalsTab,+    onInner :: inner -> UpdateM inner inner,+    memAllocsToRemove :: Names   } -type ReplaceM inner a = Reader (Env inner) a+type UpdateM inner a = Reader (Env inner) a  optimiseSeqMem :: Pass SeqMem SeqMem optimiseSeqMem = pass "short-circuit" "Array Short-Circuiting" mkCoalsTab pure replaceInParams@@ -37,6 +40,9 @@ optimiseGPUMem :: Pass GPUMem GPUMem optimiseGPUMem = pass "short-circuit-gpu" "Array Short-Circuiting (GPU)" mkCoalsTabGPU replaceInHostOp replaceInParams +optimiseMCMem :: Pass MCMem MCMem+optimiseMCMem = pass "short-circuit-mc" "Array Short-Circuiting (MC)" mkCoalsTabMC replaceInMCOp replaceInParams+ replaceInParams :: CoalsTab -> [Param FParamMem] -> (Names, [Param FParamMem]) replaceInParams coalstab fparams =   let (mem_allocs_to_remove, fparams') =@@ -56,47 +62,68 @@               (to_remove, Param attrs name (MemArray pt shp u $ ArrayIn (dstmem entry) ixf) : acc)         _ -> (to_remove, Param attrs name dec : acc) -removeStms :: Names -> Body rep -> Body rep-removeStms to_remove (Body dec stms res) =-  Body dec (stmsFromList $ filter (not . flip nameIn to_remove . head . patNames . stmPat) $ stmsToList stms) res+removeAllocsInStms :: Stms rep -> UpdateM inner (Stms rep)+removeAllocsInStms stms = do+  to_remove <- asks memAllocsToRemove+  stmsToList stms+    & filter (not . flip nameIn to_remove . head . patNames . stmPat)+    & stmsFromList+    & pure  pass ::   (Mem rep inner, LetDec rep ~ LetDecMem, CanBeAliased inner) =>   String ->   String ->-  (FunDef (Aliases rep) -> Pass.PassM CoalsTab) ->-  (inner -> ReplaceM inner inner) ->+  (Prog (Aliases rep) -> Pass.PassM (M.Map Name CoalsTab)) ->+  (inner -> UpdateM inner inner) ->   (CoalsTab -> [FParam (Aliases rep)] -> (Names, [FParam (Aliases rep)])) ->   Pass rep rep pass flag desc mk on_inner on_fparams =-  Pass flag desc $-    Pass.intraproceduralTransformationWithConsts pure $ \_ f -> do-      coaltab <- mk (AnlAls.analyseFun f)+  Pass flag desc $ \prog -> do+    coaltabs <- mk $ AnlAls.aliasAnalysis prog+    Pass.intraproceduralTransformationWithConsts pure (onFun coaltabs) prog+  where+    onFun coaltabs _ f = do+      let coaltab = coaltabs M.! funDefName f       let (mem_allocs_to_remove, new_fparams) = on_fparams coaltab $ funDefParams f       pure $         f-          { funDefBody =-              onBody (foldMap vartab $ M.elems coaltab) $-                removeStms mem_allocs_to_remove $-                  funDefBody f,+          { funDefBody = onBody coaltab mem_allocs_to_remove $ funDefBody f,             funDefParams = new_fparams           }-  where-    onBody coaltab body =-      body {bodyStms = runReader (mapM replaceInStm $ bodyStms body) (Env coaltab on_inner)} -replaceInStm :: (Mem rep inner, LetDec rep ~ LetDecMem) => Stm rep -> ReplaceM inner (Stm rep)+    onBody coaltab mem_allocs_to_remove body =+      body+        { bodyStms =+            runReader+              (updateStms $ bodyStms body)+              (Env coaltab on_inner mem_allocs_to_remove),+          bodyResult = map (replaceResMem coaltab) $ bodyResult body+        }++replaceResMem :: CoalsTab -> SubExpRes -> SubExpRes+replaceResMem coaltab res =+  case flip M.lookup coaltab =<< subExpResVName res of+    Just entry -> res {resSubExp = Var $ dstmem entry}+    Nothing -> res++updateStms :: (Mem rep inner, LetDec rep ~ LetDecMem) => Stms rep -> UpdateM inner (Stms rep)+updateStms stms = do+  stms' <- mapM replaceInStm stms+  removeAllocsInStms stms'++replaceInStm :: (Mem rep inner, LetDec rep ~ LetDecMem) => Stm rep -> UpdateM inner (Stm rep) replaceInStm (Let (Pat elems) d e) = do   elems' <- mapM replaceInPatElem elems   e' <- replaceInExp elems' e   pure $ Let (Pat elems') d e'   where-    replaceInPatElem :: PatElem LetDecMem -> ReplaceM inner (PatElem LetDecMem)+    replaceInPatElem :: PatElem LetDecMem -> UpdateM inner (PatElem LetDecMem)     replaceInPatElem p@(PatElem vname (MemArray _ _ u _)) =       fromMaybe p <$> lookupAndReplace vname PatElem u     replaceInPatElem p = pure p -replaceInExp :: (Mem rep inner, LetDec rep ~ LetDecMem) => [PatElem LetDecMem] -> Exp rep -> ReplaceM inner (Exp rep)+replaceInExp :: (Mem rep inner, LetDec rep ~ LetDecMem) => [PatElem LetDecMem] -> Exp rep -> UpdateM inner (Exp rep) replaceInExp _ e@(BasicOp _) = pure e replaceInExp pat_elems (Match cond_ses cases defbody dec) = do   defbody' <- replaceInIfBody defbody@@ -106,38 +133,52 @@   pure $ Match cond_ses cases' defbody' dec' replaceInExp _ (DoLoop loop_inits loop_form (Body dec stms res)) = do   loop_inits' <- mapM (replaceInFParam . fst) loop_inits-  stms' <- mapM replaceInStm stms-  pure $ DoLoop (zip loop_inits' $ map snd loop_inits) loop_form $ Body dec stms' res-replaceInExp _ e@(Op (Alloc _ _)) = pure e-replaceInExp _ (Op (Inner i)) = do-  on_op <- asks onInner-  Op . Inner <$> on_op i-replaceInExp _ (Op _) = error "Unreachable" -- This shouldn't be possible?+  stms' <- updateStms stms+  coalstab <- asks envCoalesceTab+  let res' = map (replaceResMem coalstab) res+  pure $ DoLoop (zip loop_inits' $ map snd loop_inits) loop_form $ Body dec stms' res'+replaceInExp _ (Op op) =+  case op of+    Inner i -> do+      on_op <- asks onInner+      Op . Inner <$> on_op i+    _ -> pure $ Op op replaceInExp _ e@WithAcc {} = pure e replaceInExp _ e@Apply {} = pure e -replaceInHostOp :: HostOp GPUMem () -> ReplaceM (HostOp GPUMem ()) (HostOp GPUMem ())-replaceInHostOp (SegOp (SegMap lvl sp tps body)) = do-  stms <- mapM replaceInStm $ kernelBodyStms body-  pure $ SegOp $ SegMap lvl sp tps $ body {kernelBodyStms = stms}-replaceInHostOp (SegOp (SegRed lvl sp binops tps body)) = do-  stms <- mapM replaceInStm $ kernelBodyStms body-  pure $ SegOp $ SegRed lvl sp binops tps $ body {kernelBodyStms = stms}-replaceInHostOp (SegOp (SegScan lvl sp binops tps body)) = do-  stms <- mapM replaceInStm $ kernelBodyStms body-  pure $ SegOp $ SegScan lvl sp binops tps $ body {kernelBodyStms = stms}-replaceInHostOp (SegOp (SegHist lvl sp hist_ops tps body)) = do-  stms <- mapM replaceInStm $ kernelBodyStms body-  pure $ SegOp $ SegHist lvl sp hist_ops tps $ body {kernelBodyStms = stms}+replaceInSegOp ::+  (Mem rep inner, LetDec rep ~ LetDecMem) =>+  SegOp lvl rep ->+  UpdateM inner (SegOp lvl rep)+replaceInSegOp (SegMap lvl sp tps body) = do+  stms <- updateStms $ kernelBodyStms body+  pure $ SegMap lvl sp tps $ body {kernelBodyStms = stms}+replaceInSegOp (SegRed lvl sp binops tps body) = do+  stms <- updateStms $ kernelBodyStms body+  pure $ SegRed lvl sp binops tps $ body {kernelBodyStms = stms}+replaceInSegOp (SegScan lvl sp binops tps body) = do+  stms <- updateStms $ kernelBodyStms body+  pure $ SegScan lvl sp binops tps $ body {kernelBodyStms = stms}+replaceInSegOp (SegHist lvl sp hist_ops tps body) = do+  stms <- updateStms $ kernelBodyStms body+  pure $ SegHist lvl sp hist_ops tps $ body {kernelBodyStms = stms}++replaceInHostOp :: HostOp GPUMem () -> UpdateM (HostOp GPUMem ()) (HostOp GPUMem ())+replaceInHostOp (SegOp op) = SegOp <$> replaceInSegOp op replaceInHostOp op = pure op -generalizeIxfun :: [PatElem dec] -> PatElem LetDecMem -> BodyReturns -> ReplaceM inner BodyReturns+replaceInMCOp :: MCOp MCMem () -> UpdateM (MCOp MCMem ()) (MCOp MCMem ())+replaceInMCOp (ParOp par_op op) =+  ParOp <$> traverse replaceInSegOp par_op <*> replaceInSegOp op+replaceInMCOp op = pure op++generalizeIxfun :: [PatElem dec] -> PatElem LetDecMem -> BodyReturns -> UpdateM inner BodyReturns generalizeIxfun   pat_elems   (PatElem vname (MemArray _ _ _ (ArrayIn mem ixf)))   m@(MemArray pt shp u _) = do     coaltab <- asks envCoalesceTab-    if vname `M.member` coaltab+    if any (M.member vname . vartab) coaltab       then         existentialiseIxFun (map patElemName pat_elems) ixf           & ReturnsInBlock mem@@ -146,12 +187,13 @@       else pure m generalizeIxfun _ _ m = pure m -replaceInIfBody :: (Mem rep inner, LetDec rep ~ LetDecMem) => Body rep -> ReplaceM inner (Body rep)-replaceInIfBody b@(Body _ stms _) = do-  stms' <- mapM replaceInStm stms-  pure $ b {bodyStms = stms'}+replaceInIfBody :: (Mem rep inner, LetDec rep ~ LetDecMem) => Body rep -> UpdateM inner (Body rep)+replaceInIfBody b@(Body _ stms res) = do+  coaltab <- asks envCoalesceTab+  stms' <- updateStms stms+  pure $ b {bodyStms = stms', bodyResult = map (replaceResMem coaltab) res} -replaceInFParam :: Param FParamMem -> ReplaceM inner (Param FParamMem)+replaceInFParam :: Param FParamMem -> UpdateM inner (Param FParamMem) replaceInFParam p@(Param _ vname (MemArray _ _ u _)) = do   fromMaybe p <$> lookupAndReplace vname (Param mempty) u replaceInFParam p = pure p@@ -160,10 +202,10 @@   VName ->   (VName -> MemBound u -> a) ->   u ->-  ReplaceM inner (Maybe a)+  UpdateM inner (Maybe a) lookupAndReplace vname f u = do   coaltab <- asks envCoalesceTab-  case M.lookup vname coaltab of+  case M.lookup vname $ foldMap vartab coaltab of     Just (Coalesced _ (MemBlock pt shp mem ixf) subs) ->       ixf         & fixPoint (substituteInIxFun subs)

@@ -1,12 +1,14 @@-{-# LANGUAGE ConstraintKinds #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE LambdaCase #-}-{-# LANGUAGE TupleSections #-} {-# LANGUAGE TypeFamilies #-}  -- | The bulk of the short-circuiting implementation.-module Futhark.Optimise.ArrayShortCircuiting.ArrayCoalescing (mkCoalsTab, CoalsTab, mkCoalsTabGPU) where+module Futhark.Optimise.ArrayShortCircuiting.ArrayCoalescing+  ( mkCoalsTab,+    CoalsTab,+    mkCoalsTabGPU,+    mkCoalsTabMC,+  )+where  import Control.Exception.Base qualified as Exc import Control.Monad.Reader@@ -21,7 +23,8 @@ import Data.Set qualified as S import Futhark.Analysis.PrimExp.Convert import Futhark.IR.Aliases-import Futhark.IR.GPUMem+import Futhark.IR.GPUMem as GPU+import Futhark.IR.MCMem as MC import Futhark.IR.Mem.IxFun qualified as IxFun import Futhark.IR.SeqMem import Futhark.MonadFreshNames@@ -33,7 +36,7 @@  -- | A helper type describing representations that can be short-circuited. type Coalesceable rep inner =-  ( CreatesNewArrOp (OpWithAliases inner),+  ( Mem rep inner,     ASTRep rep,     CanBeAliased inner,     Op rep ~ MemOp inner,@@ -49,8 +52,15 @@  type ScalarTableM rep a = Reader (ComputeScalarTableOnOp rep) a -newtype ShortCircuitReader rep = ShortCircuitReader-  { onOp :: LUTabFun -> Pat (VarAliases, LetDecMem) -> Op (Aliases rep) -> TopdownEnv rep -> BotUpEnv -> ShortCircuitM rep BotUpEnv+data ShortCircuitReader rep = ShortCircuitReader+  { onOp :: LUTabFun -> Pat (VarAliases, LetDecMem) -> Op (Aliases rep) -> TopdownEnv rep -> BotUpEnv -> ShortCircuitM rep BotUpEnv,+    ssPointFromOp ::+      LUTabFun ->+      TopdownEnv rep ->+      ScopeTab rep ->+      Pat (VarAliases, LetDecMem) ->+      Op (Aliases rep) ->+      Maybe [SSPointInfo]   }  newtype ShortCircuitM rep a = ShortCircuitM (ReaderT (ShortCircuitReader rep) (State VNameSource) a)@@ -87,52 +97,66 @@ --- Main Coalescing Transformation computes a successful coalescing table    --- -------------------------------------------------------------------------------- --- | Given a 'FunDef' in 'SegMem' representation, compute the coalescing table+-- | Given a 'Prog' in 'SegMem' representation, compute the coalescing table -- by folding over each function.-mkCoalsTab :: (MonadFreshNames m) => FunDef (Aliases SeqMem) -> m CoalsTab-mkCoalsTab =-  mkCoalsTabFun-    (snd . lastUseSeqMem)-    (ShortCircuitReader shortCircuitSeqMem)+mkCoalsTab :: (MonadFreshNames m) => Prog (Aliases SeqMem) -> m (M.Map Name CoalsTab)+mkCoalsTab prog =+  mkCoalsTabProg+    (lastUseSeqMem prog)+    (ShortCircuitReader shortCircuitSeqMem genSSPointInfoSeqMem)     (ComputeScalarTableOnOp $ const $ const $ pure mempty)+    prog --- | Given a 'FunDef' in 'GPUMem' representation, compute the coalescing table+-- | Given a 'Prog' in 'GPUMem' representation, compute the coalescing table -- by folding over each function.-mkCoalsTabGPU :: (MonadFreshNames m) => FunDef (Aliases GPUMem) -> m CoalsTab-mkCoalsTabGPU =-  mkCoalsTabFun-    (snd . lastUseGPUMem)-    (ShortCircuitReader shortCircuitGPUMem)-    (ComputeScalarTableOnOp computeScalarTableGPUMem)+mkCoalsTabGPU :: (MonadFreshNames m) => Prog (Aliases GPUMem) -> m (M.Map Name CoalsTab)+mkCoalsTabGPU prog =+  mkCoalsTabProg+    (lastUseGPUMem prog)+    (ShortCircuitReader shortCircuitGPUMem genSSPointInfoGPUMem)+    (ComputeScalarTableOnOp (computeScalarTableMemOp computeScalarTableGPUMem))+    prog +-- | Given a 'Prog' in 'MCMem' representation, compute the coalescing table+-- by folding over each function.+mkCoalsTabMC :: (MonadFreshNames m) => Prog (Aliases MCMem) -> m (M.Map Name CoalsTab)+mkCoalsTabMC prog =+  mkCoalsTabProg+    (lastUseMCMem prog)+    (ShortCircuitReader shortCircuitMCMem genSSPointInfoMCMem)+    (ComputeScalarTableOnOp (computeScalarTableMemOp computeScalarTableMCMem))+    prog+ -- | Given a function, compute the coalescing table-mkCoalsTabFun ::-  (MonadFreshNames m, Coalesceable rep inner, FParamInfo rep ~ FParamMem) =>-  (FunDef (Aliases rep) -> LUTabFun) ->+mkCoalsTabProg ::+  (MonadFreshNames m, Coalesceable rep inner) =>+  LUTabProg ->   ShortCircuitReader rep ->   ComputeScalarTableOnOp rep ->-  FunDef (Aliases rep) ->-  m CoalsTab-mkCoalsTabFun lufun r computeScalarOnOp fun@(FunDef _ _ _ _ fpars body) = do-  -- First compute last-use information-  let lutab = lufun fun-      unique_mems = getUniqueMemFParam fpars-      scalar_table =-        runReader-          ( concatMapM-              (computeScalarTable $ scopeOf fun <> scopeOf (bodyStms body))-              (stmsToList $ bodyStms body)-          )-          computeScalarOnOp-      topenv =-        emptyTopdownEnv-          { scope = scopeOfFParams fpars,-            alloc = unique_mems,-            scalarTable = scalar_table,-            nonNegatives = foldMap paramSizes fpars-          }-      ShortCircuitM m = fixPointCoalesce lutab fpars body topenv-  modifyNameSource $ runState (runReaderT m r)+  Prog (Aliases rep) ->+  m (M.Map Name CoalsTab)+mkCoalsTabProg (_, lutab_prog) r computeScalarOnOp = fmap M.fromList . mapM onFun . progFuns+  where+    onFun fun@(FunDef _ _ fname _ fpars body) = do+      -- First compute last-use information+      let unique_mems = getUniqueMemFParam fpars+          lutab = lutab_prog M.! fname+          scalar_table =+            runReader+              ( concatMapM+                  (computeScalarTable $ scopeOf fun <> scopeOf (bodyStms body))+                  (stmsToList $ bodyStms body)+              )+              computeScalarOnOp+          topenv =+            emptyTopdownEnv+              { scope = scopeOfFParams fpars,+                alloc = unique_mems,+                scalarTable = scalar_table,+                nonNegatives = foldMap paramSizes fpars+              }+          ShortCircuitM m = fixPointCoalesce lutab fpars body topenv+      (fname,) <$> modifyNameSource (runState (runReaderT m r))  paramSizes :: Param FParamMem -> Names paramSizes (Param _ _ (MemArray _ shp _ _)) = freeIn shp@@ -145,28 +169,20 @@ shortCircuitSeqMem :: LUTabFun -> Pat (VarAliases, LetDecMem) -> Op (Aliases SeqMem) -> TopdownEnv SeqMem -> BotUpEnv -> ShortCircuitM SeqMem BotUpEnv shortCircuitSeqMem _ _ _ _ = pure --- | Short-circuit handler for 'GPUMem' 'Op'.------ When the 'Op' is a 'SegOp', we handle it accordingly, otherwise we do--- nothing.-shortCircuitGPUMem ::+-- | Short-circuit handler for SegOp.+shortCircuitSegOp ::+  Coalesceable rep inner =>+  (lvl -> Bool) ->   LUTabFun ->   Pat (VarAliases, LetDecMem) ->-  Op (Aliases GPUMem) ->-  TopdownEnv GPUMem ->+  SegOp lvl (Aliases rep) ->+  TopdownEnv rep ->   BotUpEnv ->-  ShortCircuitM GPUMem BotUpEnv-shortCircuitGPUMem _ _ (Alloc _ _) _ bu_env = pure bu_env-shortCircuitGPUMem lutab pat (Inner (SegOp (SegMap lvl@SegThread {} space _ kernel_body))) td_env bu_env =-  -- No special handling necessary for 'SegMap'. Just call the helper-function.-  shortCircuitGPUMemHelper 0 lvl lutab pat space kernel_body td_env bu_env-shortCircuitGPUMem lutab pat (Inner (SegOp (SegMap lvl@SegGroup {} space _ kernel_body))) td_env bu_env =-  -- No special handling necessary for 'SegMap'. Just call the helper-function.-  shortCircuitGPUMemHelper 0 lvl lutab pat space kernel_body td_env bu_env-shortCircuitGPUMem lutab pat (Inner (SegOp (SegMap lvl@SegThreadInGroup {} space _ kernel_body))) td_env bu_env =+  ShortCircuitM rep BotUpEnv+shortCircuitSegOp lvlOK lutab pat (SegMap lvl space _ kernel_body) td_env bu_env =   -- No special handling necessary for 'SegMap'. Just call the helper-function.-  shortCircuitGPUMemHelper 0 lvl lutab pat space kernel_body td_env bu_env-shortCircuitGPUMem lutab pat (Inner (SegOp (SegRed lvl space binops _ kernel_body))) td_env bu_env =+  shortCircuitSegOpHelper 0 lvlOK lvl lutab pat space kernel_body td_env bu_env+shortCircuitSegOp lvlOK lutab pat (SegRed lvl space binops _ kernel_body) td_env bu_env =   -- When handling 'SegRed', we we first invalidate all active coalesce-entries   -- where any of the variables in 'vartab' are also free in the list of   -- 'SegBinOp'. In other words, anything that is used as part of the reduction@@ -176,26 +192,26 @@         foldl markFailedCoal (activeCoals bu_env, inhibit bu_env) $ M.keys to_fail       bu_env' = bu_env {activeCoals = active, inhibit = inh}       num_reds = length red_ts-   in shortCircuitGPUMemHelper num_reds lvl lutab pat space kernel_body td_env bu_env'+   in shortCircuitSegOpHelper num_reds lvlOK lvl lutab pat space kernel_body td_env bu_env'   where     segment_dims = init $ segSpaceDims space     red_ts = do       op <- binops       let shp = Shape segment_dims <> segBinOpShape op       map (`arrayOfShape` shp) (lambdaReturnType $ segBinOpLambda op)-shortCircuitGPUMem lutab pat (Inner (SegOp (SegScan lvl space binops _ kernel_body))) td_env bu_env =+shortCircuitSegOp lvlOK lutab pat (SegScan lvl space binops _ kernel_body) td_env bu_env =   -- Like in the handling of 'SegRed', we do not want to coalesce anything that   -- is used in the 'SegBinOp'   let to_fail = M.filter (\entry -> namesFromList (M.keys $ vartab entry) `namesIntersect` foldMap (freeIn . segBinOpLambda) binops) $ activeCoals bu_env       (active, inh) = foldl markFailedCoal (activeCoals bu_env, inhibit bu_env) $ M.keys to_fail       bu_env' = bu_env {activeCoals = active, inhibit = inh}-   in shortCircuitGPUMemHelper 0 lvl lutab pat space kernel_body td_env bu_env'-shortCircuitGPUMem lutab pat (Inner (SegOp (SegHist lvl space histops _ kernel_body))) td_env bu_env = do+   in shortCircuitSegOpHelper 0 lvlOK lvl lutab pat space kernel_body td_env bu_env'+shortCircuitSegOp lvlOK lutab pat (SegHist lvl space histops _ kernel_body) td_env bu_env = do   -- Need to take zipped patterns and histDest (flattened) and insert transitive coalesces   let to_fail = M.filter (\entry -> namesFromList (M.keys $ vartab entry) `namesIntersect` foldMap (freeIn . histOp) histops) $ activeCoals bu_env       (active, inh) = foldl markFailedCoal (activeCoals bu_env, inhibit bu_env) $ M.keys to_fail       bu_env' = bu_env {activeCoals = active, inhibit = inh}-  bu_env'' <- shortCircuitGPUMemHelper 0 lvl lutab pat space kernel_body td_env bu_env'+  bu_env'' <- shortCircuitSegOpHelper 0 lvlOK lvl lutab pat space kernel_body td_env bu_env'   pure $     foldl insertHistCoals bu_env'' $       zip (patElems pat) $@@ -216,17 +232,33 @@                     }             Nothing -> acc         _ -> acc-shortCircuitGPUMem lutab pat (Inner (GPUBody _ body)) td_env bu_env = do++-- | Short-circuit handler for 'GPUMem' 'Op'.+--+-- When the 'Op' is a 'SegOp', we handle it accordingly, otherwise we do+-- nothing.+shortCircuitGPUMem ::+  LUTabFun ->+  Pat (VarAliases, LetDecMem) ->+  Op (Aliases GPUMem) ->+  TopdownEnv GPUMem ->+  BotUpEnv ->+  ShortCircuitM GPUMem BotUpEnv+shortCircuitGPUMem _ _ (Alloc _ _) _ bu_env = pure bu_env+shortCircuitGPUMem lutab pat (Inner (GPU.SegOp op)) td_env bu_env =+  shortCircuitSegOp isSegThread lutab pat op td_env bu_env+shortCircuitGPUMem lutab pat (Inner (GPU.GPUBody _ body)) td_env bu_env = do   fresh1 <- newNameFromString "gpubody"   fresh2 <- newNameFromString "gpubody"-  shortCircuitGPUMemHelper+  shortCircuitSegOpHelper     0+    isSegThread     -- Construct a 'SegLevel' corresponding to a single thread-    ( SegThread SegNoVirt $+    ( GPU.SegThread GPU.SegNoVirt $         Just $-          KernelGrid-            (Count $ Constant $ IntValue $ Int64Value 1)-            (Count $ Constant $ IntValue $ Int64Value 1)+          GPU.KernelGrid+            (GPU.Count $ Constant $ IntValue $ Int64Value 1)+            (GPU.Count $ Constant $ IntValue $ Int64Value 1)     )     lutab     pat@@ -234,14 +266,29 @@     (bodyToKernelBody body)     td_env     bu_env-shortCircuitGPUMem _ _ (Inner (SizeOp _)) _ bu_env = pure bu_env-shortCircuitGPUMem _ _ (Inner (OtherOp ())) _ bu_env = pure bu_env+shortCircuitGPUMem _ _ (Inner (GPU.SizeOp _)) _ bu_env = pure bu_env+shortCircuitGPUMem _ _ (Inner (GPU.OtherOp ())) _ bu_env = pure bu_env +shortCircuitMCMem ::+  LUTabFun ->+  Pat (VarAliases, LetDecMem) ->+  Op (Aliases MCMem) ->+  TopdownEnv MCMem ->+  BotUpEnv ->+  ShortCircuitM MCMem BotUpEnv+shortCircuitMCMem _ _ (Alloc _ _) _ bu_env = pure bu_env+shortCircuitMCMem _ _ (Inner (MC.OtherOp ())) _ bu_env = pure bu_env+shortCircuitMCMem lutab pat (Inner (MC.ParOp (Just par_op) op)) td_env bu_env =+  shortCircuitSegOp (const True) lutab pat par_op td_env bu_env+    >>= shortCircuitSegOp (const True) lutab pat op td_env+shortCircuitMCMem lutab pat (Inner (MC.ParOp Nothing op)) td_env bu_env =+  shortCircuitSegOp (const True) lutab pat op td_env bu_env+ dropLastSegSpace :: SegSpace -> SegSpace dropLastSegSpace space = space {unSegSpace = init $ unSegSpace space} -isSegThread :: SegLevel -> Bool-isSegThread SegThread {} = True+isSegThread :: GPU.SegLevel -> Bool+isSegThread GPU.SegThread {} = True isSegThread _ = False  -- | Computes the slice written at the end of a thread in a 'SegOp'.@@ -282,18 +329,21 @@ -- -- 4. Mark active coalescings as finished, since a 'SegOp' is an array creation -- point.-shortCircuitGPUMemHelper ::+shortCircuitSegOpHelper ::+  Coalesceable rep inner =>   -- | The number of returns for which we should drop the last seg space   Int ->-  SegLevel ->+  -- | Whether we should look at a segop with this lvl.+  (lvl -> Bool) ->+  lvl ->   LUTabFun ->   Pat (VarAliases, LetDecMem) ->   SegSpace ->-  KernelBody (Aliases GPUMem) ->-  TopdownEnv GPUMem ->+  KernelBody (Aliases rep) ->+  TopdownEnv rep ->   BotUpEnv ->-  ShortCircuitM GPUMem BotUpEnv-shortCircuitGPUMemHelper num_reds lvl lutab pat@(Pat ps0) space0 kernel_body td_env bu_env = do+  ShortCircuitM rep BotUpEnv+shortCircuitSegOpHelper num_reds lvlOK lvl lutab pat@(Pat ps0) space0 kernel_body td_env bu_env = do   -- We need to drop the last element of the 'SegSpace' for pattern elements   -- that correspond to reductions.   let ps_space_and_res =@@ -311,7 +361,7 @@       (actv_return, inhibit_return) =         if num_reds > 0           then (actv0, inhibit0)-          else foldl (makeSegMapCoals lvl td_env kernel_body) (actv0, inhibit0) ps_space_and_res+          else foldl (makeSegMapCoals lvlOK lvl td_env kernel_body) (actv0, inhibit0) ps_space_and_res    -- Start from empty references, we'll update with aggregates later.   let actv0' = M.map (\etry -> etry {memrefs = mempty}) $ actv0 <> actv_return@@ -442,67 +492,61 @@  -- | Given a pattern element and the corresponding kernel result, try to put the -- kernel result directly in the memory block of pattern element-makeSegMapCoals :: SegLevel -> TopdownEnv GPUMem -> KernelBody (Aliases GPUMem) -> (CoalsTab, InhibitTab) -> (PatElem (VarAliases, LetDecMem), SegSpace, KernelResult) -> (CoalsTab, InhibitTab)-makeSegMapCoals lvl td_env kernel_body (active, inhb) (PatElem pat_name (_, MemArray _ _ _ (ArrayIn pat_mem pat_ixf)), space, Returns _ _ (Var return_name))-  | Just mb@(MemBlock tp return_shp return_mem _) <-+makeSegMapCoals ::+  (Coalesceable rep inner) =>+  (lvl -> Bool) ->+  lvl ->+  TopdownEnv rep ->+  KernelBody (Aliases rep) ->+  (CoalsTab, InhibitTab) ->+  (PatElem (VarAliases, LetDecMem), SegSpace, KernelResult) ->+  (CoalsTab, InhibitTab)+makeSegMapCoals lvlOK lvl td_env kernel_body (active, inhb) (PatElem pat_name (_, MemArray _ _ _ (ArrayIn pat_mem pat_ixf)), space, Returns _ _ (Var return_name))+  | Just (MemBlock tp return_shp return_mem _) <-       getScopeMemInfo return_name $ scope td_env <> scopeOf (kernelBodyStms kernel_body),-    isSegThread lvl,+    lvlOK lvl,     MemMem pat_space <- runReader (lookupMemInfo pat_mem) $ removeScopeAliases $ scope td_env,-    MemMem return_space <- runReader (lookupMemInfo return_mem) $ removeScopeAliases $ scope td_env <> scopeOf (kernelBodyStms kernel_body) <> scopeOfSegSpace space,+    MemMem return_space <-+      scope td_env <> scopeOf (kernelBodyStms kernel_body) <> scopeOfSegSpace space+        & removeScopeAliases+        & runReader (lookupMemInfo return_mem),     pat_space == return_space =       case M.lookup pat_mem active of         Nothing ->           -- We are not in a transitive case-          if IxFun.hasOneLmad pat_ixf-            then case ( maybe False (pat_mem `nameIn`) $ M.lookup return_mem inhb,-                        Coalesced InPlaceCoal mb mempty-                          & M.singleton return_name-                          & flip (addInvAliassesVarTab td_env) return_name-                          & fmap-                            ( M.adjust-                                ( \(Coalesced knd (MemBlock pt shp _ _) subst) ->-                                    Coalesced-                                      knd-                                      ( MemBlock pt shp pat_mem $-                                          IxFun.slice pat_ixf $-                                            fullSlice (IxFun.shape pat_ixf) $-                                              Slice $-                                                map (DimFix . TPrimExp . flip LeafExp (IntType Int64) . fst) $-                                                  unSegSpace space-                                      )-                                      subst-                                )-                                return_name-                            )-                      ) of-              (False, Just vtab) ->-                (active <> M.singleton return_mem (CoalsEntry pat_mem pat_ixf (oneName pat_mem) vtab mempty mempty), inhb)-              _ -> (active, inhb)-            else (active, inhb)-        Just trans ->-          case ( maybe False (dstmem trans `nameIn`) $ M.lookup return_mem inhb,-                 Coalesced InPlaceCoal (MemBlock tp return_shp (dstmem trans) (dstind trans)) mempty+          case ( IxFun.hasOneLmad pat_ixf+                   && maybe False (pat_mem `nameIn`) (M.lookup return_mem inhb),+                 Coalesced+                   InPlaceCoal+                   (MemBlock tp return_shp pat_mem $ resultSlice pat_ixf)+                   mempty                    & M.singleton return_name                    & flip (addInvAliassesVarTab td_env) return_name-                   & fmap-                     ( M.adjust-                         ( \(Coalesced knd (MemBlock pt shp mem ixf@(IxFun.IxFun _ base_shape _)) subst) ->-                             Coalesced-                               knd-                               ( MemBlock pt shp mem $-                                   IxFun.slice ixf $-                                     fullSlice base_shape $-                                       Slice $-                                         map (DimFix . TPrimExp . flip LeafExp (IntType Int64) . fst) $-                                           unSegSpace space-                               )-                               subst-                         )-                         return_name-                     )                ) of             (False, Just vtab) ->-              let opts = if dstmem trans == pat_mem then mempty else M.insert pat_name pat_mem $ optdeps trans+              ( active+                  <> M.singleton+                    return_mem+                    (CoalsEntry pat_mem pat_ixf (oneName pat_mem) vtab mempty mempty),+                inhb+              )+            _ -> (active, inhb)+        Just trans ->+          case ( maybe False (dstmem trans `nameIn`) $ M.lookup return_mem inhb,+                 let Coalesced _ (MemBlock _ _ trans_mem trans_ixf) _ =+                       fromMaybe (error "Impossible") $ M.lookup pat_name $ vartab trans+                  in Coalesced+                       TransitiveCoal+                       (MemBlock tp return_shp trans_mem $ resultSlice trans_ixf)+                       mempty+                       & M.singleton return_name+                       & flip (addInvAliassesVarTab td_env) return_name+               ) of+            (False, Just vtab) ->+              let opts =+                    if dstmem trans == pat_mem+                      then mempty+                      else M.insert pat_name pat_mem $ optdeps trans                in ( M.insert                       return_mem                       ( CoalsEntry@@ -517,22 +561,28 @@                     inhb                   )             _ -> (active, inhb)-makeSegMapCoals _ td_env _ x (_, _, WriteReturns _ _ return_name _) =+  where+    thread_slice =+      unSegSpace space+        & map (DimFix . TPrimExp . flip LeafExp (IntType Int64) . fst)+        & Slice+    resultSlice ixf = IxFun.slice ixf $ fullSlice (IxFun.shape ixf) thread_slice+makeSegMapCoals _ _ td_env _ x (_, _, WriteReturns _ _ return_name _) =   case getScopeMemInfo return_name $ scope td_env of     Just (MemBlock _ _ return_mem _) -> markFailedCoal x return_mem     Nothing -> error "Should not happen?"-makeSegMapCoals _ td_env _ x (_, _, result) =+makeSegMapCoals _ _ td_env _ x (_, _, result) =   freeIn result     & namesToList     & mapMaybe (flip getScopeMemInfo $ scope td_env)-    & foldr (\(MemBlock _ _ mem _) -> flip markFailedCoal mem) x+    & foldr (flip markFailedCoal . memName) x  fullSlice :: [TPrimExp Int64 VName] -> Slice (TPrimExp Int64 VName) -> Slice (TPrimExp Int64 VName) fullSlice shp (Slice slc) =   Slice $ slc ++ map (\d -> DimSlice 0 d 1) (drop (length slc) shp)  fixPointCoalesce ::-  (Coalesceable rep inner) =>+  Coalesceable rep inner =>   LUTabFun ->   [Param FParamMem] ->   Body (Aliases rep) ->@@ -593,7 +643,7 @@  -- | Perform short-circuiting on 'Stms'. mkCoalsTabStms ::-  (Coalesceable rep inner) =>+  Coalesceable rep inner =>   LUTabFun ->   Stms (Aliases rep) ->   TopdownEnv rep ->@@ -636,7 +686,7 @@ --                          then the checks should be extended to the actual --                          array-creation points. mkCoalsTabStm ::-  (Coalesceable rep inner) =>+  Coalesceable rep inner =>   LUTabFun ->   Stm (Aliases rep) ->   TopdownEnv rep ->@@ -1037,7 +1087,7 @@  -- The case of in-place update: --   @let x' = x with slice <- elm@-mkCoalsTabStm lutab stm@(Let pat@(Pat [x']) _ e@(BasicOp (Update safety x _ _elm))) td_env bu_env+mkCoalsTabStm lutab stm@(Let pat@(Pat [x']) _ (BasicOp (Update safety x _ _elm))) td_env bu_env   | [(_, MemBlock _ _ m_x _)] <- getArrMemAssoc pat =       do         -- (a) filter by the 3rd safety for @elm@ and @x'@@@ -1065,14 +1115,14 @@                         _ ->                           markFailedCoal (actv, inhbt) m_x -            -- (c) this stm is also a potential source for coalescing, so process it-            actv'' = if safety == Unsafe then mkCoalsHelper3PatternMatch pat e lutab td_env (successCoals bu_env) actv' inhbt' else actv'+        -- (c) this stm is also a potential source for coalescing, so process it+        actv'' <- if safety == Unsafe then mkCoalsHelper3PatternMatch stm lutab td_env {inhibited = inhbt'} bu_env {activeCoals = actv'} else pure actv'         pure $           bu_env {activeCoals = actv'', inhibit = inhbt'}  -- The case of flat in-place update: --   @let x' = x with flat-slice <- elm@-mkCoalsTabStm lutab stm@(Let pat@(Pat [x']) _ e@(BasicOp (FlatUpdate x _ _elm))) td_env bu_env+mkCoalsTabStm lutab stm@(Let pat@(Pat [x']) _ (BasicOp (FlatUpdate x _ _elm))) td_env bu_env   | [(_, MemBlock _ _ m_x _)] <- getArrMemAssoc pat =       do         -- (a) filter by the 3rd safety for @elm@ and @x'@@@ -1102,18 +1152,25 @@                         _ ->                           markFailedCoal (actv, inhbt) m_x -            -- (c) this stm is also a potential source for coalescing, so process it-            actv'' = mkCoalsHelper3PatternMatch pat e lutab td_env (successCoals bu_env) actv' inhbt'+        -- (c) this stm is also a potential source for coalescing, so process it+        actv'' <- mkCoalsHelper3PatternMatch stm lutab td_env {inhibited = inhbt'} bu_env {activeCoals = actv'}         pure $           bu_env {activeCoals = actv'', inhibit = inhbt'} -- mkCoalsTabStm _ (Let pat _ (BasicOp Update {})) _ _ =   error $ "In ArrayCoalescing.hs, fun mkCoalsTabStm, illegal pattern for in-place update: " ++ show pat -- default handling-mkCoalsTabStm lutab (Let pat _ (Op op)) td_env bu_env = do+mkCoalsTabStm lutab stm@(Let pat _ (Op op)) td_env bu_env = do   -- Process body   on_op <- asks onOp-  on_op lutab pat op td_env bu_env+  activeCoals' <-+    mkCoalsHelper3PatternMatch+      stm+      lutab+      td_env+      bu_env+  let bu_env' = bu_env {activeCoals = activeCoals'}+  on_op lutab pat op td_env bu_env' mkCoalsTabStm lutab stm@(Let pat _ e) td_env bu_env = do   --   i) Filter @activeCoals@ by the 3rd safety condition:   --      this is now relaxed by use of LMAD eqs:@@ -1132,8 +1189,8 @@       ((activeCoals'', inhibit''), successCoals') =         foldl (foldfun safe_4) ((activeCoals', inhibit'), successCoals bu_env) (getArrMemAssoc pat) -      -- iii) record a potentially coalesced statement in @activeCoals@-      activeCoals''' = mkCoalsHelper3PatternMatch pat e lutab td_env successCoals' activeCoals'' (inhibited td_env)+  -- iii) record a potentially coalesced statement in @activeCoals@+  activeCoals''' <- mkCoalsHelper3PatternMatch stm lutab td_env bu_env {successCoals = successCoals', activeCoals = activeCoals''}   pure bu_env {activeCoals = activeCoals''', inhibit = inhibit'', successCoals = successCoals'}   where     foldfun safe_4 ((a_acc, inhb), s_acc) (b, MemBlock tp shp mb _b_indfun) =@@ -1244,22 +1301,20 @@ -- |   Pattern matches a potentially coalesced statement and --     records a new association in @activeCoals@ mkCoalsHelper3PatternMatch ::-  HasMemBlock (Aliases rep) =>-  Pat (VarAliases, LetDecMem) ->-  Exp (Aliases rep) ->+  Coalesceable rep inner =>+  Stm (Aliases rep) ->   LUTabFun ->   TopdownEnv rep ->-  CoalsTab ->-  CoalsTab ->-  InhibitTab ->-  CoalsTab-mkCoalsHelper3PatternMatch pat e lutab td_env _ activeCoals_tab _-  | Nothing <- genCoalStmtInfo lutab (scope td_env) pat e =-      activeCoals_tab-mkCoalsHelper3PatternMatch pat e lutab td_env successCoals_tab activeCoals_tab inhibit_tab-  | Just clst <- genCoalStmtInfo lutab (scope td_env) pat e =-      foldl processNewCoalesce activeCoals_tab clst+  BotUpEnv ->+  ShortCircuitM rep CoalsTab+mkCoalsHelper3PatternMatch stm lutab td_env bu_env = do+  clst <- genCoalStmtInfo lutab td_env (scope td_env) stm+  case clst of+    Nothing -> pure activeCoals_tab+    Just clst' -> pure $ foldl processNewCoalesce activeCoals_tab clst'   where+    successCoals_tab = successCoals bu_env+    activeCoals_tab = activeCoals bu_env     processNewCoalesce acc (knd, alias_fn, x, m_x, ind_x, b, m_b, _, tp_b, shp_b) =       -- test whether we are in a transitive coalesced case, i.e.,       --      @let b = scratch ...@@@ -1286,18 +1341,18 @@                  in (m_y, alias_fn ind, oneName m_x <> y_al, x_deps0)           success0 = IxFun.hasOneLmad ind_yx           m_b_aliased_m_yx = areAnyAliased td_env m_b [m_yx] -- m_b \= m_yx-       in case (success0, not m_b_aliased_m_yx, isInScope td_env m_yx) of -- nameIn m_yx (alloc td_env)-            (True, True, True) ->-              -- Finally update the @activeCoals@ table with a fresh-              --   binding for @m_b@; if such one exists then overwrite.-              -- Also, add all variables from the alias chain of @b@ to-              --   @vartab@, for example, in the case of a sequence:-              --   @ b0 = if cond then ... else ... @-              --   @ b1 = alias0 b0 @-              --   @ b  = alias1 b1 @-              --   @ x[j] = b @-              -- Then @b1@ and @b0@ should also be added to @vartab@ if-              --   @alias1@ and @alias0@ are invertible, otherwise fail early!+       in if success0 && not m_b_aliased_m_yx && isInScope td_env m_yx -- nameIn m_yx (alloc td_env)+      -- Finally update the @activeCoals@ table with a fresh+      --   binding for @m_b@; if such one exists then overwrite.+      -- Also, add all variables from the alias chain of @b@ to+      --   @vartab@, for example, in the case of a sequence:+      --   @ b0 = if cond then ... else ... @+      --   @ b1 = alias0 b0 @+      --   @ b  = alias1 b1 @+      --   @ x[j] = b @+      -- Then @b1@ and @b0@ should also be added to @vartab@ if+      --   @alias1@ and @alias0@ are invertible, otherwise fail early!+            then               let mem_info = Coalesced knd (MemBlock tp_b shp_b m_yx ind_yx) M.empty                   opts' =                     if m_yx == m_x@@ -1306,7 +1361,7 @@                   vtab = M.singleton b mem_info                   mvtab = addInvAliassesVarTab td_env vtab b -                  is_inhibited = case M.lookup m_b inhibit_tab of+                  is_inhibited = case M.lookup m_b $ inhibited td_env of                     Just nms -> m_yx `nameIn` nms                     Nothing -> False                in case (is_inhibited, mvtab) of@@ -1323,30 +1378,139 @@                               opts'                               mempty                        in M.insert m_b coal_etry acc-            _ -> acc-mkCoalsHelper3PatternMatch _ _ _ _ _ _ _ =-  error "In ArrayCoalescing.hs, fun mkCoalsHelper3PatternMatch: Unreachable!!!"+            else acc -genCoalStmtInfo ::-  HasMemBlock (Aliases rep) =>+-- | Information about a particular short-circuit point+type SSPointInfo =+  ( CoalescedKind,+    IxFun -> IxFun,+    VName,+    VName,+    IxFun,+    VName,+    VName,+    IxFun,+    PrimType,+    Shape+  )++-- | Given an op, return a list of potential short-circuit points+type GenSSPoint rep op =   LUTabFun ->+  TopdownEnv rep ->   ScopeTab rep ->   Pat (VarAliases, LetDecMem) ->-  Exp (Aliases rep) ->-  Maybe [(CoalescedKind, IxFun -> IxFun, VName, VName, IxFun, VName, VName, IxFun, PrimType, Shape)]+  op ->+  Maybe [SSPointInfo]++genSSPointInfoSeqMem ::+  GenSSPoint SeqMem (Op (Aliases SeqMem))+genSSPointInfoSeqMem _ _ _ _ _ =+  Nothing++-- | For 'SegOp', we currently only handle 'SegMap', and only under the following+-- circumstances:+--+--  1. The 'SegMap' has only one return/pattern value.+--+--  2. The 'KernelBody' contains an 'Index' statement that is indexing an array using+--  only the values from the 'SegSpace'.+--+--  3. The array being indexed is last-used in that statement, is free in the+--  'SegMap', is unique or has been recently allocated (specifically, it should+--  not be a non-unique argument to the enclosing function), has elements with+--  the same bit-size as the pattern elements, and has the exact same 'IxFun' as+--  the pattern of the 'SegMap' statement.+--+-- There can be multiple candidate arrays, but the current implementation will+-- always just try the first one.+--+-- The first restriction could be relaxed by trying to match up arrays in the+-- 'KernelBody' with patterns of the 'SegMap', but the current implementation+-- should be enough to handle many common cases.+--+-- The result of the 'SegMap' is treated as the destination, while the candidate+-- array from inside the body is treated as the source.+genSSPointInfoSegOp ::+  Coalesceable rep inner => GenSSPoint rep (SegOp lvl (Aliases rep))+genSSPointInfoSegOp+  lutab+  td_env+  scopetab+  (Pat [PatElem dst (_, MemArray dst_pt _ _ (ArrayIn dst_mem dst_ixf))])+  (SegMap _ space _ kernel_body)+    | (src, MemBlock _ shp src_mem src_ixf) : _ <-+        mapMaybe getPotentialMapShortCircuit $+          stmsToList $+            kernelBodyStms kernel_body =+        Just [(MapCoal, id, dst, dst_mem, dst_ixf, src, src_mem, src_ixf, dst_pt, shp)]+    where+      iterators = map fst $ unSegSpace space+      frees = freeIn kernel_body++      getPotentialMapShortCircuit (Let (Pat [PatElem x _]) _ (BasicOp (Index src slc)))+        | Just inds <- sliceIndices slc,+          L.sort inds == L.sort (map Var iterators),+          Just last_uses <- M.lookup x lutab,+          src `nameIn` last_uses,+          Just memblock@(MemBlock src_pt _ src_mem src_ixf) <-+            getScopeMemInfo src scopetab,+          src_mem `nameIn` last_uses,+          -- The 'alloc' table contains allocated memory blocks, including+          -- unique memory blocks from the enclosing function. It does _not_+          -- include non-unique memory blocks from the enclosing function.+          src_mem `M.member` alloc td_env,+          src `nameIn` frees,+          src_ixf == dst_ixf,+          primBitSize src_pt == primBitSize dst_pt =+            Just (src, memblock)+      getPotentialMapShortCircuit _ = Nothing+genSSPointInfoSegOp _ _ _ _ _ =+  Nothing++genSSPointInfoMemOp ::+  GenSSPoint rep inner ->+  GenSSPoint rep (MemOp inner)+genSSPointInfoMemOp onOp lutab td_end scopetab pat (Inner op) =+  onOp lutab td_end scopetab pat op+genSSPointInfoMemOp _ _ _ _ _ _ = Nothing++genSSPointInfoGPUMem ::+  GenSSPoint GPUMem (Op (Aliases GPUMem))+genSSPointInfoGPUMem = genSSPointInfoMemOp f+  where+    f lutab td_env scopetab pat (GPU.SegOp op) =+      genSSPointInfoSegOp lutab td_env scopetab pat op+    f _ _ _ _ _ = Nothing++genSSPointInfoMCMem ::+  GenSSPoint MCMem (Op (Aliases MCMem))+genSSPointInfoMCMem = genSSPointInfoMemOp f+  where+    f lutab td_env scopetab pat (MC.ParOp Nothing op) =+      genSSPointInfoSegOp lutab td_env scopetab pat op+    f _ _ _ _ _ = Nothing++genCoalStmtInfo ::+  Coalesceable rep inner =>+  LUTabFun ->+  TopdownEnv rep ->+  ScopeTab rep ->+  Stm (Aliases rep) ->+  ShortCircuitM rep (Maybe [SSPointInfo]) -- CASE a) @let x <- copy(b^{lu})@-genCoalStmtInfo lutab scopetab pat (BasicOp (Copy b))+genCoalStmtInfo lutab _ scopetab (Let pat _ (BasicOp (Copy b)))   | Pat [PatElem x (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat =-      case (M.lookup x lutab, getScopeMemInfo b scopetab) of+      pure $ case (M.lookup x lutab, getScopeMemInfo b scopetab) of         (Just last_uses, Just (MemBlock tpb shpb m_b ind_b)) ->           if b `notNameIn` last_uses             then Nothing             else Just [(CopyCoal, id, x, m_x, ind_x, b, m_b, ind_b, tpb, shpb)]         _ -> Nothing -- CASE c) @let x[i] = b^{lu}@-genCoalStmtInfo lutab scopetab pat (BasicOp (Update _ x slice_x (Var b)))+genCoalStmtInfo lutab _ scopetab (Let pat _ (BasicOp (Update _ x slice_x (Var b))))   | Pat [PatElem x' (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat =-      case (M.lookup x' lutab, getScopeMemInfo b scopetab) of+      pure $ case (M.lookup x' lutab, getScopeMemInfo b scopetab) of         (Just last_uses, Just (MemBlock tpb shpb m_b ind_b)) ->           if b `notNameIn` last_uses             then Nothing@@ -1357,9 +1521,9 @@     updateIndFunSlice ind_fun slc_x =       let slc_x' = map (fmap pe64) $ unSlice slc_x        in IxFun.slice ind_fun $ Slice slc_x'-genCoalStmtInfo lutab scopetab pat (BasicOp (FlatUpdate x slice_x b))+genCoalStmtInfo lutab _ scopetab (Let pat _ (BasicOp (FlatUpdate x slice_x b)))   | Pat [PatElem x' (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat =-      case (M.lookup x' lutab, getScopeMemInfo b scopetab) of+      pure $ case (M.lookup x' lutab, getScopeMemInfo b scopetab) of         (Just last_uses, Just (MemBlock tpb shpb m_b ind_b)) ->           if b `notNameIn` last_uses             then Nothing@@ -1371,9 +1535,9 @@       IxFun.flatSlice ind_fun $ FlatSlice (pe64 offset) $ map (fmap pe64) dims  -- CASE b) @let x = concat(a, b^{lu})@-genCoalStmtInfo lutab scopetab pat (BasicOp (Concat concat_dim (b0 :| bs) _))+genCoalStmtInfo lutab _ scopetab (Let pat _ (BasicOp (Concat concat_dim (b0 :| bs) _)))   | Pat [PatElem x (_, MemArray _ _ _ (ArrayIn m_x ind_x))] <- pat =-      case M.lookup x lutab of+      pure $ case M.lookup x lutab of         Nothing -> Nothing         Just last_uses ->           let zero = pe64 $ intConst Int64 0@@ -1399,8 +1563,13 @@                     _ -> (acc, offs, False)               (res, _, _) = foldl markConcatParts ([], zero, True) (b0 : bs)            in if null res then Nothing else Just res--- CASE other than a), b), or c) not supported-genCoalStmtInfo _ _ _ _ = Nothing+-- case d) short-circuit points from ops. For instance, the result of a segmap+-- can be considered a short-circuit point.+genCoalStmtInfo lutab td_env scopetab (Let pat _ (Op op)) = do+  ss_op <- asks ssPointFromOp+  pure $ ss_op lutab td_env scopetab pat op+-- CASE other than a), b), c), or d) not supported+genCoalStmtInfo _ _ _ _ = pure Nothing  data MemBodyResult = MemBodyResult   { patMem :: VName,@@ -1497,7 +1666,7 @@     mki64subst _ = Nothing  computeScalarTable ::-  (Coalesceable rep inner) =>+  Coalesceable rep inner =>   ScopeTab rep ->   Stm (Aliases rep) ->   ScalarTableM rep (M.Map VName (PrimExp VName))@@ -1531,18 +1700,44 @@   on_op scope_table op computeScalarTable _ _ = pure mempty -computeScalarTableGPUMem :: ScopeTab GPUMem -> Op (Aliases GPUMem) -> ScalarTableM GPUMem (M.Map VName (PrimExp VName))-computeScalarTableGPUMem _ (Alloc _ _) = pure mempty-computeScalarTableGPUMem scope_table (Inner (SegOp segop)) = do+type ComputeScalarTable rep op =+  ScopeTab rep -> op -> ScalarTableM rep (M.Map VName (PrimExp VName))++computeScalarTableMemOp ::+  ComputeScalarTable rep inner -> ComputeScalarTable rep (MemOp inner)+computeScalarTableMemOp _ _ (Alloc _ _) = pure mempty+computeScalarTableMemOp onInner scope_table (Inner op) = onInner scope_table op++computeScalarTableSegOp ::+  Coalesceable rep inner =>+  ComputeScalarTable rep (GPU.SegOp lvl (Aliases rep))+computeScalarTableSegOp scope_table segop = do   concatMapM-    (computeScalarTable $ scope_table <> scopeOf (kernelBodyStms $ segBody segop) <> scopeOfSegSpace (segSpace segop))+    ( computeScalarTable $+        scope_table+          <> scopeOf (kernelBodyStms $ segBody segop)+          <> scopeOfSegSpace (segSpace segop)+    )     (stmsToList $ kernelBodyStms $ segBody segop)-computeScalarTableGPUMem _ (Inner (SizeOp _)) = pure mempty-computeScalarTableGPUMem _ (Inner (OtherOp ())) = pure mempty-computeScalarTableGPUMem scope_table (Inner (GPUBody _ body)) =++computeScalarTableGPUMem ::+  ComputeScalarTable GPUMem (GPU.HostOp (Aliases GPUMem) ())+computeScalarTableGPUMem scope_table (GPU.SegOp segop) =+  computeScalarTableSegOp scope_table segop+computeScalarTableGPUMem _ (GPU.SizeOp _) = pure mempty+computeScalarTableGPUMem _ (GPU.OtherOp ()) = pure mempty+computeScalarTableGPUMem scope_table (GPU.GPUBody _ body) =   concatMapM     (computeScalarTable $ scope_table <> scopeOf (bodyStms body))     (stmsToList $ bodyStms body)++computeScalarTableMCMem ::+  ComputeScalarTable MCMem (MC.MCOp (Aliases MCMem) ())+computeScalarTableMCMem _ (MC.OtherOp ()) = pure mempty+computeScalarTableMCMem scope_table (MC.ParOp par_op segop) =+  (<>)+    <$> maybe (pure mempty) (computeScalarTableSegOp scope_table) par_op+    <*> computeScalarTableSegOp scope_table segop  filterMapM1 :: (Eq k, Monad m) => (v -> m Bool) -> M.Map k v -> m (M.Map k v) filterMapM1 f m = fmap M.fromAscList $ filterM (f . snd) $ M.toAscList m

@@ -8,11 +8,7 @@     CoalescedKind (..),     ArrayMemBound (..),     AllocTab,-    AliasTab,-    LUTabFun,-    CreatesNewArrOp,     HasMemBlock,-    LUTabPrg,     ScalarTab,     CoalsTab,     ScopeTab,@@ -42,7 +38,8 @@ import Data.Maybe import Data.Set qualified as S import Futhark.IR.Aliases-import Futhark.IR.GPUMem+import Futhark.IR.GPUMem as GPU+import Futhark.IR.MCMem as MC import Futhark.IR.Mem.IxFun qualified as IxFun import Futhark.IR.SeqMem import Futhark.Util.Pretty hiding (line, sep, (</>))@@ -106,6 +103,7 @@     ConcatCoal   | -- | transitive, i.e., other variables aliased with b.     TransitiveCoal+  | MapCoal  -- | Information about a memory block: type, shape, name and ixfun. data ArrayMemBound = MemBlock@@ -169,15 +167,6 @@ type AllocTab = M.Map VName Space -- ^ the allocatted memory blocks -type AliasTab = M.Map VName Names--- ^ maps a variable or memory block to its aliases--type LUTabFun = M.Map VName Names--- ^ maps a name indentifying a stmt to the last uses in that stmt--type LUTabPrg = M.Map Name LUTabFun--- ^ maps function names to last-use tables- type ScalarTab = M.Map VName (PrimExp VName) -- ^ maps a variable name to its PrimExp scalar expression @@ -219,6 +208,7 @@   pretty InPlaceCoal = "InPlace"   pretty ConcatCoal = "Concat"   pretty TransitiveCoal = "Transitive"+  pretty MapCoal = "Map"  instance Pretty ArrayMemBound where   pretty (MemBlock ptp shp m_nm ixfn) =@@ -324,8 +314,16 @@       Just (LParamName (MemArray tp shp _ (ArrayIn m idx))) -> Just (MemBlock tp shp m idx)       _ -> Nothing +instance HasMemBlock (Aliases MCMem) where+  getScopeMemInfo r scope_env0 =+    case M.lookup r scope_env0 of+      Just (LetName (_, MemArray tp shp _ (ArrayIn m idx))) -> Just (MemBlock tp shp m idx)+      Just (FParamName (MemArray tp shp _ (ArrayIn m idx))) -> Just (MemBlock tp shp m idx)+      Just (LParamName (MemArray tp shp _ (ArrayIn m idx))) -> Just (MemBlock tp shp m idx)+      _ -> Nothing+ -- | @True@ if the expression returns a "fresh" array.-createsNewArrOK :: CreatesNewArrOp (Op rep) => Exp rep -> Bool+createsNewArrOK :: Exp rep -> Bool createsNewArrOK (BasicOp Replicate {}) = True createsNewArrOK (BasicOp Iota {}) = True createsNewArrOK (BasicOp Manifest {}) = True@@ -334,28 +332,7 @@ createsNewArrOK (BasicOp ArrayLit {}) = True createsNewArrOK (BasicOp Scratch {}) = True createsNewArrOK (BasicOp Rotate {}) = True-createsNewArrOK (Op op) = createsNewArrOp op createsNewArrOK _ = False--class CreatesNewArrOp rep where-  createsNewArrOp :: rep -> Bool--instance CreatesNewArrOp () where-  createsNewArrOp () = False--instance CreatesNewArrOp inner => CreatesNewArrOp (MemOp inner) where-  createsNewArrOp (Alloc _ _) = True-  createsNewArrOp (Inner inner) = createsNewArrOp inner--instance CreatesNewArrOp inner => CreatesNewArrOp (HostOp (Aliases GPUMem) inner) where-  createsNewArrOp (OtherOp op) = createsNewArrOp op-  createsNewArrOp (SegOp (SegMap _ _ _ kbody)) = all isReturns $ kernelBodyResult kbody-  createsNewArrOp (SizeOp _) = False-  createsNewArrOp _ = undefined--isReturns :: KernelResult -> Bool-isReturns Returns {} = True-isReturns _ = False  -- | Memory-block removal from active-coalescing table --   should only be handled via this function, it is easy

@@ -1,6 +1,5 @@-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE PartialTypeSignatures #-} {-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}  -- | Last use analysis for array short circuiting --@@ -15,57 +14,122 @@ -- This pass is different from "Futhark.Analysis.LastUse" in that memory blocks -- are used to alias arrays. For instance, an 'Update' will not result in a last -- use of the array being updated, because the result lives in the same memory.-module Futhark.Optimise.ArrayShortCircuiting.LastUse (lastUseSeqMem, lastUsePrg, lastUsePrgGPU, lastUseGPUMem) where+module Futhark.Optimise.ArrayShortCircuiting.LastUse+  ( lastUseSeqMem,+    lastUseGPUMem,+    lastUseMCMem,+    LUTabFun,+    LUTabProg,+  )+where  import Control.Monad.Reader import Control.Monad.State.Strict import Data.Bifunctor (bimap)+import Data.Function ((&)) import Data.Map.Strict qualified as M import Data.Maybe import Data.Sequence (Seq (..)) import Futhark.IR.Aliases import Futhark.IR.GPUMem+import Futhark.IR.GPUMem qualified as GPU+import Futhark.IR.MCMem+import Futhark.IR.MCMem qualified as MC import Futhark.IR.SeqMem import Futhark.Optimise.ArrayShortCircuiting.DataStructs import Futhark.Util +-- | Maps a name indentifying a Stm to the last uses in that Stm.+type LUTabFun = M.Map VName Names++-- | LU-table for the constants, and for each function.+type LUTabProg = (LUTabFun, M.Map Name LUTabFun)++type LastUseOp rep = Op (Aliases rep) -> Names -> LastUseM rep (LUTabFun, Names, Names)+ -- | 'LastUseReader' allows us to abstract over representations by supplying the -- 'onOp' function.-newtype LastUseReader rep = LastUseReader-  { onOp :: Op (Aliases rep) -> Names -> LastUseM rep (LUTabFun, Names, Names)+data LastUseReader rep = LastUseReader+  { onOp :: LastUseOp rep,+    scope :: Scope (Aliases rep)   } -type LastUseM rep a = StateT AliasTab (Reader (LastUseReader rep)) a+-- | Maps a variable or memory block to its aliases.+type AliasTab = M.Map VName Names +newtype LastUseM rep a = LastUseM (StateT AliasTab (Reader (LastUseReader rep)) a)+  deriving+    ( Monad,+      Functor,+      Applicative,+      MonadReader (LastUseReader rep),+      MonadState AliasTab+    )++instance+  (RepTypes rep, CanBeAliased (Op rep)) =>+  HasScope (Aliases rep) (LastUseM rep)+  where+  askScope = asks scope++instance+  (RepTypes rep, CanBeAliased (Op rep)) =>+  LocalScope (Aliases rep) (LastUseM rep)+  where+  localScope sc (LastUseM m) = LastUseM $ do+    local (\rd -> rd {scope = scope rd <> sc}) m++type Constraints rep =+  ( LocalScope (Aliases rep) (LastUseM rep),+    ASTRep rep,+    FreeIn (OpWithAliases (Op rep)),+    HasMemBlock (Aliases rep),+    CanBeAliased (Op rep)+  )++runLastUseM :: LastUseOp rep -> LastUseM rep a -> a+runLastUseM onOp (LastUseM m) =+  runReader (evalStateT m mempty) (LastUseReader onOp mempty)+ aliasLookup :: VName -> LastUseM rep Names aliasLookup vname =   gets $ fromMaybe mempty . M.lookup vname +lastUseProg ::+  Constraints rep =>+  Prog (Aliases rep) ->+  LastUseM rep LUTabProg+lastUseProg prog =+  let bound_in_consts =+        progConsts prog+          & concatMap (patNames . stmPat)+          & namesFromList+      consts = progConsts prog+      funs = progFuns prog+   in inScopeOf consts $ do+        (consts_lu, _) <- lastUseStms consts mempty mempty+        lus <- mapM (lastUseFun bound_in_consts) funs+        pure (consts_lu, M.fromList $ zip (map funDefName funs) lus)++lastUseFun ::+  Constraints rep =>+  Names ->+  FunDef (Aliases rep) ->+  LastUseM rep LUTabFun+lastUseFun bound_in_consts f =+  inScopeOf f $ fst <$> lastUseBody (funDefBody f) (mempty, bound_in_consts)+ -- | Perform last-use analysis on a 'Prog' in 'SeqMem'-lastUsePrg :: Prog (Aliases SeqMem) -> LUTabPrg-lastUsePrg prg = M.fromList $ map lastUseSeqMem $ progFuns prg+lastUseSeqMem :: Prog (Aliases SeqMem) -> LUTabProg+lastUseSeqMem = runLastUseM lastUseSeqOp . lastUseProg  -- | Perform last-use analysis on a 'Prog' in 'GPUMem'-lastUsePrgGPU :: Prog (Aliases GPUMem) -> LUTabPrg-lastUsePrgGPU prg = M.fromList $ map lastUseGPUMem $ progFuns prg---- | Perform last-use analysis on a 'FunDef' in 'SeqMem'-lastUseSeqMem :: FunDef (Aliases SeqMem) -> (Name, LUTabFun)-lastUseSeqMem (FunDef _ _ fname _ _ body) =-  let (res, _) =-        runReader-          (evalStateT (lastUseBody body (mempty, mempty)) mempty)-          (LastUseReader lastUseSeqOp)-   in (fname, res)+lastUseGPUMem :: Prog (Aliases GPUMem) -> LUTabProg+lastUseGPUMem = runLastUseM (lastUseMemOp lastUseGPUOp) . lastUseProg --- | Perform last-use analysis on a 'FunDef' in 'GPUMem'-lastUseGPUMem :: FunDef (Aliases GPUMem) -> (Name, LUTabFun)-lastUseGPUMem (FunDef _ _ fname _ _ body) =-  let (res, _) =-        runReader-          (evalStateT (lastUseBody body (mempty, mempty)) mempty)-          (LastUseReader lastUseGPUOp)-   in (fname, res)+-- | Perform last-use analysis on a 'Prog' in 'MCMem'+lastUseMCMem :: Prog (Aliases MCMem) -> LUTabProg+lastUseMCMem = runLastUseM (lastUseMemOp lastUseMCOp) . lastUseProg  -- | Performing the last-use analysis on a body. --@@ -74,7 +138,7 @@ -- difference between the free-variables in that stmt and the set of variables -- known to be used after that statement. lastUseBody ::-  (ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>+  Constraints rep =>   -- | The body of statements   Body (Aliases rep) ->   -- | The current last-use table, tupled with the known set of already used names@@ -83,18 +147,19 @@   --      (i) an updated last-use table,   --     (ii) an updated set of used names (including the binding).   LastUseM rep (LUTabFun, Names)-lastUseBody bdy@(Body _ stms result) (lutab, used_nms) = do+lastUseBody bdy@(Body _ stms result) (lutab, used_nms) =   -- perform analysis bottom-up in bindings: results are known to be used,   -- hence they are added to the used_nms set.-  (lutab', _) <--    lastUseStms stms (lutab, used_nms) $-      namesToList $-        freeIn $-          map resSubExp result-  -- Clean up the used names by recomputing the aliasing transitive-closure-  -- of the free names in body based on the current alias table @alstab@.-  used_in_body <- aliasTransitiveClosure $ freeIn bdy-  pure (lutab', used_nms <> used_in_body)+  inScopeOf stms $ do+    (lutab', _) <-+      lastUseStms stms (lutab, used_nms) $+        namesToList $+          freeIn $+            map resSubExp result+    -- Clean up the used names by recomputing the aliasing transitive-closure+    -- of the free names in body based on the current alias table @alstab@.+    used_in_body <- aliasTransitiveClosure $ freeIn bdy+    pure (lutab', used_nms <> used_in_body)  -- | Performing the last-use analysis on a body. --@@ -103,7 +168,7 @@ -- difference between the free-variables in that stmt and the set of variables -- known to be used after that statement. lastUseKernelBody ::-  (CanBeAliased (Op rep), ASTRep rep) =>+  Constraints rep =>   -- | The body of statements   KernelBody (Aliases rep) ->   -- | The current last-use table, tupled with the known set of already used names@@ -112,62 +177,71 @@   --      (i) an updated last-use table,   --     (ii) an updated set of used names (including the binding).   LastUseM rep (LUTabFun, Names)-lastUseKernelBody bdy@(KernelBody _ stms result) (lutab, used_nms) = do-  -- perform analysis bottom-up in bindings: results are known to be used,-  -- hence they are added to the used_nms set.-  (lutab', _) <--    lastUseStms stms (lutab, used_nms) $ namesToList $ freeIn result-  -- Clean up the used names by recomputing the aliasing transitive-closure-  -- of the free names in body based on the current alias table @alstab@.-  used_in_body <- aliasTransitiveClosure $ freeIn bdy-  pure (lutab', used_nms <> used_in_body)+lastUseKernelBody bdy@(KernelBody _ stms result) (lutab, used_nms) =+  inScopeOf stms $ do+    -- perform analysis bottom-up in bindings: results are known to be used,+    -- hence they are added to the used_nms set.+    (lutab', _) <-+      lastUseStms stms (lutab, used_nms) $ namesToList $ freeIn result+    -- Clean up the used names by recomputing the aliasing transitive-closure+    -- of the free names in body based on the current alias table @alstab@.+    used_in_body <- aliasTransitiveClosure $ freeIn bdy+    pure (lutab', used_nms <> used_in_body)  lastUseStms ::-  (ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>+  Constraints rep =>   Stms (Aliases rep) ->   (LUTabFun, Names) ->   [VName] ->   LastUseM rep (LUTabFun, Names) lastUseStms Empty (lutab, nms) res_nms = do   aliases <- concatMapM aliasLookup res_nms-  pure (lutab, nms <> aliases)-lastUseStms (stm@(Let pat _ e) :<| stms) (lutab, nms) res_nms = do-  let extra_alias = case e of-        BasicOp (Update _ old _ _) -> oneName old-        BasicOp (FlatUpdate old _ _) -> oneName old-        _ -> mempty-  -- We build up aliases top-down-  updateAliasing extra_alias pat-  -- But compute last use bottom-up-  (lutab', nms') <- lastUseStms stms (lutab, nms) res_nms-  (lutab'', nms'') <- lastUseStm stm (lutab', nms')-  pure (lutab'', nms'')+  pure (lutab, nms <> aliases <> namesFromList res_nms)+lastUseStms (stm@(Let pat _ e) :<| stms) (lutab, nms) res_nms =+  inScopeOf stm $ do+    let extra_alias = case e of+          BasicOp (Update _ old _ _) -> oneName old+          BasicOp (FlatUpdate old _ _) -> oneName old+          _ -> mempty+    -- We build up aliases top-down+    updateAliasing extra_alias pat+    -- But compute last use bottom-up+    (lutab', nms') <- lastUseStms stms (lutab, nms) res_nms+    (lutab'', nms'') <- lastUseStm stm (lutab', nms')+    pure (lutab'', nms'')  lastUseStm ::-  (ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>+  Constraints rep =>   Stm (Aliases rep) ->   (LUTabFun, Names) ->   LastUseM rep (LUTabFun, Names)-lastUseStm (Let pat _ e) (lutab, used_nms) =-  do-    -- analyse the expression and get the-    --  (i)  a new last-use table (in case the @e@ contains bodies of stmts)-    -- (ii) the set of variables lastly used in the current binding.-    -- (iii)  aliased transitive-closure of used names, and-    (lutab', last_uses, used_nms') <- lastUseExp e used_nms-    -- filter-out the binded names from the set of used variables,-    -- since they go out of scope, and update the last-use table.-    let patnms = patNames pat-        used_nms'' = used_nms' `namesSubtract` namesFromList patnms-        lutab'' =-          M.union lutab' $ M.insert (head patnms) last_uses lutab-    pure (lutab'', used_nms'')+lastUseStm (Let pat _ e) (lutab, used_nms) = do+  -- analyse the expression and get the+  --  (i)  a new last-use table (in case the @e@ contains bodies of stmts)+  -- (ii) the set of variables lastly used in the current binding.+  -- (iii)  aliased transitive-closure of used names, and+  (lutab', last_uses, used_nms') <- lastUseExp e used_nms+  sc <- asks scope+  let lu_mems =+        namesToList last_uses+          & mapMaybe (`getScopeMemInfo` sc)+          & map memName+          & namesFromList+          & flip namesSubtract used_nms +  -- filter-out the binded names from the set of used variables,+  -- since they go out of scope, and update the last-use table.+  let patnms = patNames pat+      used_nms'' = used_nms' `namesSubtract` namesFromList patnms+      lutab'' =+        M.union lutab' $ M.insert (head patnms) (last_uses <> lu_mems) lutab+  pure (lutab'', used_nms'')+ --------------------------------  -- | Last-Use Analysis for an expression. lastUseExp ::-  (ASTRep rep, FreeIn (OpWithAliases (Op rep))) =>+  Constraints rep =>   -- | The expression to analyse   Exp (Aliases rep) ->   -- | The set of used names "after" this expression@@ -191,7 +265,7 @@   let used_nms' = used_cases <> body_used_nms   (_, last_used_arrs) <- lastUsedInNames used_nms $ free_in_body <> free_in_cases   pure (lutab_cases <> lutab', last_used_arrs, used_nms')-lastUseExp (DoLoop var_ses _ body) used_nms0 = do+lastUseExp (DoLoop var_ses lf body) used_nms0 = inScopeOf lf $ do   free_in_body <- aliasTransitiveClosure $ freeIn body   -- compute the aliasing transitive closure of initializers that are not last-uses   var_inis <- catMaybes <$> mapM (initHelper (free_in_body <> used_nms0)) var_ses@@ -230,56 +304,93 @@   (used_nms', lu_vars) <- lastUsedInNames used_nms free_in_e   pure (M.empty, lu_vars, used_nms') -lastUseGPUOp :: Op (Aliases GPUMem) -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)-lastUseGPUOp (Alloc se sp) used_nms = do+lastUseMemOp ::+  (inner -> Names -> LastUseM rep (LUTabFun, Names, Names)) ->+  MemOp inner ->+  Names ->+  LastUseM rep (LUTabFun, Names, Names)+lastUseMemOp _ (Alloc se sp) used_nms = do   let free_in_e = freeIn se <> freeIn sp   (used_nms', lu_vars) <- lastUsedInNames used_nms free_in_e   pure (M.empty, lu_vars, used_nms')-lastUseGPUOp (Inner (OtherOp ())) used_nms =-  pure (mempty, mempty, used_nms)-lastUseGPUOp (Inner (SizeOp sop)) used_nms = do-  (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn sop-  pure (mempty, lu_vars, used_nms')-lastUseGPUOp (Inner (SegOp (SegMap _ _ tps kbody))) used_nms = do+lastUseMemOp onInner (Inner op) used_nms = onInner op used_nms++lastUseSegOp ::+  Constraints rep =>+  SegOp lvl (Aliases rep) ->+  Names ->+  LastUseM rep (LUTabFun, Names, Names)+lastUseSegOp (SegMap _ _ tps kbody) used_nms = do   (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn tps   (body_lutab, used_nms'') <- lastUseKernelBody kbody (mempty, used_nms')   pure (body_lutab, lu_vars, used_nms' <> used_nms'')-lastUseGPUOp (Inner (SegOp (SegRed _ _ sbos tps kbody))) used_nms = do+lastUseSegOp (SegRed _ _ sbos tps kbody) used_nms = do   (lutab_sbo, lu_vars_sbo, used_nms_sbo) <- lastUseSegBinOp sbos used_nms   (used_nms', lu_vars) <- lastUsedInNames used_nms_sbo $ freeIn tps   (body_lutab, used_nms'') <- lastUseKernelBody kbody (mempty, used_nms')   pure (M.union lutab_sbo body_lutab, lu_vars <> lu_vars_sbo, used_nms_sbo <> used_nms' <> used_nms'')-lastUseGPUOp (Inner (SegOp (SegScan _ _ sbos tps kbody))) used_nms = do+lastUseSegOp (SegScan _ _ sbos tps kbody) used_nms = do   (lutab_sbo, lu_vars_sbo, used_nms_sbo) <- lastUseSegBinOp sbos used_nms   (used_nms', lu_vars) <- lastUsedInNames used_nms_sbo $ freeIn tps   (body_lutab, used_nms'') <- lastUseKernelBody kbody (mempty, used_nms')   pure (M.union lutab_sbo body_lutab, lu_vars <> lu_vars_sbo, used_nms_sbo <> used_nms' <> used_nms'')-lastUseGPUOp (Inner (SegOp (SegHist _ _ hos tps kbody))) used_nms = do+lastUseSegOp (SegHist _ _ hos tps kbody) used_nms = do   (lutab_sbo, lu_vars_sbo, used_nms_sbo) <- lastUseHistOp hos used_nms   (used_nms', lu_vars) <- lastUsedInNames used_nms_sbo $ freeIn tps   (body_lutab, used_nms'') <- lastUseKernelBody kbody (mempty, used_nms')   pure (M.union lutab_sbo body_lutab, lu_vars <> lu_vars_sbo, used_nms_sbo <> used_nms' <> used_nms'')-lastUseGPUOp (Inner (GPUBody tps body)) used_nms = do++lastUseGPUOp :: HostOp (Aliases GPUMem) () -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)+lastUseGPUOp (GPU.OtherOp ()) used_nms =+  pure (mempty, mempty, used_nms)+lastUseGPUOp (SizeOp sop) used_nms = do+  (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn sop+  pure (mempty, lu_vars, used_nms')+lastUseGPUOp (GPUBody tps body) used_nms = do   (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn tps   (body_lutab, used_nms'') <- lastUseBody body (mempty, used_nms')   pure (body_lutab, lu_vars, used_nms' <> used_nms'')+lastUseGPUOp (SegOp op) used_nms =+  lastUseSegOp op used_nms -lastUseSegBinOp :: [SegBinOp (Aliases GPUMem)] -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)+lastUseMCOp :: MCOp (Aliases MCMem) () -> Names -> LastUseM MCMem (LUTabFun, Names, Names)+lastUseMCOp (MC.OtherOp ()) used_nms =+  pure (mempty, mempty, used_nms)+lastUseMCOp (MC.ParOp par_op op) used_nms = do+  (lutab_par_op, lu_vars_par_op, used_names_par_op) <-+    maybe (pure mempty) (`lastUseSegOp` used_nms) par_op+  (lutab_op, lu_vars_op, used_names_op) <-+    lastUseSegOp op used_nms+  pure+    ( lutab_par_op <> lutab_op,+      lu_vars_par_op <> lu_vars_op,+      used_names_par_op <> used_names_op+    )++lastUseSegBinOp ::+  Constraints rep =>+  [SegBinOp (Aliases rep)] ->+  Names ->+  LastUseM rep (LUTabFun, Names, Names) lastUseSegBinOp sbos used_nms = do   (lutab, lu_vars, used_nms') <- unzip3 <$> mapM helper sbos   pure (mconcat lutab, mconcat lu_vars, mconcat used_nms')   where-    helper (SegBinOp _ (Lambda _ body _) neutral shp) = do+    helper (SegBinOp _ l@(Lambda _ body _) neutral shp) = inScopeOf l $ do       (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn neutral <> freeIn shp       (body_lutab, used_nms'') <- lastUseBody body (mempty, used_nms')       pure (body_lutab, lu_vars, used_nms'') -lastUseHistOp :: [HistOp (Aliases GPUMem)] -> Names -> LastUseM GPUMem (LUTabFun, Names, Names)+lastUseHistOp ::+  Constraints rep =>+  [HistOp (Aliases rep)] ->+  Names ->+  LastUseM rep (LUTabFun, Names, Names) lastUseHistOp hos used_nms = do   (lutab, lu_vars, used_nms') <- unzip3 <$> mapM helper hos   pure (mconcat lutab, mconcat lu_vars, mconcat used_nms')   where-    helper (HistOp shp rf dest neutral shp' (Lambda _ body _)) = do+    helper (HistOp shp rf dest neutral shp' l@(Lambda _ body _)) = inScopeOf l $ do       (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn shp <> freeIn rf <> freeIn dest <> freeIn neutral <> freeIn shp'       (body_lutab, used_nms'') <- lastUseBody body (mempty, used_nms')       pure (body_lutab, lu_vars, used_nms'')

@@ -228,7 +228,7 @@             if m_b `nameIn` original_mem_aliases               then (wrt_lmads' <> lmads'', Set mempty)               else (wrt_lmads', lmads'')-          no_overlap = noMemOverlap td_env prev_use wrt_lmads''+          no_overlap = noMemOverlap td_env (lmads <> prev_use) wrt_lmads''           wrt_lmads =             if no_overlap               then Just wrt_lmads''

@@ -22,7 +22,8 @@ import Data.Maybe import Futhark.Analysis.PrimExp.Convert import Futhark.IR.Aliases-import Futhark.IR.GPUMem+import Futhark.IR.GPUMem as GPU+import Futhark.IR.MCMem as MC import Futhark.IR.Mem.IxFun qualified as IxFun import Futhark.Optimise.ArrayShortCircuiting.DataStructs @@ -131,19 +132,40 @@    scopeHelper :: inner -> Scope rep +instance TopDownHelper (SegOp lvl rep) where+  innerNonNegatives _ op =+    foldMap (oneName . fst) $ unSegSpace $ segSpace op++  innerKnownLessThan op =+    map (fmap $ primExpFromSubExp $ IntType Int64) $ unSegSpace $ segSpace op++  scopeHelper op = scopeOfSegSpace $ segSpace op+ instance TopDownHelper (HostOp (Aliases GPUMem) ()) where-  innerNonNegatives _ (SegOp seg_op) =-    foldMap (oneName . fst) $ unSegSpace $ segSpace seg_op+  innerNonNegatives vs (SegOp op) = innerNonNegatives vs op   innerNonNegatives [vname] (SizeOp (GetSize _ _)) = oneName vname   innerNonNegatives [vname] (SizeOp (GetSizeMax _)) = oneName vname   innerNonNegatives _ _ = mempty -  innerKnownLessThan (SegOp seg_op) =-    map (fmap $ primExpFromSubExp $ IntType Int64) $ unSegSpace $ segSpace seg_op+  innerKnownLessThan (SegOp op) = innerKnownLessThan op   innerKnownLessThan _ = mempty -  scopeHelper (SegOp seg_op) = scopeOfSegSpace $ segSpace seg_op+  scopeHelper (SegOp op) = scopeHelper op   scopeHelper _ = mempty++instance TopDownHelper inner => TopDownHelper (MC.MCOp (Aliases MCMem) inner) where+  innerNonNegatives vs (ParOp par_op op) =+    maybe mempty (innerNonNegatives vs) par_op+      <> innerNonNegatives vs op+  innerNonNegatives vs (MC.OtherOp op) =+    innerNonNegatives vs op+  innerKnownLessThan (ParOp par_op op) =+    maybe mempty innerKnownLessThan par_op <> innerKnownLessThan op+  innerKnownLessThan (MC.OtherOp op) =+    innerKnownLessThan op+  scopeHelper (ParOp par_op op) =+    maybe mempty scopeHelper par_op <> scopeHelper op+  scopeHelper MC.OtherOp {} = mempty  instance TopDownHelper () where   innerNonNegatives _ () = mempty

@@ -6,11 +6,17 @@ -- | This pass attempts to lift allocations as far towards the top in their body -- as possible. It does not try to hoist allocations outside across body -- boundaries.-module Futhark.Pass.LiftAllocations (liftAllocationsSeqMem, liftAllocationsGPUMem) where+module Futhark.Pass.LiftAllocations+  ( liftAllocationsSeqMem,+    liftAllocationsGPUMem,+    liftAllocationsMCMem,+  )+where  import Control.Monad.Reader import Data.Sequence (Seq (..)) import Futhark.IR.GPUMem+import Futhark.IR.MCMem import Futhark.IR.SeqMem import Futhark.Pass (Pass (..)) @@ -40,6 +46,19 @@               progFuns         } +liftAllocationsMCMem :: Pass MCMem MCMem+liftAllocationsMCMem =+  Pass "lift allocations mc" "lift allocations mc" $ \prog@Prog {progFuns} ->+    pure $+      prog+        { progFuns =+            fmap+              ( \f@FunDef {funDefBody} ->+                  f {funDefBody = runReader (liftAllocationsInBody funDefBody) (Env liftAllocationsInMCOp)}+              )+              progFuns+        }+ newtype Env inner = Env   {onInner :: inner -> LiftM inner inner} @@ -104,17 +123,28 @@     then liftAllocationsInStms stms (stm :<| lifted) acc ((to_lift `namesSubtract` pat_names) <> freeIn stm)     else liftAllocationsInStms stms lifted (stm :<| acc) to_lift -liftAllocationsInHostOp :: HostOp GPUMem () -> LiftM (HostOp GPUMem ()) (HostOp GPUMem ())-liftAllocationsInHostOp (SegOp (SegMap lvl sp tps body)) = do+liftAllocationsInSegOp ::+  (Mem rep inner, LetDec rep ~ LetDecMem) =>+  SegOp lvl rep ->+  LiftM inner (SegOp lvl rep)+liftAllocationsInSegOp (SegMap lvl sp tps body) = do   stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty-  pure $ SegOp $ SegMap lvl sp tps $ body {kernelBodyStms = stms}-liftAllocationsInHostOp (SegOp (SegRed lvl sp binops tps body)) = do+  pure $ SegMap lvl sp tps $ body {kernelBodyStms = stms}+liftAllocationsInSegOp (SegRed lvl sp binops tps body) = do   stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty-  pure $ SegOp $ SegRed lvl sp binops tps $ body {kernelBodyStms = stms}-liftAllocationsInHostOp (SegOp (SegScan lvl sp binops tps body)) = do+  pure $ SegRed lvl sp binops tps $ body {kernelBodyStms = stms}+liftAllocationsInSegOp (SegScan lvl sp binops tps body) = do   stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty-  pure $ SegOp $ SegScan lvl sp binops tps $ body {kernelBodyStms = stms}-liftAllocationsInHostOp (SegOp (SegHist lvl sp histops tps body)) = do+  pure $ SegScan lvl sp binops tps $ body {kernelBodyStms = stms}+liftAllocationsInSegOp (SegHist lvl sp histops tps body) = do   stms <- liftAllocationsInStms (kernelBodyStms body) mempty mempty mempty-  pure $ SegOp $ SegHist lvl sp histops tps $ body {kernelBodyStms = stms}+  pure $ SegHist lvl sp histops tps $ body {kernelBodyStms = stms}++liftAllocationsInHostOp :: HostOp GPUMem () -> LiftM (HostOp GPUMem ()) (HostOp GPUMem ())+liftAllocationsInHostOp (SegOp op) = SegOp <$> liftAllocationsInSegOp op liftAllocationsInHostOp op = pure op++liftAllocationsInMCOp :: MCOp MCMem () -> LiftM (MCOp MCMem ()) (MCOp MCMem ())+liftAllocationsInMCOp (ParOp par op) =+  ParOp <$> traverse liftAllocationsInSegOp par <*> liftAllocationsInSegOp op+liftAllocationsInMCOp op = pure op

@@ -5,7 +5,12 @@  -- | This pass attempts to lower allocations as far towards the bottom of their -- body as possible.-module Futhark.Pass.LowerAllocations (lowerAllocationsSeqMem, lowerAllocationsGPUMem) where+module Futhark.Pass.LowerAllocations+  ( lowerAllocationsSeqMem,+    lowerAllocationsGPUMem,+    lowerAllocationsMCMem,+  )+where  import Control.Monad.Reader import Data.Function ((&))@@ -13,6 +18,7 @@ import Data.Sequence (Seq (..)) import Data.Sequence qualified as Seq import Futhark.IR.GPUMem+import Futhark.IR.MCMem import Futhark.IR.SeqMem import Futhark.Pass (Pass (..)) @@ -42,6 +48,19 @@               progFuns         } +lowerAllocationsMCMem :: Pass MCMem MCMem+lowerAllocationsMCMem =+  Pass "lower allocations mc" "lower allocations mc" $ \prog@Prog {progFuns} ->+    pure $+      prog+        { progFuns =+            fmap+              ( \f@FunDef {funDefBody} ->+                  f {funDefBody = runReader (lowerAllocationsInBody funDefBody) (Env lowerAllocationsInMCOp)}+              )+              progFuns+        }+ newtype Env inner = Env   {onInner :: inner -> LowerM inner inner} @@ -98,17 +117,28 @@       )       (alloc, acc) -lowerAllocationsInHostOp :: HostOp GPUMem () -> LowerM (HostOp GPUMem ()) (HostOp GPUMem ())-lowerAllocationsInHostOp (SegOp (SegMap lvl sp tps body)) = do+lowerAllocationsInSegOp ::+  (Mem rep inner, LetDec rep ~ LetDecMem) =>+  SegOp lvl rep ->+  LowerM inner (SegOp lvl rep)+lowerAllocationsInSegOp (SegMap lvl sp tps body) = do   stms <- lowerAllocationsInStms (kernelBodyStms body) mempty mempty-  pure $ SegOp $ SegMap lvl sp tps $ body {kernelBodyStms = stms}-lowerAllocationsInHostOp (SegOp (SegRed lvl sp binops tps body)) = do+  pure $ SegMap lvl sp tps $ body {kernelBodyStms = stms}+lowerAllocationsInSegOp (SegRed lvl sp binops tps body) = do   stms <- lowerAllocationsInStms (kernelBodyStms body) mempty mempty-  pure $ SegOp $ SegRed lvl sp binops tps $ body {kernelBodyStms = stms}-lowerAllocationsInHostOp (SegOp (SegScan lvl sp binops tps body)) = do+  pure $ SegRed lvl sp binops tps $ body {kernelBodyStms = stms}+lowerAllocationsInSegOp (SegScan lvl sp binops tps body) = do   stms <- lowerAllocationsInStms (kernelBodyStms body) mempty mempty-  pure $ SegOp $ SegScan lvl sp binops tps $ body {kernelBodyStms = stms}-lowerAllocationsInHostOp (SegOp (SegHist lvl sp histops tps body)) = do+  pure $ SegScan lvl sp binops tps $ body {kernelBodyStms = stms}+lowerAllocationsInSegOp (SegHist lvl sp histops tps body) = do   stms <- lowerAllocationsInStms (kernelBodyStms body) mempty mempty-  pure $ SegOp $ SegHist lvl sp histops tps $ body {kernelBodyStms = stms}+  pure $ SegHist lvl sp histops tps $ body {kernelBodyStms = stms}++lowerAllocationsInHostOp :: HostOp GPUMem () -> LowerM (HostOp GPUMem ()) (HostOp GPUMem ())+lowerAllocationsInHostOp (SegOp op) = SegOp <$> lowerAllocationsInSegOp op lowerAllocationsInHostOp op = pure op++lowerAllocationsInMCOp :: MCOp MCMem () -> LowerM (MCOp MCMem ()) (MCOp MCMem ())+lowerAllocationsInMCOp (ParOp par op) =+  ParOp <$> traverse lowerAllocationsInSegOp par <*> lowerAllocationsInSegOp op+lowerAllocationsInMCOp op = pure op

@@ -195,5 +195,15 @@         simplifyMCMem,         entryPointMemMC,         doubleBufferMC,+        simplifyMCMem,+        performCSE False,+        LiftAllocations.liftAllocationsMCMem,+        simplifyMCMem,+        ArrayShortCircuiting.optimiseMCMem,+        simplifyMCMem,+        performCSE False,+        simplifyMCMem,+        LowerAllocations.lowerAllocationsMCMem,+        performCSE False,         simplifyMCMem       ]

@@ -86,9 +86,11 @@ -- | How a program can be transformed. data StructurePipeline   = GpuPipeline+  | MCPipeline   | SOACSPipeline   | SeqMemPipeline   | GpuMemPipeline+  | MCMemPipeline   | NoPipeline   deriving (Show) @@ -342,6 +344,8 @@   choice     [ lexeme sep "gpu-mem" $> GpuMemPipeline,       lexeme sep "gpu" $> GpuPipeline,+      lexeme sep "mc-mem" $> MCMemPipeline,+      lexeme sep "mc" $> MCPipeline,       lexeme sep "seq-mem" $> SeqMemPipeline,       lexeme sep "internalised" $> NoPipeline,       pure SOACSPipeline

@@ -155,10 +155,10 @@ prettyValue :: Value m -> Doc a prettyValue = prettyValueWith pprPrim   where-    pprPrim (UnsignedValue (Int8Value v)) = pretty v-    pprPrim (UnsignedValue (Int16Value v)) = pretty v-    pprPrim (UnsignedValue (Int32Value v)) = pretty v-    pprPrim (UnsignedValue (Int64Value v)) = pretty v+    pprPrim (UnsignedValue (Int8Value v)) = pretty (fromIntegral v :: Word8)+    pprPrim (UnsignedValue (Int16Value v)) = pretty (fromIntegral v :: Word16)+    pprPrim (UnsignedValue (Int32Value v)) = pretty (fromIntegral v :: Word32)+    pprPrim (UnsignedValue (Int64Value v)) = pretty (fromIntegral v :: Word64)     pprPrim (SignedValue (Int8Value v)) = pretty v     pprPrim (SignedValue (Int16Value v)) = pretty v     pprPrim (SignedValue (Int32Value v)) = pretty v

@@ -31,6 +31,7 @@     PrimValue (..),     primValueType,     blankPrimValue,+    onePrimValue,      -- * Operations     Overflow (..),@@ -357,6 +358,19 @@ blankPrimValue (FloatType Float64) = FloatValue $ Float64Value 0.0 blankPrimValue Bool = BoolValue False blankPrimValue Unit = UnitValue++-- | A one value of the given primitive type - this is one+-- whatever is close to it.+onePrimValue :: PrimType -> PrimValue+onePrimValue (IntType Int8) = IntValue $ Int8Value 1+onePrimValue (IntType Int16) = IntValue $ Int16Value 1+onePrimValue (IntType Int32) = IntValue $ Int32Value 1+onePrimValue (IntType Int64) = IntValue $ Int64Value 1+onePrimValue (FloatType Float16) = FloatValue $ Float16Value 1.0+onePrimValue (FloatType Float32) = FloatValue $ Float32Value 1.0+onePrimValue (FloatType Float64) = FloatValue $ Float64Value 1.0+onePrimValue Bool = BoolValue True+onePrimValue Unit = UnitValue  -- | Various unary operators.  It is a bit ad-hoc what is a unary -- operator and what is a built-in function.  Perhaps these should all

@@ -25,7 +25,7 @@  -- | Is this character a valid member of an identifier? constituent :: Char -> Bool-constituent c = isAlphaNum c || (c `elem` ("_/'+-=!&^.<>*|" :: String))+constituent c = isAlphaNum c || (c `elem` ("_/'+-=!&^.<>*|%" :: String))  -- | Consume whitespace (including skipping line comments). whitespace :: Parsec Void T.Text ()

@@ -21,6 +21,7 @@     prettyStacktrace,     progHoles,     defaultEntryPoint,+    paramName,      -- * Queries on expressions     typeOf,@@ -347,6 +348,11 @@     (combineTypeShapes (Scalar et1) (Scalar et2) `setAliases` mempty) combineTypeShapes _ new_tp = new_tp +-- | The name, if any.+paramName :: PName -> Maybe VName+paramName (Named v) = Just v+paramName Unnamed = Nothing+ -- | Match the dimensions of otherwise assumed-equal types.  The -- combining function is also passed the names bound within the type -- (from named parameters or return types).@@ -380,7 +386,7 @@         ( Scalar (Arrow als1 p1 a1 (RetType dims1 b1)),           Scalar (Arrow als2 p2 a2 (RetType dims2 b2))           ) ->-            let bound' = mapMaybe maybePName [p1, p2] <> dims1 <> dims2 <> bound+            let bound' = mapMaybe paramName [p1, p2] <> dims1 <> dims2 <> bound              in Scalar                   <$> ( Arrow (als1 <> als2) p1                           <$> matchDims' bound' a1 a2@@ -397,9 +403,6 @@     matchTypeArg bound (TypeArgDim x loc) (TypeArgDim y _) =       TypeArgDim <$> onDims bound x y <*> pure loc     matchTypeArg _ a _ = pure a--    maybePName (Named v) = Just v-    maybePName Unnamed = Nothing      onShapes bound shape1 shape2 =       Shape <$> zipWithM (onDims bound) (shapeDims shape1) (shapeDims shape2)

@@ -642,17 +642,26 @@   case entry' of     Just _       | any isTypeParam tparams' ->-          typeError loc mempty "Entry point functions may not be polymorphic."+          typeError loc mempty $+            withIndexLink+              "polymorphic-entry"+              "Entry point functions may not be polymorphic."       | not (all patternOrderZero params')           || not (all orderZero rettype_params)           || not (orderZero rettype') ->-          typeError loc mempty "Entry point functions may not be higher-order."+          typeError loc mempty $+            withIndexLink+              "higher-order-entry"+              "Entry point functions may not be higher-order."       | sizes_only_in_ret <-           S.fromList (map typeParamName tparams')             `S.intersection` freeInType rettype'             `S.difference` foldMap freeInType (map patternStructType params' ++ rettype_params),         not $ S.null sizes_only_in_ret ->-          typeError loc mempty "Entry point functions must not be size-polymorphic in their return type."+          typeError loc mempty $+            withIndexLink+              "size-polymorphic-entry"+              "Entry point functions must not be size-polymorphic in their return type."       | p : _ <- filter nastyParameter params' ->           warn loc $             "Entry point parameter\n"

@@ -37,7 +37,7 @@ overloadedTypeVars :: Constraints -> Names overloadedTypeVars = mconcat . map f . M.elems   where-    f (_, HasFields fs _) = mconcat $ map typeVars $ M.elems fs+    f (_, HasFields _ fs _) = mconcat $ map typeVars $ M.elems fs     f _ = mempty  --- Basic checking@@ -375,7 +375,7 @@     zeroOrderType       (mkUsage loc "returning value of this type from 'if' expression")       "type returned from branch"-      t'+      (toStruct t')      pure $ AppExp (If e1' e2' e3' loc) (Info $ AppRes t' retext)   where@@ -445,8 +445,10 @@     t <- expType e'     case anyConsumption e_occs of       Just c ->-        let msg = "type computed with consumption at " <> locText (location c)-         in zeroOrderType (mkUsage loc "consumption in right-hand side of 'let'-binding") msg t+        zeroOrderType+          (mkUsage loc "consumption in right-hand side of 'let'-binding")+          ("type computed with consumption at " <> locText (location c))+          (toStruct t)       _ -> pure ()      incLevel . bindingSizes sizes $ \sizes' ->@@ -499,8 +501,7 @@             )             (Info $ AppRes body_t ext) checkExp (AppExp (LetWith dest src slice ve body loc) _) =-  sequentially (checkIdent src) $ \src' _ -> do-    slice' <- checkSlice slice+  sequentially ((,) <$> checkIdent src <*> checkSlice slice) $ \(src', slice') _ -> do     (t, _) <- newArrayType (srclocOf src) "src" $ sliceDims slice'     unify (mkUsage loc "type of target array") t $ toStruct $ unInfo $ identType src' @@ -719,7 +720,7 @@     zeroOrderType       (mkUsage loc "being returned 'match'")       "type returned from pattern match"-      t+      (toStruct t)     pure $ AppExp (Match e' cs' loc) (Info $ AppRes t retext) checkExp (Attr info e loc) =   Attr <$> checkAttr info <*> checkExp e <*> pure loc@@ -1258,14 +1259,14 @@       typeError usage mempty . withIndexLink "ambiguous-type" $         "Type is ambiguous (must be equality type)."           </> "Add a type annotation to disambiguate the type."-    fixOverloaded (_, HasFields fs usage) =+    fixOverloaded (_, HasFields _ fs usage) =       typeError usage mempty . withIndexLink "ambiguous-type" $         "Type is ambiguous.  Must be record with fields:"           </> indent 2 (stack $ map field $ M.toList fs)           </> "Add a type annotation to disambiguate the type."       where         field (l, t) = pretty l <> colon <+> align (pretty t)-    fixOverloaded (_, HasConstrs cs usage) =+    fixOverloaded (_, HasConstrs _ cs usage) =       typeError usage mempty . withIndexLink "ambiguous-type" $         "Type is ambiguous (must be a sum type with constructors:"           <+> pretty (Sum cs) <> ")."

@@ -206,6 +206,7 @@     zeroOrderType       (mkUsage (srclocOf mergeexp) "use as loop variable")       "type used as loop variable"+      . toStruct       =<< expTypeFully mergeexp'      -- The handling of dimension sizes is a bit intricate, but very

@@ -235,7 +235,6 @@   (Pat -> TermTypeM a) ->   TermTypeM a bindingPat sizes p t m = do-  checkForDuplicateNames (map sizeBinderToParam sizes) [p]   checkPat sizes p t $ \p' -> binding True (S.toList $ patIdents p') $ do     -- Perform an observation of every declared dimension.  This     -- prevents unused-name warnings for otherwise unused dimensions.@@ -354,6 +353,13 @@   pure $ PatLit l (Info (fromStruct t')) loc checkPat' sizes (PatConstr n NoInfo ps loc) (Ascribed (Scalar (Sum cs)))   | Just ts <- M.lookup n cs = do+      when (length ps /= length ts) $+        typeError loc mempty $+          "Pattern #" <> pretty n <> " expects"+            <+> pretty (length ps)+            <+> "constructor arguments, but type provides"+            <+> pretty (length ts)+            <+> "arguments."       ps' <- zipWithM (checkPat' sizes) ps $ map Ascribed ts       pure $ PatConstr n (Info (Scalar (Sum cs))) ps' loc checkPat' sizes (PatConstr n NoInfo ps loc) (Ascribed t) = do

@@ -122,9 +122,9 @@   | ParamType Liftedness SrcLoc   | Constraint StructRetType Usage   | Overloaded [PrimType] Usage-  | HasFields (M.Map Name StructType) Usage+  | HasFields Liftedness (M.Map Name StructType) Usage   | Equality Usage-  | HasConstrs (M.Map Name [StructType]) Usage+  | HasConstrs Liftedness (M.Map Name [StructType]) Usage   | ParamSize SrcLoc   | -- | Is not actually a type, but a term-level size,     -- possibly already set to something specific.@@ -141,9 +141,9 @@   locOf (ParamType _ usage) = locOf usage   locOf (Constraint _ usage) = locOf usage   locOf (Overloaded _ usage) = locOf usage-  locOf (HasFields _ usage) = locOf usage+  locOf (HasFields _ _ usage) = locOf usage   locOf (Equality usage) = locOf usage-  locOf (HasConstrs _ usage) = locOf usage+  locOf (HasConstrs _ _ usage) = locOf usage   locOf (ParamSize loc) = locOf loc   locOf (Size _ usage) = locOf usage   locOf (UnknowableSize loc _) = locOf loc@@ -270,7 +270,7 @@ typeVarNotes :: MonadUnify m => VName -> m Notes typeVarNotes v = maybe mempty (note . snd) . M.lookup v <$> getConstraints   where-    note (HasConstrs cs _) =+    note (HasConstrs _ cs _) =       aNote $         prettyName v           <+> "="@@ -278,7 +278,7 @@           <+> "..."     note (Overloaded ts _) =       aNote $ prettyName v <+> "must be one of" <+> mconcat (punctuate ", " (map pretty ts))-    note (HasFields fs _) =+    note (HasFields _ fs _) =       aNote $         prettyName v           <+> "="@@ -678,21 +678,21 @@                 <> commasep (map (dquotes . prettyName) problems)                 <> " used as size(s) would go out of scope." +  let unliftedBcs unlifted_usage =+        breadCrumb+          ( Matching $+              "When verifying that"+                <+> dquotes (prettyName vn)+                <+> textwrap "is not instantiated with a function type, due to"+                <+> pretty unlifted_usage+          )+          bcs+   case snd <$> M.lookup vn constraints of     Just (NoConstraint Unlifted unlift_usage) -> do-      let bcs' =-            breadCrumb-              ( Matching $-                  "When verifying that"-                    <+> dquotes (prettyName vn)-                    <+> textwrap "is not instantiated with a function type, due to"-                    <+> pretty unlift_usage-              )-              bcs-       link -      arrayElemTypeWith usage bcs' tp+      arrayElemTypeWith usage (unliftedBcs unlift_usage) tp       when (any (`elem` bound) (freeInType tp)) $         unifyError usage mempty bcs $           "Type variable"@@ -722,7 +722,8 @@                   <+> commasep (map pretty ts)                   </> "due to"                   <+> pretty old_usage <> "."-    Just (HasFields required_fields old_usage) -> do+    Just (HasFields l required_fields old_usage) -> do+      when (l == Unlifted) $ arrayElemTypeWith usage (unliftedBcs old_usage) tp       case tp of         Scalar (Record tp_fields)           | all (`M.member` tp_fields) $ M.keys required_fields -> do@@ -734,7 +735,7 @@               unifySharedFields onDims usage bound bcs required_fields' tp_fields         Scalar (TypeVar _ _ (QualName [] v) []) -> do           case M.lookup v constraints of-            Just (_, HasFields tp_fields _) ->+            Just (_, HasFields _ tp_fields _) ->               unifySharedFields onDims usage bound bcs required_fields tp_fields             Just (_, NoConstraint {}) -> pure ()             Just (_, Equality {}) -> pure ()@@ -746,10 +747,10 @@             M.insertWith               combineFields               v-              (lvl, HasFields required_fields old_usage)+              (lvl, HasFields l required_fields old_usage)           where-            combineFields (_, HasFields fs1 usage1) (_, HasFields fs2 _) =-              (lvl, HasFields (M.union fs1 fs2) usage1)+            combineFields (_, HasFields l1 fs1 usage1) (_, HasFields l2 fs2 _) =+              (lvl, HasFields (l1 `min` l2) (M.union fs1 fs2) usage1)             combineFields hasfs _ = hasfs         _ ->           unifyError usage mempty bcs $@@ -764,7 +765,8 @@               </> "due to"               <+> pretty old_usage <> "."     -- See Note [Linking variables to sum types]-    Just (HasConstrs required_cs old_usage) ->+    Just (HasConstrs l required_cs old_usage) -> do+      when (l == Unlifted) $ arrayElemTypeWith usage (unliftedBcs old_usage) tp       case tp of         Scalar (Sum ts)           | all (`M.member` ts) $ M.keys required_cs -> do@@ -775,7 +777,7 @@               unifySharedConstructors onDims usage bound bcs required_cs ts         Scalar (TypeVar _ _ (QualName [] v) []) -> do           case M.lookup v constraints of-            Just (_, HasConstrs v_cs _) ->+            Just (_, HasConstrs _ v_cs _) ->               unifySharedConstructors onDims usage bound bcs required_cs v_cs             Just (_, NoConstraint {}) -> pure ()             Just (_, Equality {}) -> pure ()@@ -787,10 +789,10 @@             M.insertWith               combineConstrs               v-              (lvl, HasConstrs required_cs old_usage)+              (lvl, HasConstrs l required_cs old_usage)           where-            combineConstrs (_, HasConstrs cs1 usage1) (_, HasConstrs cs2 _) =-              (lvl, HasConstrs (M.union cs1 cs2) usage1)+            combineConstrs (_, HasConstrs l1 cs1 usage1) (_, HasConstrs l2 cs2 _) =+              (lvl, HasConstrs (l1 `min` l2) (M.union cs1 cs2) usage1)             combineConstrs hasCs _ = hasCs         _ -> noSumType mempty     _ -> link@@ -875,14 +877,14 @@               <+> "due to"               <+> pretty vn_usage <> "."         ts' -> modifyConstraints $ M.insert vn (lvl, Overloaded ts' usage)-    Just (_, HasConstrs _ vn_usage) ->+    Just (_, HasConstrs _ _ vn_usage) ->       unifyError usage mempty noBreadCrumbs $         "Type constrained to one of"           <+> commasep (map pretty ts)             <> ", but also inferred to be sum type due to"           <+> pretty vn_usage             <> "."-    Just (_, HasFields _ vn_usage) ->+    Just (_, HasFields _ _ vn_usage) ->       unifyError usage mempty noBreadCrumbs $         "Type constrained to one of"           <+> commasep (map pretty ts)@@ -927,29 +929,33 @@           pure () -- All primtypes support equality.         Just (_, Equality {}) ->           pure ()-        Just (_, HasConstrs cs _) ->+        Just (_, HasConstrs _ cs _) ->           mapM_ (equalityType usage) $ concat $ M.elems cs         _ ->           unifyError usage mempty noBreadCrumbs $             "Type" <+> prettyName vn <+> "does not support equality."  zeroOrderTypeWith ::-  (MonadUnify m, Pretty (Shape dim), Monoid as) =>+  MonadUnify m =>   Usage ->   BreadCrumbs ->-  TypeBase dim as ->+  StructType ->   m () zeroOrderTypeWith usage bcs t = do   unless (orderZero t) $     unifyError usage mempty bcs $       "Type" </> indent 2 (pretty t) </> "found to be functional."-  mapM_ mustBeZeroOrder . S.toList . typeVars $ t+  mapM_ mustBeZeroOrder . S.toList . typeVars =<< normType t   where     mustBeZeroOrder vn = do       constraints <- getConstraints       case M.lookup vn constraints of         Just (lvl, NoConstraint _ _) ->           modifyConstraints $ M.insert vn (lvl, NoConstraint Unlifted usage)+        Just (lvl, HasFields _ fs _) ->+          modifyConstraints $ M.insert vn (lvl, HasFields Unlifted fs usage)+        Just (lvl, HasConstrs _ cs _) ->+          modifyConstraints $ M.insert vn (lvl, HasConstrs Unlifted cs usage)         Just (_, ParamType Lifted ploc) ->           unifyError usage mempty bcs $             "Type parameter"@@ -961,11 +967,7 @@  -- | Assert that this type must be zero-order. zeroOrderType ::-  (MonadUnify m, Pretty (Shape dim), Monoid as) =>-  Usage ->-  T.Text ->-  TypeBase dim as ->-  m ()+  MonadUnify m => Usage -> T.Text -> StructType -> m () zeroOrderType usage desc =   zeroOrderTypeWith usage $ breadCrumb bc noBreadCrumbs   where@@ -1057,12 +1059,14 @@   constraints <- getConstraints   case t of     Scalar (TypeVar _ _ (QualName _ tn) [])-      | Just (lvl, NoConstraint {}) <- M.lookup tn constraints -> do+      | Just (lvl, NoConstraint l _) <- M.lookup tn constraints -> do           mapM_ (scopeCheck usage noBreadCrumbs tn lvl) fs-          modifyConstraints $ M.insert tn (lvl, HasConstrs (M.singleton c fs) usage)-      | Just (lvl, HasConstrs cs _) <- M.lookup tn constraints ->+          modifyConstraints $ M.insert tn (lvl, HasConstrs l (M.singleton c fs) usage)+      | Just (lvl, HasConstrs l cs _) <- M.lookup tn constraints ->           case M.lookup c cs of-            Nothing -> modifyConstraints $ M.insert tn (lvl, HasConstrs (M.insert c fs cs) usage)+            Nothing ->+              modifyConstraints $+                M.insert tn (lvl, HasConstrs l (M.insert c fs cs) usage)             Just fs'               | length fs == length fs' -> zipWithM_ (unify usage) fs fs'               | otherwise ->@@ -1098,16 +1102,16 @@     Scalar (TypeVar _ _ (QualName _ tn) [])       | Just (lvl, NoConstraint {}) <- M.lookup tn constraints -> do           scopeCheck usage bcs tn lvl l_type-          modifyConstraints $ M.insert tn (lvl, HasFields (M.singleton l l_type) usage)+          modifyConstraints $ M.insert tn (lvl, HasFields Lifted (M.singleton l l_type) usage)           pure l_type'-      | Just (lvl, HasFields fields _) <- M.lookup tn constraints -> do+      | Just (lvl, HasFields lifted fields _) <- M.lookup tn constraints -> do           case M.lookup l fields of             Just t' -> unifyWith onDims usage bound bcs l_type t'             Nothing ->               modifyConstraints $                 M.insert                   tn-                  (lvl, HasFields (M.insert l l_type fields) usage)+                  (lvl, HasFields lifted (M.insert l l_type fields) usage)           pure l_type'     Scalar (Record fields)       | Just t' <- M.lookup l fields -> do