diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/docs/error-index.rst b/docs/error-index.rst --- a/docs/error-index.rst +++ b/docs/error-index.rst @@ -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 diff --git a/docs/language-reference.rst b/docs/language-reference.rst --- a/docs/language-reference.rst +++ b/docs/language-reference.rst @@ -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.. 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 diff --git a/src/Futhark/AD/Rev/SOAC.hs b/src/Futhark/AD/Rev/SOAC.hs --- a/src/Futhark/AD/Rev/SOAC.hs +++ b/src/Futhark/AD/Rev/SOAC.hs @@ -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 diff --git a/src/Futhark/Analysis/HORep/MapNest.hs b/src/Futhark/Analysis/HORep/MapNest.hs --- a/src/Futhark/Analysis/HORep/MapNest.hs +++ b/src/Futhark/Analysis/HORep/MapNest.hs @@ -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 diff --git a/src/Futhark/Analysis/HORep/SOAC.hs b/src/Futhark/Analysis/HORep/SOAC.hs --- a/src/Futhark/Analysis/HORep/SOAC.hs +++ b/src/Futhark/Analysis/HORep/SOAC.hs @@ -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 diff --git a/src/Futhark/Fmt/Printer.hs b/src/Futhark/Fmt/Printer.hs --- a/src/Futhark/Fmt/Printer.hs +++ b/src/Futhark/Fmt/Printer.hs @@ -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 diff --git a/src/Futhark/IR/Parse.hs b/src/Futhark/IR/Parse.hs --- a/src/Futhark/IR/Parse.hs +++ b/src/Futhark/IR/Parse.hs @@ -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 = diff --git a/src/Futhark/IR/SOACS.hs b/src/Futhark/IR/SOACS.hs --- a/src/Futhark/IR/SOACS.hs +++ b/src/Futhark/IR/SOACS.hs @@ -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 diff --git a/src/Futhark/IR/SOACS/SOAC.hs b/src/Futhark/IR/SOACS/SOAC.hs --- a/src/Futhark/IR/SOACS/SOAC.hs +++ b/src/Futhark/IR/SOACS/SOAC.hs @@ -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. diff --git a/src/Futhark/IR/SOACS/Simplify.hs b/src/Futhark/IR/SOACS/Simplify.hs --- a/src/Futhark/IR/SOACS/Simplify.hs +++ b/src/Futhark/IR/SOACS/Simplify.hs @@ -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 diff --git a/src/Futhark/Internalise/Exps.hs b/src/Futhark/Internalise/Exps.hs --- a/src/Futhark/Internalise/Exps.hs +++ b/src/Futhark/Internalise/Exps.hs @@ -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 diff --git a/src/Futhark/Internalise/LiftLambdas.hs b/src/Futhark/Internalise/LiftLambdas.hs --- a/src/Futhark/Internalise/LiftLambdas.hs +++ b/src/Futhark/Internalise/LiftLambdas.hs @@ -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 () diff --git a/src/Futhark/Internalise/Monomorphise.hs b/src/Futhark/Internalise/Monomorphise.hs --- a/src/Futhark/Internalise/Monomorphise.hs +++ b/src/Futhark/Internalise/Monomorphise.hs @@ -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' diff --git a/src/Futhark/Optimise/Fusion.hs b/src/Futhark/Optimise/Fusion.hs --- a/src/Futhark/Optimise/Fusion.hs +++ b/src/Futhark/Optimise/Fusion.hs @@ -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) diff --git a/src/Futhark/Optimise/Fusion/RulesWithAccs.hs b/src/Futhark/Optimise/Fusion/RulesWithAccs.hs --- a/src/Futhark/Optimise/Fusion/RulesWithAccs.hs +++ b/src/Futhark/Optimise/Fusion/RulesWithAccs.hs @@ -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 diff --git a/src/Futhark/Optimise/Fusion/TryFusion.hs b/src/Futhark/Optimise/Fusion/TryFusion.hs --- a/src/Futhark/Optimise/Fusion/TryFusion.hs +++ b/src/Futhark/Optimise/Fusion/TryFusion.hs @@ -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 _ _ _ = diff --git a/src/Futhark/Optimise/GenRedOpt.hs b/src/Futhark/Optimise/GenRedOpt.hs --- a/src/Futhark/Optimise/GenRedOpt.hs +++ b/src/Futhark/Optimise/GenRedOpt.hs @@ -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) diff --git a/src/Futhark/Pass/AD.hs b/src/Futhark/Pass/AD.hs --- a/src/Futhark/Pass/AD.hs +++ b/src/Futhark/Pass/AD.hs @@ -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 = diff --git a/src/Futhark/Pass/ExtractKernels.hs b/src/Futhark/Pass/ExtractKernels.hs --- a/src/Futhark/Pass/ExtractKernels.hs +++ b/src/Futhark/Pass/ExtractKernels.hs @@ -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') + diff --git a/src/Futhark/Tools.hs b/src/Futhark/Tools.hs --- a/src/Futhark/Tools.hs +++ b/src/Futhark/Tools.hs @@ -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 diff --git a/src/Futhark/Transform/FirstOrderTransform.hs b/src/Futhark/Transform/FirstOrderTransform.hs --- a/src/Futhark/Transform/FirstOrderTransform.hs +++ b/src/Futhark/Transform/FirstOrderTransform.hs @@ -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 diff --git a/src/Language/Futhark/Interpreter.hs b/src/Language/Futhark/Interpreter.hs --- a/src/Language/Futhark/Interpreter.hs +++ b/src/Language/Futhark/Interpreter.hs @@ -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 diff --git a/src/Language/Futhark/Interpreter/AD.hs b/src/Language/Futhark/Interpreter/AD.hs --- a/src/Language/Futhark/Interpreter/AD.hs +++ b/src/Language/Futhark/Interpreter/AD.hs @@ -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 diff --git a/src/Language/Futhark/Interpreter/Values.hs b/src/Language/Futhark/Interpreter/Values.hs --- a/src/Language/Futhark/Interpreter/Values.hs +++ b/src/Language/Futhark/Interpreter/Values.hs @@ -152,10 +152,14 @@ pprPrec _ ValueAcc {} = "#" 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 diff --git a/src/Language/Futhark/Prop.hs b/src/Language/Futhark/Prop.hs --- a/src/Language/Futhark/Prop.hs +++ b/src/Language/Futhark/Prop.hs @@ -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 diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs --- a/src/Language/Futhark/TypeChecker/Terms.hs +++ b/src/Language/Futhark/TypeChecker/Terms.hs @@ -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 diff --git a/src/Language/Futhark/TypeChecker/Unify.hs b/src/Language/Futhark/TypeChecker/Unify.hs --- a/src/Language/Futhark/TypeChecker/Unify.hs +++ b/src/Language/Futhark/TypeChecker/Unify.hs @@ -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