futhark 0.25.24 → 0.25.25

raw patch · 30 files changed

+356/−366 lines, 30 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• CHANGELOG.md +22/−0
• docs/error-index.rst +0/−63
• docs/language-reference.rst +10/−5
• futhark.cabal +1/−1
• src/Futhark/AD/Fwd.hs +3/−3
• src/Futhark/AD/Rev/SOAC.hs +1/−1
• src/Futhark/Analysis/HORep/MapNest.hs +1/−1
• src/Futhark/Analysis/HORep/SOAC.hs +3/−3
• src/Futhark/Fmt/Printer.hs +26/−11
• src/Futhark/IR/Parse.hs +15/−14
• src/Futhark/IR/SOACS.hs +1/−1
• src/Futhark/IR/SOACS/SOAC.hs +71/−84
• src/Futhark/IR/SOACS/Simplify.hs +17/−19
• src/Futhark/Internalise/Exps.hs +6/−6
• src/Futhark/Internalise/LiftLambdas.hs +13/−4
• src/Futhark/Internalise/Monomorphise.hs +2/−14
• src/Futhark/Optimise/Fusion.hs +5/−5
• src/Futhark/Optimise/Fusion/RulesWithAccs.hs +2/−2
• src/Futhark/Optimise/Fusion/TryFusion.hs +12/−5
• src/Futhark/Optimise/GenRedOpt.hs +1/−1
• src/Futhark/Pass/AD.hs +4/−4
• src/Futhark/Pass/ExtractKernels.hs +3/−3
• src/Futhark/Tools.hs +3/−5
• src/Futhark/Transform/FirstOrderTransform.hs +1/−1
• src/Language/Futhark/Interpreter.hs +7/−7
• src/Language/Futhark/Interpreter/AD.hs +6/−2
• src/Language/Futhark/Interpreter/Values.hs +6/−2
• src/Language/Futhark/Prop.hs +24/−0
• src/Language/Futhark/TypeChecker/Terms.hs +1/−88
• src/Language/Futhark/TypeChecker/Unify.hs +89/−11

@@ -5,6 +5,28 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html). +## [0.25.25]++### Added++* Improvements to `futhark fmt`.++### Fixed++* Sizes that go out of scope due to use of higher order functions will+  now work in more cases by adding existentials. (#2193)++* Tracing inside AD operators with the interpreter now prints values+  properly.++* Compiled and interpreted code now have same treatment of inclusive+  ranges with start==end and negative step size, e.g. `1..0...1`+  produces `[1]` rather than an invalid range error.++* Inconsistent handling of types in lambda lifting (#2197).++* Invalid primal results from `vjp2` in interpreter (#2199).+ ## [0.25.24]  ### Added

@@ -566,69 +566,6 @@ which is erroneous, but with workarounds, as explained in :ref:`unused-existential`. -.. _unify-param-existential:--"Parameter *x* used as size would go out of scope."-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~--This error tends to happen when higher-order functions are used in a-way that causes a size requirement to become impossible to express.-Real programs that encounter this issue tend to be complicated, but to-illustrate the problem, consider the following contrived function:--.. code-block:: futhark--  def f (n: i64) (m: i64) (b: [n][m]bool) = b[0,0]--We have the following type:--.. code-block:: futhark--  val f : (n: i64) -> (m: i64) -> (b: [n][m]bool) -> bool--Now suppose we say:--.. code-block:: futhark--  def g = uncurry f--What should be the type of ``g``?  Intuitively, something like this:--.. code-block:: futhark--  val g : (n: i64, m: i64) -> (b: [n][m]bool) -> bool--But this is *not* expressible in the Futhark type system - and even if-it were, it would not be easy to infer this in general, as it depends-on exactly what ``uncurry`` does, which the type checker does not-know.--As a workaround, we can use explicit type annotation and size-coercions to give ``g`` an acceptable type:--.. code-block:: futhark--  def g [a][b] (n,m) (b: [a][b]bool) = f n m (b :> [n][m]bool)--Another workaround, which is often the right one in cases not as-contrived as above, is to modify ``f`` itself to produce a *witness*-of the constraint, in the form of an array of shape ``[n][m]``:--.. code-block:: futhark--  def f (n: i64) (m: i64) : ([n][m](), [n][m]bool -> bool) =-    (replicate n (replicate m ()), \b -> b[0,0])--Then ``uncurry f`` works just fine and has the following type:--.. code-block:: futhark--  (i64, i64) -> ?[n][m].([n][m](), [n][m]bool -> bool)--Programming with such *explicit size witnesses* is a fairly advanced-technique, but often necessary when writing advanced size-dependent-code.- .. _unify-consuming-param:  "Parameter types *x* and *y* are incompatible regarding consuming their arguments

@@ -758,17 +758,22 @@  Construct a signed integer array whose first element is ``x`` and which proceeds with a stride of ``y-x`` until reaching ``z``-(inclusive).  The ``..y`` part can be elided in which case a stride of-1 is used.  A run-time error occurs if ``z`` is less than ``x`` or-``y``, or if ``x`` and ``y`` are the same value.+(inclusive). The ``..y`` part can be elided in which case a stride of+``1`` is used. All components must be of the same unsigned integer+type. +A run-time error occurs if ``z`` is less than ``x`` or ``y``, or if+``x`` and ``y`` are the same value.+ In the general case, the size of the array produced by a range is unknown (see `Size types`_).  In a few cases, the size is known statically: -  * ``1..2...n`` has size ``n``+  * ``0..<n`` has size ``n``. -This holds only if ``n`` is a variable or constant.+  * ``0..1..<n`` has size ``n``.++  * ``1..2...n`` has size ``n``  .. _range_upto: 

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.25.24+version:        0.25.25 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to

@@ -264,13 +264,13 @@   letExp "zero" $ zeroExp t  fwdSOAC :: Pat Type -> StmAux () -> SOAC SOACS -> ADM ()-fwdSOAC pat aux (Screma size xs (ScremaForm scs reds f)) = do+fwdSOAC pat aux (Screma size xs (ScremaForm f scs reds)) = do   pat' <- bundleNewPat pat   xs' <- bundleTangents xs+  f' <- fwdLambda f   scs' <- mapM fwdScan scs   reds' <- mapM fwdRed reds-  f' <- fwdLambda f-  addStm $ Let pat' aux $ Op $ Screma size xs' $ ScremaForm scs' reds' f'+  addStm $ Let pat' aux $ Op $ Screma size xs' $ ScremaForm f' scs' reds'   where     fwdScan :: Scan SOACS -> ADM (Scan SOACS)     fwdScan sc = do

@@ -227,7 +227,7 @@     cs == mempty,     [map_stm] <- stmsToList (bodyStms lam_body),     (Let (Pat [pe]) _ (Op scrm)) <- map_stm,-    (Screma _ [a1, a2] (ScremaForm [] [] map_lam)) <- scrm,+    (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

@@ -76,7 +76,7 @@   [Ident] ->   SOAC rep ->   m (Maybe (MapNest rep))-fromSOAC' bound (SOAC.Screma w inps (SOAC.ScremaForm [] [] lam)) = do+fromSOAC' bound (SOAC.Screma w inps (SOAC.ScremaForm lam [] [])) = do   maybenest <- case ( stmsToList $ bodyStms $ lambdaBody lam,                       bodyResult $ lambdaBody lam                     ) of

@@ -435,7 +435,7 @@ lambda :: SOAC rep -> Lambda rep lambda (Stream _ _ _ lam) = lam lambda (Scatter _len _ivs _spec lam) = lam-lambda (Screma _ _ (ScremaForm _ _ lam)) = lam+lambda (Screma _ _ (ScremaForm lam _ _)) = lam lambda (Hist _ _ _ lam) = lam  -- | Set the lambda used in the SOAC.@@ -444,8 +444,8 @@   Stream w arrs nes lam setLambda lam (Scatter len arrs spec _lam) =   Scatter len arrs spec lam-setLambda lam (Screma w arrs (ScremaForm scan red _)) =-  Screma w arrs (ScremaForm scan red lam)+setLambda lam (Screma w arrs (ScremaForm _ scan red)) =+  Screma w arrs (ScremaForm lam scan red) setLambda lam (Hist w ops inps _) =   Hist w ops inps lam 

@@ -226,6 +226,29 @@       BoolValue False -> "false"       FloatValue v -> prettyText v +updates ::+  UncheckedExp ->+  (UncheckedExp, [(Fmt, Fmt)])+updates (RecordUpdate src fs ve _ _) = second (++ [(fs', ve')]) $ updates src+  where+    fs' = sep "." $ fmt <$> fs+    ve' = fmt ve+updates (Update src is ve _) = second (++ [(is', ve')]) $ updates src+  where+    is' = brackets $ sep ("," <> space) $ map fmt is+    ve' = fmt ve+updates e = (e, [])++fmtUpdate :: UncheckedExp -> Fmt+fmtUpdate e =+  -- Special case multiple chained Updates/RecordUpdates.+  let (root, us) = updates e+      loc = srclocOf e+   in addComments loc . localLayout loc $+        fmt root <+> align (sep line (map fmtWith us))+  where+    fmtWith (fs', v) = "with" <+> fs' <+> "=" <+> v+ instance Format UncheckedExp where   fmt (Var name _ loc) = addComments loc $ fmtQualName name   fmt (Hole _ loc) = addComments loc "???"@@ -246,16 +269,8 @@   fmt (Project k e _ loc) = addComments loc $ fmt e <> "." <> fmt k   fmt (Negate e loc) = addComments loc $ "-" <> fmt e   fmt (Not e loc) = addComments loc $ "!" <> fmt e-  fmt (Update src idxs ve loc) =-    addComments loc $-      fmt src <+> "with" <+> idxs' <+> stdNest ("=" </> fmt ve)-    where-      idxs' = brackets $ sep ("," <> space) $ map fmt idxs-  fmt (RecordUpdate src fs ve _ loc) =-    addComments loc $-      fmt src <+> "with" <+> fs' <+> stdNest ("=" </> fmt ve)-    where-      fs' = sep "." $ fmt <$> fs+  fmt e@Update {} = fmtUpdate e+  fmt e@RecordUpdate {} = fmtUpdate e   fmt (Assert e1 e2 _ loc) =     addComments loc $ "assert" <+> fmt e1 <+> fmt e2   fmt (Lambda params body rettype _ loc) =@@ -303,7 +318,7 @@  instance Format (AppExpBase NoInfo Name) where   fmt (BinOp (bop, _) _ (x, _) (y, _) loc) =-    addComments loc $ fmt x </> fmtBinOp bop <+> fmt y+    addComments loc $ align (fmt x) </> fmtBinOp bop <+> align (fmt y)   fmt (Match e cs loc) =     addComments loc $ "match" <+> fmt e </> sep line (map fmt $ toList cs)   -- need some way to omit the inital value expression, when this it's trivial

@@ -727,24 +727,25 @@           <*> p     pScremaForm =       SOAC.ScremaForm-        <$> braces (pScan pr `sepBy` pComma)+        <$> pLambda pr         <* pComma-        <*> braces (pReduce pr `sepBy` pComma)+        <*> braces (pScan pr `sepBy` pComma)         <* pComma-        <*> pLambda pr+        <*> braces (pReduce pr `sepBy` pComma)     pRedomapForm =-      SOAC.ScremaForm mempty-        <$> braces (pReduce pr `sepBy` pComma)+      SOAC.ScremaForm+        <$> pLambda pr+        <*> pure []         <* pComma-        <*> pLambda pr+        <*> braces (pReduce pr `sepBy` pComma)     pScanomapForm =       SOAC.ScremaForm-        <$> braces (pScan pr `sepBy` pComma)+        <$> pLambda pr         <* pComma-        <*> pure mempty-        <*> pLambda pr+        <*> braces (pScan pr `sepBy` pComma)+        <*> pure []     pMapForm =-      SOAC.ScremaForm mempty mempty <$> pLambda pr+      SOAC.ScremaForm <$> pLambda pr <*> pure mempty <*> pure mempty     pScatter =       keyword "scatter"         *> parens@@ -797,19 +798,19 @@     pVJP =       parens $         SOAC.VJP-          <$> pLambda pr+          <$> braces (pSubExp `sepBy` pComma)           <* pComma           <*> braces (pSubExp `sepBy` pComma)           <* pComma-          <*> braces (pSubExp `sepBy` pComma)+          <*> pLambda pr     pJVP =       parens $         SOAC.JVP-          <$> pLambda pr+          <$> braces (pSubExp `sepBy` pComma)           <* pComma           <*> braces (pSubExp `sepBy` pComma)           <* pComma-          <*> braces (pSubExp `sepBy` pComma)+          <*> pLambda pr  pSizeClass :: Parser GPU.SizeClass pSizeClass =

@@ -55,7 +55,7 @@     expUsesAD (Op JVP {}) = True     expUsesAD (Op VJP {}) = True     expUsesAD (Op (Stream _ _ _ lam)) = lamUsesAD lam-    expUsesAD (Op (Screma _ _ (ScremaForm scans reds lam))) =+    expUsesAD (Op (Screma _ _ (ScremaForm lam scans reds))) =       lamUsesAD lam         || any (lamUsesAD . scanLambda) scans         || any (lamUsesAD . redLambda) reds

@@ -125,9 +125,9 @@     -- The final lambda produces indexes and values for the 'HistOp's.     Hist SubExp [VName] [HistOp rep] (Lambda rep)   | -- FIXME: this should not be here-    JVP (Lambda rep) [SubExp] [SubExp]+    JVP [SubExp] [SubExp] (Lambda rep)   | -- FIXME: this should not be here-    VJP (Lambda rep) [SubExp] [SubExp]+    VJP [SubExp] [SubExp] (Lambda rep)   | -- | A combination of scan, reduction, and map.  The first     -- t'SubExp' is the size of the input arrays.     Screma SubExp [VName] (ScremaForm rep)@@ -148,14 +148,14 @@ -- | The essential parts of a 'Screma' factored out (everything -- except the input arrays). data ScremaForm rep = ScremaForm-  { scremaScans :: [Scan rep],-    scremaReduces :: [Reduce rep],-    -- | The "main" lambda of the Screma. For a map, this is+  { -- | The "main" lambda of the Screma. For a map, this is     -- equivalent to 'isMapSOAC'. Note that the meaning of the return     -- value of this lambda depends crucially on exactly which Screma     -- this is. The parameters will correspond exactly to elements of     -- the input arrays, however.-    scremaLambda :: Lambda rep+    scremaLambda :: Lambda rep,+    scremaScans :: [Scan rep],+    scremaReduces :: [Reduce rep]   }   deriving (Eq, Ord, Show) @@ -221,7 +221,7 @@ -- | The types produced by a single 'Screma', given the size of the -- input array. scremaType :: SubExp -> ScremaForm rep -> [Type]-scremaType w (ScremaForm scans reds map_lam) =+scremaType w (ScremaForm map_lam scans reds) =   scan_tps ++ red_tps ++ map (`arrayOfRow` w) map_tps   where     scan_tps =@@ -258,12 +258,12 @@ -- | Construct a Screma with possibly multiple scans, and -- the given map function. scanomapSOAC :: [Scan rep] -> Lambda rep -> ScremaForm rep-scanomapSOAC scans = ScremaForm scans []+scanomapSOAC scans lam = ScremaForm lam scans []  -- | Construct a Screma with possibly multiple reductions, and -- the given map function. redomapSOAC :: [Reduce rep] -> Lambda rep -> ScremaForm rep-redomapSOAC = ScremaForm []+redomapSOAC reds lam = ScremaForm lam [] reds  -- | Construct a Screma with possibly multiple scans, and identity map -- function.@@ -287,11 +287,11 @@  -- | Construct a Screma corresponding to a map. mapSOAC :: Lambda rep -> ScremaForm rep-mapSOAC = ScremaForm [] []+mapSOAC lam = ScremaForm lam [] []  -- | Does this Screma correspond to a scan-map composition? isScanomapSOAC :: ScremaForm rep -> Maybe ([Scan rep], Lambda rep)-isScanomapSOAC (ScremaForm scans reds map_lam) = do+isScanomapSOAC (ScremaForm map_lam scans reds) = do   guard $ null reds   guard $ not $ null scans   pure (scans, map_lam)@@ -305,7 +305,7 @@  -- | Does this Screma correspond to a reduce-map composition? isRedomapSOAC :: ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)-isRedomapSOAC (ScremaForm scans reds map_lam) = do+isRedomapSOAC (ScremaForm map_lam scans reds) = do   guard $ null scans   guard $ not $ null reds   pure (reds, map_lam)@@ -320,7 +320,7 @@ -- | Does this Screma correspond to a simple map, without any -- reduction or scan results? isMapSOAC :: ScremaForm rep -> Maybe (Lambda rep)-isMapSOAC (ScremaForm scans reds map_lam) = do+isMapSOAC (ScremaForm map_lam scans reds) = do   guard $ null scans   guard $ null reds   pure map_lam@@ -399,16 +399,16 @@   SOACMapper frep trep m ->   SOAC frep ->   m (SOAC trep)-mapSOACM tv (JVP lam args vec) =+mapSOACM tv (JVP args vec lam) =   JVP-    <$> mapOnSOACLambda tv lam-    <*> mapM (mapOnSOACSubExp tv) args+    <$> mapM (mapOnSOACSubExp tv) args     <*> mapM (mapOnSOACSubExp tv) vec-mapSOACM tv (VJP lam args vec) =+    <*> mapOnSOACLambda tv lam+mapSOACM tv (VJP args vec lam) =   VJP-    <$> mapOnSOACLambda tv lam-    <*> mapM (mapOnSOACSubExp tv) args+    <$> mapM (mapOnSOACSubExp tv) args     <*> mapM (mapOnSOACSubExp tv) vec+    <*> mapOnSOACLambda tv lam mapSOACM tv (Stream size arrs accs lam) =   Stream     <$> mapOnSOACSubExp tv size@@ -443,12 +443,13 @@       )       ops     <*> mapOnSOACLambda tv bucket_fun-mapSOACM tv (Screma w arrs (ScremaForm scans reds map_lam)) =+mapSOACM tv (Screma w arrs (ScremaForm map_lam scans reds)) =   Screma     <$> mapOnSOACSubExp tv w     <*> mapM (mapOnSOACVName tv) arrs     <*> ( ScremaForm-            <$> forM+            <$> mapOnSOACLambda tv map_lam+            <*> forM               scans               ( \(Scan red_lam red_nes) ->                   Scan@@ -462,7 +463,6 @@                     <$> mapOnSOACLambda tv red_lam                     <*> mapM (mapOnSOACSubExp tv) red_nes               )-            <*> mapOnSOACLambda tv map_lam         )  -- | A helper for defining 'TraverseOpStms'.@@ -514,12 +514,10 @@  -- | The type of a SOAC. soacType :: (Typed (LParamInfo rep)) => SOAC rep -> [Type]-soacType (JVP lam _ _) =-  lambdaReturnType lam-    ++ lambdaReturnType lam-soacType (VJP lam _ _) =-  lambdaReturnType lam-    ++ map paramType (lambdaParams lam)+soacType (JVP _ _ lam) =+  lambdaReturnType lam ++ lambdaReturnType lam+soacType (VJP _ _ lam) =+  lambdaReturnType lam ++ map paramType (lambdaParams lam) soacType (Stream outersize _ accs lam) =   map (substNamesInType substs) rtp   where@@ -547,7 +545,7 @@   consumedInOp VJP {} = mempty   -- Only map functions can consume anything.  The operands to scan   -- and reduce functions are always considered "fresh".-  consumedInOp (Screma _ arrs (ScremaForm _ _ map_lam)) =+  consumedInOp (Screma _ arrs (ScremaForm map_lam _ _)) =     mapNames consumedArray $ consumedByLambda map_lam     where       consumedArray v = fromMaybe v $ lookup v params_to_arrs@@ -572,10 +570,10 @@   HistOp w rf dests nes $ f lam  instance CanBeAliased SOAC where-  addOpAliases aliases (JVP lam args vec) =-    JVP (Alias.analyseLambda aliases lam) args vec-  addOpAliases aliases (VJP lam args vec) =-    VJP (Alias.analyseLambda aliases lam) args vec+  addOpAliases aliases (JVP args vec lam) =+    JVP args vec (Alias.analyseLambda aliases lam)+  addOpAliases aliases (VJP args vec lam) =+    VJP args vec (Alias.analyseLambda aliases lam)   addOpAliases aliases (Stream size arr accs lam) =     Stream size arr accs $ Alias.analyseLambda aliases lam   addOpAliases aliases (Scatter len arrs dests lam) =@@ -586,12 +584,12 @@       arrs       (map (mapHistOp (Alias.analyseLambda aliases)) ops)       (Alias.analyseLambda aliases bucket_fun)-  addOpAliases aliases (Screma w arrs (ScremaForm scans reds map_lam)) =+  addOpAliases aliases (Screma w arrs (ScremaForm map_lam scans reds)) =     Screma w arrs $       ScremaForm+        (Alias.analyseLambda aliases map_lam)         (map onScan scans)         (map onRed reds)-        (Alias.analyseLambda aliases map_lam)     where       onRed red = red {redLambda = Alias.analyseLambda aliases $ redLambda red}       onScan scan = scan {scanLambda = Alias.analyseLambda aliases $ scanLambda scan}@@ -631,18 +629,18 @@     where       flattenBlocks (_, arr, ivs) =         oneName arr <> mconcat (map (mconcat . fst) ivs) <> mconcat (map snd ivs)-  opDependencies (JVP lam args vec) =+  opDependencies (JVP args vec lam) =     mconcat $       replicate 2 $         lambdaDependencies mempty lam $           zipWith (<>) (map depsOf' args) (map depsOf' vec)-  opDependencies (VJP lam args vec) =+  opDependencies (VJP args vec lam) =     lambdaDependencies       mempty       lam       (zipWith (<>) (map depsOf' args) (map depsOf' vec))       <> map (const $ freeIn args <> freeIn lam) (lambdaParams lam)-  opDependencies (Screma w arrs (ScremaForm scans reds map_lam)) =+  opDependencies (Screma w arrs (ScremaForm map_lam scans reds)) =     let (scans_in, reds_in, map_deps) =           splitAt3 (scanResults scans) (redResults reds) $             lambdaDependencies mempty map_lam (depsOfArrays w arrs)@@ -682,7 +680,7 @@       SubExpRes _ (Var v) -> uncurry (flip ST.Indexed) <$> M.lookup v arr_indexes'       _ -> Nothing     where-      lambdaAndSubExp (Screma _ arrs (ScremaForm scans reds map_lam)) =+      lambdaAndSubExp (Screma _ arrs (ScremaForm map_lam scans reds)) =         nthMapOut (scanResults scans + redResults reds) map_lam arrs       lambdaAndSubExp _ =         Nothing@@ -713,7 +711,7 @@  -- | Type-check a SOAC. typeCheckSOAC :: (TC.Checkable rep) => SOAC (Aliases rep) -> TC.TypeM rep ()-typeCheckSOAC (VJP lam args vec) = do+typeCheckSOAC (VJP args vec lam) = do   args' <- mapM TC.checkArg args   TC.checkLambda lam $ map TC.noArgAliases args'   vec_ts <- mapM TC.checkSubExp vec@@ -723,7 +721,7 @@         </> PP.indent 2 (pretty (lambdaReturnType lam))         </> "does not match type of seed vector"         </> PP.indent 2 (pretty vec_ts)-typeCheckSOAC (JVP lam args vec) = do+typeCheckSOAC (JVP args vec lam) = do   args' <- mapM TC.checkArg args   TC.checkLambda lam $ map TC.noArgAliases args'   vec_ts <- mapM TC.checkSubExp vec@@ -849,7 +847,7 @@         <> prettyTuple (lambdaReturnType bucket_fun)         <> " but should have type "         <> prettyTuple bucket_ret_t-typeCheckSOAC (Screma w arrs (ScremaForm scans reds map_lam)) = do+typeCheckSOAC (Screma w arrs (ScremaForm map_lam scans reds)) = do   TC.require [Prim int64] w   arrs' <- TC.checkSOACArrayArgs w arrs   TC.checkLambda map_lam arrs'@@ -893,10 +891,10 @@       <> " wrong for given scan and reduction functions."  instance RephraseOp SOAC where-  rephraseInOp r (VJP lam args vec) =-    VJP <$> rephraseLambda r lam <*> pure args <*> pure vec-  rephraseInOp r (JVP lam args vec) =-    JVP <$> rephraseLambda r lam <*> pure args <*> pure vec+  rephraseInOp r (VJP args vec lam) =+    VJP args vec <$> rephraseLambda r lam+  rephraseInOp r (JVP args vec lam) =+    JVP args vec <$> rephraseLambda r lam   rephraseInOp r (Stream w arrs acc lam) =     Stream w arrs acc <$> rephraseLambda r lam   rephraseInOp r (Scatter w arrs dests lam) =@@ -906,21 +904,21 @@     where       onOp (HistOp dest_shape rf dests nes op) =         HistOp dest_shape rf dests nes <$> rephraseLambda r op-  rephraseInOp r (Screma w arrs (ScremaForm scans red lam)) =+  rephraseInOp r (Screma w arrs (ScremaForm lam scans red)) =     Screma w arrs       <$> ( ScremaForm-              <$> mapM onScan scans+              <$> rephraseLambda r lam+              <*> mapM onScan scans               <*> mapM onRed red-              <*> rephraseLambda r lam           )     where       onScan (Scan op nes) = Scan <$> rephraseLambda r op <*> pure nes       onRed (Reduce comm op nes) = Reduce comm <$> rephraseLambda r op <*> pure nes  instance (OpMetrics (Op rep)) => OpMetrics (SOAC rep) where-  opMetrics (VJP lam _ _) =+  opMetrics (VJP _ _ lam) =     inside "VJP" $ lambdaMetrics lam-  opMetrics (JVP lam _ _) =+  opMetrics (JVP _ _ lam) =     inside "JVP" $ lambdaMetrics lam   opMetrics (Stream _ _ _ lam) =     inside "Stream" $ lambdaMetrics lam@@ -928,32 +926,28 @@     inside "Scatter" $ lambdaMetrics lam   opMetrics (Hist _ _ ops bucket_fun) =     inside "Hist" $ mapM_ (lambdaMetrics . histOp) ops >> lambdaMetrics bucket_fun-  opMetrics (Screma _ _ (ScremaForm scans reds map_lam)) =+  opMetrics (Screma _ _ (ScremaForm map_lam scans reds)) =     inside "Screma" $ do+      lambdaMetrics map_lam       mapM_ (lambdaMetrics . scanLambda) scans       mapM_ (lambdaMetrics . redLambda) reds-      lambdaMetrics map_lam  instance (PrettyRep rep) => PP.Pretty (SOAC rep) where-  pretty (VJP lam args vec) =+  pretty (VJP args vec lam) =     "vjp"       <> parens         ( PP.align $-            pretty lam-              <> comma-                </> PP.braces (commasep $ map pretty args)-              <> comma-                </> PP.braces (commasep $ map pretty vec)+            PP.braces (commasep $ map pretty args)+              <> comma </> PP.braces (commasep $ map pretty vec)+              <> comma </> pretty lam         )-  pretty (JVP lam args vec) =+  pretty (JVP args vec lam) =     "jvp"       <> parens         ( PP.align $-            pretty lam-              <> comma-                </> PP.braces (commasep $ map pretty args)-              <> comma-                </> PP.braces (commasep $ map pretty vec)+            PP.braces (commasep $ map pretty args)+              <> comma </> PP.braces (commasep $ map pretty vec)+              <> comma </> pretty lam         )   pretty (Stream size arrs acc lam) =     ppStream size arrs acc lam@@ -961,56 +955,49 @@     ppScatter w arrs dests lam   pretty (Hist w arrs ops bucket_fun) =     ppHist w arrs ops bucket_fun-  pretty (Screma w arrs (ScremaForm scans reds map_lam))+  pretty (Screma w arrs (ScremaForm map_lam scans reds))     | null scans,       null reds =         "map"           <> (parens . align)             ( pretty w-                <> comma-                  </> ppTuple' (map pretty arrs)-                <> comma-                  </> pretty map_lam+                <> comma </> ppTuple' (map pretty arrs)+                <> comma </> pretty map_lam             )     | null scans =         "redomap"           <> (parens . align)             ( pretty w-                <> comma-                  </> ppTuple' (map pretty arrs)+                <> comma </> ppTuple' (map pretty arrs)+                <> comma </> pretty map_lam                 <> comma                   </> PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty reds)-                <> comma-                  </> pretty map_lam             )     | null reds =         "scanomap"           <> (parens . align)             ( pretty w-                <> comma-                  </> ppTuple' (map pretty arrs)-                <> comma-                  </> PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty scans)+                <> comma </> ppTuple' (map pretty arrs)+                <> comma </> pretty map_lam                 <> comma-                  </> pretty map_lam+                  </> PP.braces+                    (mconcat $ intersperse (comma <> PP.line) $ map pretty scans)             )   pretty (Screma w arrs form) = ppScrema w arrs form  -- | Prettyprint the given Screma. ppScrema ::   (PrettyRep rep, Pretty inp) => SubExp -> [inp] -> ScremaForm rep -> Doc ann-ppScrema w arrs (ScremaForm scans reds map_lam) =+ppScrema w arrs (ScremaForm map_lam scans reds) =   "screma"     <> (parens . align)       ( pretty w-          <> comma-            </> ppTuple' (map pretty arrs)+          <> comma </> ppTuple' (map pretty arrs)+          <> comma </> pretty map_lam           <> comma             </> PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty scans)           <> comma             </> PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty reds)-          <> comma-            </> pretty map_lam       )  -- | Prettyprint the given Stream.

@@ -82,16 +82,16 @@ simplifySOAC ::   (Simplify.SimplifiableRep rep) =>   Simplify.SimplifyOp rep (SOAC (Wise rep))-simplifySOAC (VJP lam arr vec) = do+simplifySOAC (VJP arr vec lam) = do   (lam', hoisted) <- Engine.simplifyLambda mempty lam   arr' <- mapM Engine.simplify arr   vec' <- mapM Engine.simplify vec-  pure (VJP lam' arr' vec', hoisted)-simplifySOAC (JVP lam arr vec) = do+  pure (VJP arr' vec' lam', hoisted)+simplifySOAC (JVP arr vec lam) = do   (lam', hoisted) <- Engine.simplifyLambda mempty lam   arr' <- mapM Engine.simplify arr   vec' <- mapM Engine.simplify vec-  pure (JVP lam' arr' vec', hoisted)+  pure (JVP arr' vec' lam', hoisted) simplifySOAC (Stream outerdim arr nes lam) = do   outerdim' <- Engine.simplify outerdim   nes' <- mapM Engine.simplify nes@@ -117,7 +117,7 @@   imgs' <- mapM Engine.simplify imgs   (bfun', bfun_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty bfun   pure (Hist w' imgs' ops' bfun', mconcat hoisted <> bfun_hoisted)-simplifySOAC (Screma w arrs (ScremaForm scans reds map_lam)) = do+simplifySOAC (Screma w arrs (ScremaForm map_lam scans reds)) = do   (scans', scans_hoisted) <- fmap unzip $     forM scans $ \(Scan lam nes) -> do       (lam', hoisted) <- Engine.simplifyLambda mempty lam@@ -136,7 +136,7 @@     <$> ( Screma             <$> Engine.simplify w             <*> Engine.simplify arrs-            <*> pure (ScremaForm scans' reds' map_lam')+            <*> pure (ScremaForm map_lam' scans' reds')         )     <*> pure (mconcat scans_hoisted <> mconcat reds_hoisted <> map_lam_hoisted) @@ -399,10 +399,10 @@   TopDownRuleOp rep removeUnusedSOACInput _ pat aux op   | Just (Screma w arrs form :: SOAC rep) <- asSOAC op,-    ScremaForm scan reduce map_lam <- form,+    ScremaForm map_lam scan reduce <- form,     Just (used_arrs, map_lam') <- remove map_lam arrs =       Simplify . auxing aux . letBind pat . Op $-        soacOp (Screma w used_arrs (ScremaForm scan reduce map_lam'))+        soacOp (Screma w used_arrs (ScremaForm map_lam' scan reduce))   | Just (Scatter w arrs dests map_lam :: SOAC rep) <- asSOAC op,     Just (used_arrs, map_lam') <- remove map_lam arrs =       Simplify . auxing aux . letBind pat . Op $@@ -418,7 +418,7 @@ removeUnusedSOACInput _ _ _ _ = Skip  removeDeadMapping :: BottomUpRuleOp (Wise SOACS)-removeDeadMapping (_, used) (Pat pes) aux (Screma w arrs (ScremaForm scans reds lam))+removeDeadMapping (_, used) (Pat pes) aux (Screma w arrs (ScremaForm lam scans reds))   | (nonmap_pes, map_pes) <- splitAt num_nonmap_res pes,     not $ null map_pes =       let (nonmap_res, map_res) = splitAt num_nonmap_res $ bodyResult $ lambdaBody lam@@ -434,10 +434,8 @@        in if map_pes /= map_pes'             then               Simplify . auxing aux $-                letBind (Pat $ nonmap_pes <> map_pes') $-                  Op $-                    Screma w arrs $-                      ScremaForm scans reds lam'+                letBind (Pat $ nonmap_pes <> map_pes') . Op $+                  Screma w arrs (ScremaForm lam' scans reds)             else Skip   where     num_nonmap_res = scanResults scans + redResults reds@@ -642,7 +640,7 @@   (Buildable rep, BuilderOps rep, HasSOAC rep) =>   TopDownRuleOp rep simplifyKnownIterationSOAC _ pat _ op-  | Just (Screma (Constant k) arrs (ScremaForm scans reds map_lam)) <- asSOAC op,+  | Just (Screma (Constant k) arrs (ScremaForm map_lam scans reds)) <- asSOAC op,     oneIsh k = Simplify $ do       let (Reduce _ red_lam red_nes) = singleReduce reds           (Scan scan_lam scan_nes) = singleScan scans@@ -697,7 +695,7 @@         certifying cs $ letBindNames [v] $ BasicOp $ SubExp se -- simplifyKnownIterationSOAC _ pat aux op-  | Just (Screma (Constant (IntValue (Int64Value k))) arrs (ScremaForm [] [] map_lam)) <- asSOAC op,+  | Just (Screma (Constant (IntValue (Int64Value k))) arrs (ScremaForm map_lam [] [])) <- asSOAC op,     "unroll" `inAttrs` stmAuxAttrs aux = Simplify $ do       arrs_elems <- fmap transpose . forM [0 .. k - 1] $ \i -> do         map_lam' <- renameLambda map_lam@@ -836,7 +834,7 @@   (Buildable rep, BuilderOps rep, HasSOAC rep) =>   TopDownRuleOp rep simplifyMapIota vtable screma_pat aux op-  | Just (Screma w arrs (ScremaForm scan reduce map_lam) :: SOAC rep) <- asSOAC op,+  | Just (Screma w arrs (ScremaForm map_lam scan reduce) :: SOAC rep) <- asSOAC op,     Just (p, _) <- find isIota (zip (lambdaParams map_lam) arrs),     indexings <-       mapMaybe (indexesWith (paramName p)) . S.toList $@@ -855,7 +853,7 @@               }        auxing aux . letBind screma_pat . Op . soacOp $-        Screma w (arrs <> more_arrs) (ScremaForm scan reduce map_lam')+        Screma w (arrs <> more_arrs) (ScremaForm map_lam' scan reduce)   where     isIota (_, arr) = case ST.lookupBasicOp arr vtable of       Just (Iota _ (Constant o) (Constant s) _, _) ->@@ -908,7 +906,7 @@ -- corresponding to that transformation performed on the rows of the -- full array. moveTransformToInput :: TopDownRuleOp (Wise SOACS)-moveTransformToInput vtable screma_pat aux soac@(Screma w arrs (ScremaForm scan reduce map_lam))+moveTransformToInput vtable screma_pat aux soac@(Screma w arrs (ScremaForm map_lam scan reduce))   | ops <- filter arrayIsMapParam $ S.toList $ arrayOps mempty $ lambdaBody map_lam,     not $ null ops = Simplify $ do       (more_arrs, more_params, replacements) <-@@ -923,7 +921,7 @@               }        auxing aux . letBind screma_pat . Op $-        Screma w (arrs <> more_arrs) (ScremaForm scan reduce map_lam')+        Screma w (arrs <> more_arrs) (ScremaForm map_lam' scan reduce)   where     -- It is not safe to move the transform if the root array is being     -- consumed by the Screma.  This is a bit too conservative - it's

@@ -215,10 +215,10 @@                )             ++ [ErrorVal int64 end'_i64, " is invalid."] -  (it, le_op, lt_op) <-+  (it, lt_op) <-     case E.typeOf start of-      E.Scalar (E.Prim (E.Signed it)) -> pure (it, CmpSle it, CmpSlt it)-      E.Scalar (E.Prim (E.Unsigned it)) -> pure (it, CmpUle it, CmpUlt it)+      E.Scalar (E.Prim (E.Signed it)) -> pure (it, CmpSlt it)+      E.Scalar (E.Prim (E.Unsigned it)) -> pure (it, CmpUlt it)       start_t -> error $ "Start value in range has type " ++ prettyString start_t    let one = intConst it 1@@ -245,7 +245,7 @@   bounds_invalid_downwards <-     letSubExp "bounds_invalid_downwards" $       I.BasicOp $-        I.CmpOp le_op start' end'+        I.CmpOp lt_op start' end'   bounds_invalid_upwards <-     letSubExp "bounds_invalid_upwards" $       I.BasicOp $@@ -1710,8 +1710,8 @@           lam <- internaliseLambdaCoerce f =<< mapM subExpType x'           fmap (map I.Var) . letTupExp desc . Op $             case fname of-              "jvp2" -> JVP lam x' v'-              _ -> VJP lam x' v'+              "jvp2" -> JVP x' v' lam+              _ -> VJP x' v' lam     handleAD _ _ = Nothing      handleRest [a, si, v] "scatter" = Just $ scatterF 1 a si v

@@ -6,6 +6,7 @@ import Control.Monad.Reader import Control.Monad.State import Data.Bifunctor+import Data.Bitraversable import Data.Foldable import Data.List (partition) import Data.Map.Strict qualified as M@@ -144,6 +145,12 @@ transformSubExps :: ASTMapper LiftM transformSubExps = identityMapper {mapOnExp = transformExp} +transformType :: TypeBase Exp u -> LiftM (TypeBase Exp u)+transformType = bitraverse transformExp pure++transformPat :: PatBase Info VName (TypeBase Exp u) -> LiftM (PatBase Info VName (TypeBase Exp u))+transformPat = traverse transformType+ transformExp :: Exp -> LiftM Exp transformExp (AppExp (LetFun fname (tparams, params, _, Info ret, funbody) body _) _) = do   funbody' <- bindingParams (map typeParamName tparams) params $ transformExp funbody@@ -156,8 +163,9 @@   liftFunction fname [] params ret body' <*> pure (typeOf e) transformExp (AppExp (LetPat sizes pat e body loc) appres) = do   e' <- transformExp e-  body' <- bindingLetPat (map sizeName sizes) pat $ transformExp body-  pure $ AppExp (LetPat sizes pat e' body' loc) appres+  pat' <- transformPat pat+  body' <- bindingLetPat (map sizeName sizes) pat' $ transformExp body+  pure $ AppExp (LetPat sizes pat' e' body' loc) appres transformExp (AppExp (Match e cases loc) appres) = do   e' <- transformExp e   cases' <- mapM transformCase cases@@ -173,10 +181,11 @@     form' <- astMap transformSubExps form     body' <- bindingForm form' $ transformExp body     pure $ AppExp (Loop sizes pat (LoopInitExplicit args') form' body' loc) appres-transformExp e@(Var v (Info t) _) =+transformExp (Var v (Info t) loc) = do+  t' <- transformType t   -- Note that function-typed variables can only occur in expressions,   -- not in other places where VNames/QualNames can occur.-  asks $ maybe e ($ t) . M.lookup (qualLeaf v) . envReplace+  asks $ maybe (Var v (Info t') loc) ($ t') . M.lookup (qualLeaf v) . envReplace transformExp e = astMap transformSubExps e  transformValBind :: ValBind -> LiftM ()

@@ -236,20 +236,8 @@ calculateDims body repl =   foldCalc top_repl $ expReplace top_repl body   where-    -- list of strict sub-expressions of e-    subExps e-      | Just e' <- stripExp e = subExps e'-      | otherwise = astMap mapper e `execState` mempty-      where-        mapOnExp e'-          | Just e'' <- stripExp e' = mapOnExp e''-          | otherwise = do-              modify (ReplacedExp e' :)-              astMap mapper e'-        mapper = identityMapper {mapOnExp}-    depends (a, _) (b, _) = b `elem` subExps (unReplaced a)-    top_repl =-      topologicalSort depends repl+    depends (a, _) (b, _) = unReplaced b `elem` subExps (unReplaced a)+    top_repl = topologicalSort depends repl      ---- Calculus insertion     foldCalc [] body' = pure body'

@@ -433,7 +433,7 @@  removeOutputsExcept :: [VName] -> NodeT -> NodeT removeOutputsExcept toKeep s = case s of-  SoacNode ots (Pat pats1) soac@(H.Screma _ _ (ScremaForm scans_1 red_1 lam_1)) aux1 ->+  SoacNode ots (Pat pats1) soac@(H.Screma _ _ (ScremaForm lam_1 scans_1 red_1)) aux1 ->     SoacNode ots (Pat $ pats_unchanged <> pats_new) (H.setLambda lam_new soac) aux1     where       scan_output_size = Futhark.scanResults scans_1@@ -541,12 +541,12 @@     cases' <- mapM (traverse $ renameBody <=< (`doFusionWithDelayed` to_fuse)) cases     defbody' <- doFusionWithDelayed defbody to_fuse     pure (incoming, node, MatchNode (Let pat aux (Match cond cases' defbody' dec)) [], outgoing)-  StmNode (Let pat aux (Op (Futhark.VJP lam args vec))) -> doFuseScans $ do+  StmNode (Let pat aux (Op (Futhark.VJP args vec lam))) -> doFuseScans $ do     lam' <- fst <$> doFusionInLambda lam-    pure (incoming, node, StmNode (Let pat aux (Op (Futhark.VJP lam' args vec))), outgoing)-  StmNode (Let pat aux (Op (Futhark.JVP lam args vec))) -> doFuseScans $ do+    pure (incoming, node, StmNode (Let pat aux (Op (Futhark.VJP args vec lam'))), outgoing)+  StmNode (Let pat aux (Op (Futhark.JVP args vec lam))) -> doFuseScans $ do     lam' <- fst <$> doFusionInLambda lam-    pure (incoming, node, StmNode (Let pat aux (Op (Futhark.JVP lam' args vec))), outgoing)+    pure (incoming, node, StmNode (Let pat aux (Op (Futhark.JVP args vec lam'))), outgoing)   StmNode (Let pat aux (WithAcc inputs lam)) -> doFuseScans $ do     lam' <- fst <$> doFusionInLambda lam     pure (incoming, node, StmNode (Let pat aux (WithAcc inputs lam')), outgoing)

@@ -210,7 +210,7 @@     isMap nT       | SoacNode out_trsfs _pat soac _ <- nT,         H.Screma _ _ form <- soac,-        ScremaForm [] [] _ <- form =+        ScremaForm _ [] [] <- form =           H.nullTransforms out_trsfs     isMap _ = False checkSafeAndProfitable _ _ _ _ = False@@ -341,7 +341,7 @@       mkWithAccBdy' static_arg dims (dim : dims_rev) (iot_par_nms ++ [paramName iota_p]) rshp_ps' cons_ps'     let map_lam = Lambda (rshp_ps' ++ [iota_p] ++ cons_ps') (map paramDec cons_ps') map_lam_bdy         map_inps = map paramName rshp_ps ++ [iota_arr] ++ map paramName cons_ps-        map_soac = F.Screma dim map_inps $ ScremaForm [] [] map_lam+        map_soac = F.Screma dim map_inps $ ScremaForm map_lam [] []     res_nms <- letTupExp "acc_res" $ Op map_soac     pure $ map (subExpRes . Var) res_nms 

@@ -262,11 +262,18 @@       | unfus_set /= mempty,         not (SOAC.nullTransforms $ fsOutputTransform ker) ->           fail "Cannot perform diagonal fusion in the presence of output transforms."-    ( SOAC.Screma _ _ (ScremaForm scans_c reds_c _),-      SOAC.Screma _ _ (ScremaForm scans_p reds_p _),+    ( SOAC.Screma _ _ (ScremaForm _ scans_c reds_c),+      SOAC.Screma _ _ (ScremaForm _ scans_p reds_p),       _       )-        | scremaFusionOK (splitAt (Futhark.scanResults scans_p + Futhark.redResults reds_p) outVars) ker -> do+        | scremaFusionOK+            ( splitAt+                ( Futhark.scanResults scans_p+                    + Futhark.redResults reds_p+                )+                outVars+            )+            ker -> do             let red_nes_p = concatMap redNeutral reds_p                 red_nes_c = concatMap redNeutral reds_c                 scan_nes_p = concatMap scanNeutral scans_p@@ -300,7 +307,7 @@               $ SOAC.Screma                 w                 new_inp-                (ScremaForm (scans_p ++ scans_c) (reds_p ++ reds_c) res_lam')+                (ScremaForm res_lam' (scans_p ++ scans_c) (reds_p ++ reds_c))      ------------------     -- Scatter fusion --@@ -587,7 +594,7 @@             t : _ -> 1 : 0 : [2 .. arrayRank t]        pure-        ( SOAC.Screma map_w map_arrs' (ScremaForm [] [] map_fun'),+        ( SOAC.Screma map_w map_arrs' (mapSOAC map_fun'),           ots SOAC.|> SOAC.Rearrange map_cs perm         ) iswim _ _ _ =

@@ -150,7 +150,7 @@       map_lam <- renameLambda map_lam0       (k1_res, ker1_stms) <- runBuilderT' $ do         iota <- letExp "iota" $ BasicOp $ Iota inv_dim_len (intConst Int64 0) (intConst Int64 1) Int64-        let op_exp = Op (OtherOp (Screma inv_dim_len [iota] (ScremaForm [] [red] map_lam)))+        let op_exp = Op (OtherOp (Screma inv_dim_len [iota] (ScremaForm map_lam [] [red])))         res_redmap <- letTupExp "res_mapred" op_exp         letSubExp (baseString pat_acc_nm ++ "_big_update") $           BasicOp (UpdateAcc safety acc_nm acc_inds $ map Var res_redmap)

@@ -36,20 +36,20 @@     certifying cs $ letBindNames [v] $ BasicOp $ SubExp se  onStm :: Mode -> Scope SOACS -> Stm SOACS -> PassM (Stms SOACS)-onStm mode scope (Let pat aux (Op (VJP lam args vec))) = do+onStm mode scope (Let pat aux (Op (VJP args vec lam))) = do   lam' <- onLambda mode scope lam   if mode == All || lam == lam'     then do       lam'' <- (`runReaderT` scope) . simplifyLambda =<< revVJP scope lam'       runBuilderT_ (bindLambda pat aux lam'' $ args ++ vec) scope-    else pure $ oneStm $ Let pat aux $ Op $ VJP lam' args vec-onStm mode scope (Let pat aux (Op (JVP lam args vec))) = do+    else pure $ oneStm $ Let pat aux $ Op $ VJP args vec lam'+onStm mode scope (Let pat aux (Op (JVP args vec lam))) = do   lam' <- onLambda mode scope lam   if mode == All || lam == lam'     then do       lam'' <- fwdJVP scope lam'       runBuilderT_ (bindLambda pat aux lam'' $ args ++ vec) scope-    else pure $ oneStm $ Let pat aux $ Op $ JVP lam' args vec+    else pure $ oneStm $ Let pat aux $ Op $ JVP args vec lam' onStm mode scope (Let pat aux e) = oneStm . Let pat aux <$> mapExpM mapper e   where     mapper =

@@ -596,7 +596,7 @@             max             (bodyInterest defbody)             (map (bodyInterest . caseBody) cases)-      | Op (Screma w _ (ScremaForm _ _ lam')) <- stmExp stm =+      | Op (Screma w _ (ScremaForm lam' _ _)) <- stmExp stm =           zeroIfTooSmall w + bodyInterest (lambdaBody lam')       | Op (Stream _ _ _ lam') <- stmExp stm =           bodyInterest $ lambdaBody lam'@@ -625,7 +625,7 @@     interest depth stm       | "sequential" `inAttrs` attrs =           0 :: Int-      | Op (Screma _ _ form@(ScremaForm _ _ lam')) <- stmExp stm,+      | Op (Screma _ _ form@(ScremaForm lam' _ _)) <- stmExp stm,         isJust $ isMapSOAC form =           if sequential_inner             then 0@@ -636,7 +636,7 @@           bodyInterest (depth + 1) body * 10       | WithAcc _ withacc_lam <- stmExp stm =           bodyInterest (depth + 1) (lambdaBody withacc_lam)-      | Op (Screma _ _ form@(ScremaForm _ _ lam')) <- stmExp stm =+      | Op (Screma _ _ form@(ScremaForm lam' _ _)) <- stmExp stm =           1             + bodyInterest (depth + 1) (lambdaBody lam')             +

@@ -106,18 +106,16 @@   ScremaForm (Rep m) ->   [VName] ->   m ()-dissectScrema pat w (ScremaForm scans reds map_lam) arrs = do+dissectScrema pat w (ScremaForm map_lam scans reds) arrs = do   let num_reds = redResults reds       num_scans = scanResults scans-      (scan_res, red_res, map_res) =-        splitAt3 num_scans num_reds $ patNames pat+      (scan_res, red_res, map_res) = splitAt3 num_scans num_reds $ patNames pat    to_red <- replicateM num_reds $ newVName "to_red"    let scanomap = scanomapSOAC scans map_lam   letBindNames (scan_res <> to_red <> map_res) $-    Op $-      Screma w arrs scanomap+    Op (Screma w arrs scanomap)    reduce <- reduceSOAC reds   letBindNames red_res $ Op $ Screma w to_red reduce

@@ -125,7 +125,7 @@   error "transformSOAC: unhandled JVP" transformSOAC _ VJP {} =   error "transformSOAC: unhandled VJP"-transformSOAC pat (Screma w arrs form@(ScremaForm scans reds map_lam)) = do+transformSOAC pat (Screma w arrs form@(ScremaForm map_lam scans reds)) = do   -- See Note [Translation of Screma].   --   -- Start by combining all the reduction and scan parts into a single

@@ -264,8 +264,8 @@ asInteger (ValuePrim (SignedValue v)) = P.valueIntegral v asInteger (ValuePrim (UnsignedValue v)) =   toInteger (P.valueIntegral (P.doZExt v Int64) :: Word64)-asInteger (ValueAD d v)-  | P.IntValue v' <- AD.primitive $ AD.primal $ AD.Variable d v =+asInteger (ValueAD _ v)+  | P.IntValue v' <- AD.varPrimal v =       P.valueIntegral v' asInteger v = error $ "Unexpectedly not an integer: " <> show v @@ -274,8 +274,8 @@  asSigned :: Value -> IntValue asSigned (ValuePrim (SignedValue v)) = v-asSigned (ValueAD d v)-  | P.IntValue v' <- AD.primitive $ AD.primal $ AD.Variable d v = v'+asSigned (ValueAD _ v)+  | P.IntValue v' <- AD.varPrimal v = v' asSigned v = error $ "Unexpectedly not a signed integer: " <> show v  asInt64 :: Value -> Int64@@ -283,8 +283,8 @@  asBool :: Value -> Bool asBool (ValuePrim (BoolValue x)) = x-asBool (ValueAD d v)-  | P.BoolValue v' <- AD.primitive $ AD.primal $ AD.Variable d v = v'+asBool (ValueAD _ v)+  | P.BoolValue v' <- AD.varPrimal v = v' asBool v = error $ "Unexpectedly not a boolean: " <> show v  lookupInEnv ::@@ -2010,7 +2010,7 @@         let drvs = M.map (Just . putAD) $ M.unionsWith add $ map snd m          -- Extract the output values, and the partial derivatives-        let ov = modifyValue (\i _ -> fst $ m !! i) o+        let ov = modifyValue (\i _ -> fst $ m !! (length m - 1 - i)) o         let od =               fromMaybe (error "vjp: differentiation failed") $                 modifyValueM (\i vo -> M.findWithDefault (ValuePrim . putV . P.blankPrimValue . P.primValueType . AD.primitive <$> getAD vo) i drvs) v

@@ -7,9 +7,9 @@     JVPValue (..),     doOp,     addFor,-    primal,     tapePrimal,     primitive,+    varPrimal,     deriveTape,   ) where@@ -97,8 +97,12 @@ primal (Constant v) = Constant v  primitive :: ADValue -> PrimValue-primitive v@(Variable _ _) = primitive $ primal v+primitive (Variable _ v) = varPrimal v primitive (Constant v) = v++varPrimal :: ADVariable -> PrimValue+varPrimal (VJP (VJPValue t)) = primitive $ tapePrimal t+varPrimal (JVP (JVPValue v _)) = primitive $ primal v  -- Evaluates a PrimExp using doOp evalPrimExp :: M.Map VName ADValue -> PrimExp VName -> Maybe ADValue

@@ -152,10 +152,14 @@     pprPrec _ ValueAcc {} = "#<acc>"     pprPrec p (ValueSum _ n vs) =       parensIf (p > (0 :: Int)) $ "#" <> sep (pretty n : map (pprPrec 1) vs)-    -- TODO: This could be prettier. Perhaps add pretty printing for ADVariable / ADValues-    pprPrec _ (ValueAD d v) = pretty $ "d[" ++ show d ++ "]" ++ show v+    pprPrec _ (ValueAD _ v) = pprPrim $ putV $ AD.varPrimal v     pprElem v@ValueArray {} = pprPrec 0 v     pprElem v = group $ pprPrec 0 v++    putV (P.IntValue x) = SignedValue x+    putV (P.FloatValue x) = FloatValue x+    putV (P.BoolValue x) = BoolValue x+    putV P.UnitValue = BoolValue True  -- | Prettyprint value. prettyValue :: Value m -> Doc a

@@ -31,7 +31,9 @@     valBindBound,     funType,     stripExp,+    subExps,     similarExps,+    sameExp,      -- * Queries on patterns and params     patIdents,@@ -1371,6 +1373,20 @@ stripExp (Ascript e _ _) = stripExp e `mplus` Just e stripExp _ = Nothing +-- | All non-trivial subexpressions (as by stripExp) of some+-- expression, not including the expression itself.+subExps :: Exp -> [Exp]+subExps e+  | Just e' <- stripExp e = subExps e'+  | otherwise = astMap mapper e `execState` mempty+  where+    mapOnExp e'+      | Just e'' <- stripExp e' = mapOnExp e''+      | otherwise = do+          modify (e' :)+          astMap mapper e'+    mapper = identityMapper {mapOnExp}+ similarSlices :: Slice -> Slice -> Maybe [(Exp, Exp)] similarSlices slice1 slice2   | length slice1 == length slice2 = do@@ -1453,6 +1469,14 @@ similarExps (IndexSection slice1 _ _) (IndexSection slice2 _ _) =   similarSlices slice1 slice2 similarExps _ _ = Nothing++-- | Are these the same expression as per recursively invoking+-- 'similarExps'?+sameExp :: Exp -> Exp -> Bool+sameExp e1 e2+  | Just es <- similarExps e1 e2 =+      all (uncurry sameExp) es+  | otherwise = False  -- | An identifier with type- and aliasing information. type Ident = IdentBase Info VName

@@ -15,7 +15,6 @@  import Control.Monad import Control.Monad.Except-import Control.Monad.Identity import Control.Monad.Reader import Control.Monad.State.Strict import Data.Bifunctor@@ -27,7 +26,7 @@ import Data.Map.Strict qualified as M import Data.Maybe import Data.Set qualified as S-import Futhark.Util (mapAccumLM, nubOrd, topologicalSort)+import Futhark.Util (mapAccumLM, nubOrd) import Futhark.Util.Pretty hiding (space) import Language.Futhark import Language.Futhark.Primitive (intByteSize)@@ -229,92 +228,6 @@               loc               "a size coercion where the underlying expression size cannot be determined"           pure $ sizeFromName (qualName v) (srclocOf d)--sameExp :: Exp -> Exp -> Bool-sameExp e1 e2-  | Just es <- similarExps e1 e2 =-      all (uncurry sameExp) es-  | otherwise = False---- All non-trivial subexpressions (as by stripExp) of some expression,--- not including the expression itself.-subExps :: Exp -> [Exp]-subExps e-  | Just e' <- stripExp e = subExps e'-  | otherwise = astMap mapper e `execState` mempty-  where-    mapOnExp e'-      | Just e'' <- stripExp e' = mapOnExp e''-      | otherwise = do-          modify (e' :)-          astMap mapper e'-    mapper = identityMapper {mapOnExp}---- Expressions witnessed by type, topologically sorted.-topWit :: TypeBase Exp u -> [Exp]-topWit = topologicalSort depends . witnessedExps-  where-    witnessedExps t = execState (traverseDims onDim t) mempty-      where-        onDim _ PosImmediate e = modify (e :)-        onDim _ _ _ = pure ()-    depends a b = any (sameExp b) $ subExps a--sizeFree ::-  SrcLoc ->-  (Exp -> Maybe VName) ->-  TypeBase Size u ->-  TermTypeM (TypeBase Size u, [VName])-sizeFree tloc expKiller orig_t = do-  runReaderT (toBeReplaced orig_t $ onType orig_t) mempty `runStateT` mempty-  where-    lookReplacement e repl = snd <$> find (sameExp e . fst) repl-    expReplace mapping e-      | Just e' <- lookReplacement e mapping = e'-      | otherwise = runIdentity $ astMap mapper e-      where-        mapper = identityMapper {mapOnExp = pure . expReplace mapping}--    replacing e = do-      e' <- asks (`expReplace` e)-      case expKiller e' of-        Nothing -> pure e'-        Just cause -> do-          vn <- lift $ lift $ newRigidDim tloc (RigidOutOfScope (locOf e) cause) "d"-          modify (vn :)-          pure $ sizeFromName (qualName vn) (srclocOf e)--    toBeReplaced t m' = foldl f m' $ topWit t-      where-        f m e = do-          e' <- replacing e-          local ((e, e') :) m--    onScalar (Record fs) =-      Record <$> traverse onType fs-    onScalar (Sum cs) =-      Sum <$> (traverse . traverse) onType cs-    onScalar (Arrow as pn d argT (RetType dims retT)) = do-      argT' <- onType argT-      old_bound <- get-      retT' <- toBeReplaced retT $ onType retT-      rl <- state $ partition (`notElem` old_bound)-      let dims' = dims <> rl-      pure $ Arrow as pn d argT' (RetType dims' retT')-    onScalar (TypeVar u v args) =-      TypeVar u v <$> mapM onTypeArg args-      where-        onTypeArg (TypeArgDim d) = TypeArgDim <$> replacing d-        onTypeArg (TypeArgType ty) = TypeArgType <$> onType ty-    onScalar (Prim pt) = pure $ Prim pt--    onType ::-      TypeBase Size u ->-      ReaderT [(Exp, Exp)] (StateT [VName] TermTypeM) (TypeBase Size u)-    onType (Array u shape scalar) =-      Array u <$> traverse replacing shape <*> onScalar scalar-    onType (Scalar ty) =-      Scalar <$> onScalar ty  -- Used to remove unknown sizes from function body types before we -- perform let-generalisation.  This is because if a function is

@@ -11,6 +11,7 @@     Rigidity (..),     RigidSource (..),     BreadCrumbs,+    sizeFree,     noBreadCrumbs,     hasNoBreadCrumbs,     dimNotes,@@ -30,14 +31,18 @@  import Control.Monad import Control.Monad.Except+import Control.Monad.Identity+import Control.Monad.Reader import Control.Monad.State import Data.List qualified as L import Data.Map.Strict qualified as M import Data.Maybe import Data.Set qualified as S import Data.Text qualified as T+import Futhark.Util (topologicalSort) import Futhark.Util.Pretty import Language.Futhark+import Language.Futhark.Traversals import Language.Futhark.TypeChecker.Monad hiding (BoundV) import Language.Futhark.TypeChecker.Types @@ -616,6 +621,74 @@           <+> dquotes (prettyName v)           <+> "is rigidly bound in a deeper scope." +-- Expressions witnessed by type, topologically sorted.+topWit :: TypeBase Exp u -> [Exp]+topWit = topologicalSort depends . witnessedExps+  where+    witnessedExps t = execState (traverseDims onDim t) mempty+      where+        onDim _ PosImmediate e = modify (e :)+        onDim _ _ _ = pure ()+    depends a b = any (sameExp b) $ subExps a++sizeFree ::+  (MonadUnify m) =>+  SrcLoc ->+  (Exp -> Maybe VName) ->+  TypeBase Size u ->+  m (TypeBase Size u, [VName])+sizeFree tloc expKiller orig_t = do+  runReaderT (toBeReplaced orig_t $ onType orig_t) mempty `runStateT` mempty+  where+    lookReplacement e repl = snd <$> L.find (sameExp e . fst) repl+    expReplace mapping e+      | Just e' <- lookReplacement e mapping = e'+      | otherwise = runIdentity $ astMap mapper e+      where+        mapper = identityMapper {mapOnExp = pure . expReplace mapping}++    replacing e = do+      e' <- asks (`expReplace` e)+      case expKiller e' of+        Nothing -> pure e'+        Just cause -> do+          vn <- lift $ lift $ newRigidDim tloc (RigidOutOfScope (locOf e) cause) "d"+          modify (vn :)+          pure $ sizeFromName (qualName vn) (srclocOf e)++    toBeReplaced t m' = foldl f m' $ topWit t+      where+        f m e = do+          e' <- replacing e+          local ((e, e') :) m++    onScalar (Record fs) =+      Record <$> traverse onType fs+    onScalar (Sum cs) =+      Sum <$> (traverse . traverse) onType cs+    onScalar (Arrow as pn d argT (RetType dims retT)) = do+      argT' <- onType argT+      old_bound <- get+      retT' <- toBeReplaced retT $ onType retT+      rl <- state $ L.partition (`notElem` old_bound)+      let dims' = dims <> rl+      pure $ Arrow as pn d argT' (RetType dims' retT')+    onScalar (TypeVar u v args) =+      TypeVar u v <$> mapM onTypeArg args+      where+        onTypeArg (TypeArgDim d) = TypeArgDim <$> replacing d+        onTypeArg (TypeArgType ty) = TypeArgType <$> onType ty+    onScalar (Prim pt) = pure $ Prim pt++    onType ::+      (MonadUnify m) =>+      TypeBase Size u ->+      ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Size u)+    onType (Array u shape scalar) =+      Array u <$> traverse replacing shape <*> onScalar scalar+    onType (Scalar ty) =+      Scalar <$> onScalar ty+ linkVarToType ::   (MonadUnify m) =>   UnifySizes m ->@@ -633,21 +706,25 @@   occursCheck usage bcs vn tp   scopeCheck usage bcs vn lvl tp -  constraints <- getConstraints   let link = do         let (witnessed, not_witnessed) = determineSizeWitnesses tp             used v = v `S.member` witnessed || v `S.member` not_witnessed-            ext = filter used bound-        case filter (`notElem` witnessed) ext of-          [] ->-            modifyConstraints $-              M.insert vn (lvl, Constraint (RetType ext tp) usage)-          problems ->-            unifyError usage mempty bcs . withIndexLink "unify-param-existential" $-              "Parameter(s) "-                <> commasep (map (dquotes . prettyName) problems)-                <> " used as size(s) would go out of scope."+            (ext_witnessed, ext_not_witnessed) =+              L.partition (`elem` witnessed) $ filter used bound +            -- Any size that uses an ext_not_witnessed variable must+            -- be replaced with a fresh existential.+            problematic e =+              L.find (`elem` ext_not_witnessed) $+                S.toList $+                  fvVars $+                    freeInExp e++        (tp', ext_new) <- sizeFree (srclocOf usage) problematic tp++        modifyConstraints $+          M.insert vn (lvl, Constraint (RetType (ext_new <> ext_witnessed) tp') usage)+   let unliftedBcs unlifted_usage =         breadCrumb           ( Matching $@@ -658,6 +735,7 @@           )           bcs +  constraints <- getConstraints   case snd <$> M.lookup vn constraints of     Just (NoConstraint Unlifted unlift_usage) -> do       link