diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,13 +5,32 @@ 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.28] +## [0.25.29] -### Added +### Fixed -### Removed +* AD would in some cases produce code that would cause a compiler + crash (#2228). -### Changed +* Slight error in the definition of the partial derivatives for the + `**` operator could cause NaNs in the interpreter when using + forward-mode AD (#2229). + +* The magical machinery for inferring external API types did not + handle arrays with uniqueness annotations consistently, resulting in + incompatible entry point types being generated, leading to a + compiler crash. (#2231) + +* A simplification rule for array slices would in some cases produce + type-incorrect code. (#2232) + +* A bug in the defunctionaliser could cause a compiler crash in code + that used both higher order functions and size expressions in clever + ways (#2234). + +* Fusion could crash after AD in some circumstances (#2236). + +## [0.25.28] ### Fixed diff --git a/futhark.cabal b/futhark.cabal --- a/futhark.cabal +++ b/futhark.cabal @@ -1,6 +1,6 @@ cabal-version: 2.4 name: futhark -version: 0.25.28 +version: 0.25.29 synopsis: An optimising compiler for a functional, array-oriented language. description: Futhark is a small programming language designed to be compiled to diff --git a/src/Futhark/AD/Derivatives.hs b/src/Futhark/AD/Derivatives.hs --- a/src/Futhark/AD/Derivatives.hs +++ b/src/Futhark/AD/Derivatives.hs @@ -136,7 +136,7 @@ where derivs x y = ( y * (x ** (y - 1)), - condExp (x .<=. 0) 0 ((x ** y) * log x) + (x ** y) * condExp (x .<=. 0) 0 (log x) ) pdBinOp (FMax ft) a b = floatBinOp derivs derivs derivs ft a b diff --git a/src/Futhark/AD/Rev.hs b/src/Futhark/AD/Rev.hs --- a/src/Futhark/AD/Rev.hs +++ b/src/Futhark/AD/Rev.hs @@ -28,6 +28,15 @@ patName (Pat [pe]) = pure $ patElemName pe patName pat = error $ "Expected single-element pattern: " ++ prettyString pat +copyIfArray :: VName -> ADM VName +copyIfArray v = do + v_t <- lookupType v + case v_t of + Array {} -> + letExp (baseString v <> "_copy") . BasicOp $ + Replicate mempty (Var v) + _ -> pure v + -- The vast majority of BasicOps require no special treatment in the -- forward pass and produce one value (and hence one adjoint). We -- deal with that case here. @@ -202,13 +211,7 @@ (_pat_v, pat_adj) <- commonBasicOp pat aux e m returnSweepCode $ do v_adj <- letExp "update_val_adj" $ BasicOp $ Index pat_adj slice - t <- lookupType v_adj - v_adj_copy <- - case t of - Array {} -> - letExp "update_val_adj_copy" . BasicOp $ - Replicate mempty (Var v_adj) - _ -> pure v_adj + v_adj_copy <- copyIfArray v_adj updateSubExpAdj v v_adj_copy zeroes <- letSubExp "update_zero" . zeroExp =<< subExpType v void $ @@ -278,7 +281,9 @@ ses (map (fmap $ diffBody adjs branches_free) cases) (diffBody adjs branches_free defbody) - zipWithM_ insAdj branches_free branches_free_adj + -- See Note [Array Adjoints of Match] + forM_ (zip branches_free branches_free_adj) $ \(v, v_adj) -> + insAdj v =<< copyIfArray v_adj diffStm (Let pat aux (Op soac)) m = vjpSOAC vjpOps pat aux soac m diffStm (Let pat aux loop@Loop {}) m = @@ -439,3 +444,43 @@ -- our current translation rules, they will be dead code. As long as -- we are careful to run dead code elimination after revVJP, we should -- be good. + +-- Note [Array Adjoints of Match] +-- +-- Some unusual, but sadly not completely contrived, contain Match +-- expressions that return multiple arrays, and there the arrays +-- returned by one branch have overlapping aliases with another +-- branch, although in different places. As an example consider this: +-- +-- let (X,Y) = if c +-- then (A, B) +-- else (B, A) +-- +-- Because our aliasing representation cannot express mutually +-- exclusive aliases, we will consider X and Y to be aliased to each +-- other. In practice, this means it is unlikely for X or Y to be +-- consumed, because it would also consume the other (although it's +-- possible for carefully written code). +-- +-- When producing adjoints for this, it will be something like +-- +-- let (X_adj,Y_adj) = if c +-- then (A_adj, B_adj) +-- else (B_adj, A_adj) +-- +-- which completely reflects the primal code. However, while it is +-- unlikely that any consumption takes place for the original primal +-- variables, it is almost guaranteed that X_adj and Y_adj will be +-- consumed (that is the main way we use adjoints after all), and due +-- to the conservative aliasing, when one is consumed, so is the +-- other! To avoid this tragic fate, we are forced to copy any +-- array-typed adjoints returned by a Match. This can be quite costly. +-- However: +-- +-- 1) Futhark has pretty OK copy removal, so maybe it can get rid of +-- these by using information not available to the AD pass. +-- +-- 2) In many cases, arrays will have accumulator adjoints, which are +-- not subject to this problem. +-- +-- Issue #2228 was caused by neglecting to do this. diff --git a/src/Futhark/IR/Prop/Rearrange.hs b/src/Futhark/IR/Prop/Rearrange.hs --- a/src/Futhark/IR/Prop/Rearrange.hs +++ b/src/Futhark/IR/Prop/Rearrange.hs @@ -6,6 +6,7 @@ rearrangeReach, rearrangeCompose, isPermutationOf, + isIdentityPerm, transposeIndex, isMapTranspose, ) @@ -61,6 +62,11 @@ | otherwise = do (xs', v) <- pick (i + 1) xs y pure (x : xs', v) + +-- | Is this an identify permutation? An identity permutation is of +-- the form @[0, 1, ..., k]@. +isIdentityPerm :: [Int] -> Bool +isIdentityPerm perm = perm == [0 .. length perm - 1] -- | If @l@ is an index into the array @a@, then @transposeIndex k n -- l@ is an index to the same element in the array @transposeArray k n diff --git a/src/Futhark/Internalise/Defunctionalise.hs b/src/Futhark/Internalise/Defunctionalise.hs --- a/src/Futhark/Internalise/Defunctionalise.hs +++ b/src/Futhark/Internalise/Defunctionalise.hs @@ -310,8 +310,11 @@ f _ = Nothing sizesToRename :: StaticVal -> S.Set VName -sizesToRename (DynamicFun (_, sv1) sv2) = - sizesToRename sv1 <> sizesToRename sv2 +sizesToRename (DynamicFun (_, sv1) _sv2) = + -- It is intentional that we do not look at sv2 here, as some names + -- that are free in sv2 are actually bound by the parameters in sv1. + -- See #2234. + sizesToRename sv1 sizesToRename IntrinsicSV = mempty sizesToRename HoleSV {} = @@ -882,8 +885,9 @@ globals <- asks fst let bound_sizes = S.fromList dims' <> globals pats' <- instAnySizes pats + let dims'' = dims' ++ unboundSizes bound_sizes pats' - liftValDec fname (RetType [] rettype') (dims' ++ unboundSizes bound_sizes pats') pats' e0 + liftValDec fname (RetType [] rettype') dims'' pats' e0 pure ( Var (qualName fname) diff --git a/src/Futhark/Internalise/Entry.hs b/src/Futhark/Internalise/Entry.hs --- a/src/Futhark/Internalise/Entry.hs +++ b/src/Futhark/Internalise/Entry.hs @@ -185,6 +185,12 @@ entryPointTypeName (I.TypeOpaque v) = v entryPointTypeName (I.TypeTransparent {}) = error "entryPointTypeName: TypeTransparent" +elemTypeExp :: E.TypeExp E.Exp VName -> Maybe (E.TypeExp E.Exp VName) +elemTypeExp (E.TEArray _ te _) = Just te +elemTypeExp (E.TEUnique te _) = elemTypeExp te +elemTypeExp (E.TEParens te _) = elemTypeExp te +elemTypeExp _ = Nothing + entryPointType :: VisibleTypes -> E.EntryType -> @@ -221,17 +227,13 @@ rank = E.shapeRank shape ts' = map (strip rank) ts record_t = E.Scalar (E.Record fs) - record_te = case E.entryAscribed t of - Just (E.TEArray _ te _) -> Just te - _ -> Nothing + record_te = elemTypeExp =<< E.entryAscribed t ept <- snd <$> entryPointType types (E.EntryType record_t record_te) ts' addType desc . I.OpaqueRecordArray rank (entryPointTypeName ept) =<< opaqueRecordArray types rank fs' ts E.Array _ shape et -> do let ts' = map (strip (E.shapeRank shape)) ts - elem_te = case E.entryAscribed t of - Just (E.TEArray _ te _) -> Just te - _ -> Nothing + elem_te = elemTypeExp =<< E.entryAscribed t ept <- snd <$> entryPointType types (E.EntryType (E.Scalar et) elem_te) ts' addType desc . I.OpaqueArray (E.shapeRank shape) (entryPointTypeName ept) $ map valueType ts @@ -265,9 +267,15 @@ E.entryAscribed eret ) of (Just ts, Just (E.TETuple e_ts _)) -> - zipWithM (entryPointType types) (zipWith E.EntryType ts (map Just e_ts)) crets + zipWithM + (entryPointType types) + (zipWith E.EntryType ts (map Just e_ts)) + crets (Just ts, Nothing) -> - zipWithM (entryPointType types) (map (`E.EntryType` Nothing) ts) crets + zipWithM + (entryPointType types) + (map (`E.EntryType` Nothing) ts) + crets _ -> pure <$> entryPointType types eret (concat crets) ) 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 @@ -841,12 +841,12 @@ -- the thing that should be fixed, but it requires fiddling with -- the defunctorisation of size-lifted types. - onExps bound (Var v _ _) e - | Just rexp <- lookup (qualLeaf v) named1 = - onExps mempty (unReplaced rexp) e - | otherwise = - unless (any (`elem` bound) $ freeVarsInExp e) $ - modify (M.insert (qualLeaf v) e) + onExps bound (Var v _ _) e = do + unless (any (`elem` bound) $ freeVarsInExp e) $ + modify (M.insert (qualLeaf v) e) + case lookup (qualLeaf v) named1 of + Just rexp -> onExps mempty (unReplaced rexp) e + Nothing -> pure () onExps _bound e (Var v _ _) | Just rexp <- lookup (qualLeaf v) named2 = onExps mempty e (unReplaced rexp) @@ -997,7 +997,7 @@ substTypesAny (fmap (fmap (second (const mempty))) . (`M.lookup` substs')) params' = map (substPat substStructType) params params'' <- withArgs shape_names $ mapM transformPat params' - exp_naming <- paramGetClean + exp_naming <- get <* put mempty let args = S.fromList $ foldMap patNames params arg_params = map snd exp_naming @@ -1006,7 +1006,7 @@ withParams exp_naming $ withArgs (args <> shape_names) $ hardTransformRetType (applySubst (`M.lookup` substs') rettype) - extNaming <- paramGetClean + extNaming <- get <* put mempty scope <- S.union shape_names <$> askScope' let (rettype'', new_params) = arrowArg scope args arg_params rettype' bind_t' = substTypesAny (`M.lookup` substs') bind_t @@ -1062,11 +1062,6 @@ shape_params = filter (not . isTypeParam) tparams updateExpTypes substs = astMap (mapper substs) - - paramGetClean = do - ret <- get - put mempty - pure ret hardTransformRetType (RetType dims ty) = do ty' <- transformType ty 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 @@ -387,18 +387,19 @@ | wacc2_cons_nms <- namesFromList $ concatMap (\(_, nms, _) -> nms) w2_inps, wacc1_indep_nms <- map getName is1, all (`notNameIn` wacc2_cons_nms) wacc1_indep_nms = do - -- \^ the other safety checks are done inside `tryFuseWithAccs` + -- the other safety checks are done inside `tryFuseWithAccs` lam1' <- fst <$> doFusionInLambda lam1 lam2' <- fst <$> doFusionInLambda lam2 let stm1 = Let pat1 aux1 (WithAcc w1_inps lam1') stm2 = Let pat2 aux2 (WithAcc w2_inps lam2') - mstm <- SF.tryFuseWithAccs infusible stm1 stm2 + mstm <- sequence $ SF.tryFuseWithAccs infusible stm1 stm2 case mstm of + Nothing -> pure Nothing Just (Let pat aux (WithAcc w_inps wlam)) -> do (wlam', success) <- doFusionInLambda wlam let new_stm = Let pat aux (WithAcc w_inps wlam') if success then fusedSomething (StmNode new_stm) else pure Nothing - _ -> error "Illegal result of tryFuseWithAccs called from vFuseNodeT." + Just _ -> error "Illegal result of tryFuseWithAccs called from vFuseNodeT." -- vFuseNodeT _ _ _ _ = pure Nothing 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 @@ -400,7 +400,7 @@ [VName] -> Stm SOACS -> Stm SOACS -> - m (Maybe (Stm SOACS)) + Maybe (m (Stm SOACS)) tryFuseWithAccs infusible (Let pat1 aux1 (WithAcc w_inps1 lam1)) @@ -424,7 +424,7 @@ all (`notElem` infusible) bs, -- safety 3: cs <- namesFromList $ concatMap ((\(_, xs, _) -> xs) . accTup2) acc_tup2, - all ((`notNameIn` cs) . patElemName . fst) other_pr1 = do + all ((`notNameIn` cs) . patElemName . fst) other_pr1 = Just $ do let getCertPairs (t1, t2) = (paramName (accTup3 t2), paramName (accTup3 t1)) tab_certs = M.fromList $ map getCertPairs tup_common lam2_bdy' = substituteNames tab_certs (lambdaBody lam2) @@ -472,8 +472,7 @@ ++ map fst (other_pr1 ++ other_pr2) res_w_inps = map (accTup2 . fst) tup_common ++ map accTup2 (acc_tup1' ++ acc_tup2') res_w_inps' <- mapM renameLamInWAccInp res_w_inps - let stm_res = Let (Pat res_pat) (aux1 <> aux2) $ WithAcc res_w_inps' res_lam' - pure $ Just stm_res + pure $ Let (Pat res_pat) (aux1 <> aux2) $ WithAcc res_w_inps' res_lam' where -- local helpers: @@ -535,7 +534,7 @@ renameLamInWAccInp winp = pure winp -- tryFuseWithAccs _ _ _ = - pure Nothing + Nothing ------------------------------- --- simple helper functions --- diff --git a/src/Futhark/Optimise/Simplify/Rules/Index.hs b/src/Futhark/Optimise/Simplify/Rules/Index.hs --- a/src/Futhark/Optimise/Simplify/Rules/Index.hs +++ b/src/Futhark/Optimise/Simplify/Rules/Index.hs @@ -8,6 +8,8 @@ where import Control.Monad (guard) +import Data.Bifunctor (first) +import Data.List qualified as L import Data.List.NonEmpty (NonEmpty (..)) import Data.Maybe import Futhark.Analysis.PrimExp.Convert @@ -33,6 +35,8 @@ | SubExpResult Certs SubExp -- Fake expressions that we can recognise. +-- +-- See Note [Simplifying a Slice]. fakeIndices :: [TPrimExp Int64 VName] fakeIndices = map f [0 :: Int ..] where @@ -61,6 +65,7 @@ worthInlining e, all (`ST.elem` vtable) (unCerts cs) -> Just $ SubExpResult cs <$> toSubExp "index_primexp" e + -- For the two cases below, see Note [Simplifying a Slice]. | Just inds' <- sliceIndices (Slice inds), Just (ST.IndexedArray cs arr inds'') <- ST.index idd inds' vtable, all (worthInlining . untyped) inds'', @@ -76,17 +81,28 @@ all (`ST.elem` vtable) (unCerts cs), not consuming, not $ consumed arr, - Just inds''' <- mapM okIdx inds'' -> do - Just $ IndexResult cs arr . Slice <$> sequence inds''' + Just (ordering, inds''') <- first concat . unzip <$> mapM okIdx inds'', + Just perm <- L.sort ordering `isPermutationOf` ordering -> + if isIdentityPerm perm + then Just $ IndexResult cs arr . Slice <$> sequence inds''' + else Just $ do + arr_sliced <- + certifying cs $ + letExp (baseString arr <> "_sliced") . BasicOp . Index arr . Slice + =<< sequence inds''' + arr_sliced_tr <- + letSubExp (baseString arr_sliced <> "_tr") $ + BasicOp (Rearrange perm arr_sliced) + pure $ SubExpResult mempty arr_sliced_tr where - matches = zip fakeIndices $ sliceDims $ Slice inds + matches = zip fakeIndices $ zip [0 :: Int ..] $ sliceDims $ Slice inds okIdx i = case lookup i matches of - Just w -> - Just $ pure $ DimSlice (constant (0 :: Int64)) w (constant (1 :: Int64)) + Just (j, w) -> + Just ([j], pure $ DimSlice (constant (0 :: Int64)) w (constant (1 :: Int64))) Nothing -> do guard $ not $ any ((`namesIntersect` freeIn i) . freeIn . fst) matches - Just $ DimFix <$> toSubExp "index_primexp" i + Just ([], DimFix <$> toSubExp "index_primexp" i) Nothing -> Nothing Just (SubExp (Var v), cs) -> Just $ pure $ IndexResult cs v $ Slice inds @@ -258,3 +274,48 @@ True | otherwise = False + +-- Note [Simplifying a Slice] +-- +-- The 'indexOp' simplification machinery permits simplification of +-- full indexing (i.e., where every component of the Slice is a +-- DimFix). We use this in a creative way to also simplify slices. +-- For example, for a slice +-- +-- A[i:j:+n] +-- +-- we synthesize a uniquely recognizable index expression "ie" (see +-- 'fakeIndices'), we use `indexOp` to simplify the full indexing +-- +-- A[ie] +-- +-- and if that produces a simplification result +-- +-- B[ie] +-- +-- then we can replace the "ie" DimFix with our original slice and +-- produce +-- +-- B[i:j:+n] +-- +-- While the above case is trivial, this is useful for cases that +-- intermix indexing and slicing. We must be careful, however: If we +-- have an original expression +-- +-- A[i0:j0:+n0,i1:j1:+n1] +-- +-- for which we then synthesize the expression +-- +-- A[ie0, ie1] +-- +-- then if we receive back the result +-- +-- B[ie1, ie0] +-- +-- we cannot just replace the indexes with the original slices, as +-- that would change the shape (and semantics) of the result: +-- +-- B[i1:j1:+n1,i0:j0:+n0] +-- +-- In such cases we must actually insert a Rearrange operation to move +-- the dimensions of the result appropriately. 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 @@ -35,6 +35,7 @@ import Data.Array import Data.Bifunctor import Data.Bitraversable +import Data.Either (fromRight) import Data.List ( find, foldl', @@ -1351,8 +1352,9 @@ adToPrim v = putV $ AD.primitive v - adBinOp op x y = AD.doOp op [x, y] - adUnOp op x = AD.doOp op [x] + adBinOp op x y = + either (const Nothing) Just $ AD.doOp op [x, y] + adUnOp op x = either (const Nothing) Just $ AD.doOp op [x] fun1 f = TermValue Nothing $ ValueFun $ \x -> f x @@ -1421,7 +1423,7 @@ _ | Just x' <- getAD x, Just y' <- getAD y, - Just z <- msum $ map (`bopDefAD'` (x', y')) fs -> do + Just z <- msum $ map (`bopDefAD` (x', y')) fs -> do breakOnNaN [adToPrim x', adToPrim y'] $ adToPrim z pure $ putAD z _ -> @@ -1436,7 +1438,7 @@ x' <- valf x y' <- valf y retf =<< op x' y' - bopDefAD' (_, _, _, dop) (x, y) = dop x y + bopDefAD (_, _, _, dop) (x, y) = dop x y unopDef fs = fun1 $ \x -> case x of @@ -1471,7 +1473,7 @@ Just [x, y] | Just x' <- getAD x, Just y' <- getAD y, - Just z <- AD.doOp op [x', y'] -> do + Right z <- AD.doOp op [x', y'] -> do breakOnNaN [adToPrim x', adToPrim y'] $ adToPrim z pure $ putAD z _ -> @@ -2029,7 +2031,7 @@ -- TODO: Perhaps this could be fully abstracted by AD? -- Making addFor private would be nice.. add x y = - fromMaybe (error "jvp: illtyped add") $ + fromRight (error "jvp: illtyped add") $ AD.doOp (AD.OpBin $ AD.addFor $ P.primValueType $ AD.primitive x) [x, y] def "jvp2" = Just $ -- TODO: This could be much better. Currently, it is very inefficient 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 @@ -15,7 +15,7 @@ where import Control.Monad (foldM, zipWithM) -import Data.Either (isRight) +import Data.Either (fromRight, isRight) import Data.List (find, foldl') import Data.Map qualified as M import Data.Maybe (fromMaybe) @@ -70,12 +70,14 @@ mulFor Bool = LogAnd mulFor t = error $ "mulFor: " ++ show t +type Depth = Int + -- Types and utility functions-- -- When taking the partial derivative of a function, we -- must differentiate between the values which are kept -- constant, and those which are not data ADValue - = Variable Int ADVariable + = Variable Depth ADVariable | Constant PrimValue deriving (Show) @@ -88,7 +90,7 @@ | JVP JVPValue deriving (Show) -depth :: ADValue -> Int +depth :: ADValue -> Depth depth (Variable d _) = d depth (Constant _) = 0 @@ -97,6 +99,12 @@ primal (Variable _ (JVP (JVPValue v _))) = primal v primal (Constant v) = Constant v +primalFor :: Depth -> ADValue -> ADValue +primalFor cur v@(Variable tag _) | cur /= tag = v +primalFor _ (Variable _ (VJP (VJPValue t))) = tapePrimal t +primalFor cur (Variable _ (JVP (JVPValue v _))) = primalFor cur v +primalFor _ (Constant v) = Constant v + primitive :: ADValue -> PrimValue primitive (Variable _ v) = varPrimal v primitive (Constant v) = v @@ -106,9 +114,11 @@ varPrimal (JVP (JVPValue v _)) = primitive $ primal v -- Evaluates a PrimExp using doOp -evalPrimExp :: M.Map VName ADValue -> PrimExp VName -> Maybe ADValue -evalPrimExp m (LeafExp n _) = M.lookup n m -evalPrimExp _ (ValueExp pv) = Just $ Constant pv +evalPrimExp :: M.Map VName ADValue -> PrimExp VName -> Either String ADValue +evalPrimExp m (LeafExp n _) = + maybe (Left $ "Unknown variable " <> show n) Right $ M.lookup n m +evalPrimExp _ (ValueExp pv) = + Right $ Constant pv evalPrimExp m (BinOpExp op x y) = do x' <- evalPrimExp m x y' <- evalPrimExp m y @@ -141,23 +151,25 @@ -- This function performs a mathematical operation on a -- list of operands, performing automatic differentiation -- if one or more operands is a Variable (of depth > 0) -doOp :: Op -> [ADValue] -> Maybe ADValue +doOp :: Op -> [ADValue] -> Either String ADValue doOp op o | not $ opTypeMatch op (map primValueType pv) = -- This function may be called with arguments of invalid types, -- because it is used as part of an overloaded operator. - Nothing + Left $ unwords ["invalid types for op", show op, "and operands", show o] | otherwise = do let dep = case op of OpCmp _ -> 0 -- AD is not well-defined for comparason operations -- There are no derivatives for those written in -- PrimExp (check lookupPDs) _ -> maximum (map depth o) - if dep == 0 then constCase else nonconstCase dep + if dep == 0 + then maybe (Left "failed to evaluate const") Right constCase + else nonconstCase dep where pv = map primitive o - divideDepths :: Int -> ADValue -> Either ADValue ADVariable + divideDepths :: Depth -> ADValue -> Either ADValue ADVariable divideDepths _ v@(Constant {}) = Left v divideDepths d v@(Variable d' v') = if d' < d then Left v else Right v' @@ -198,7 +210,7 @@ -- have to perform the necessary steps for AD -- First, we calculate the value for the previous depth - let oprev = map primal o + let oprev = map (primalFor dep) o vprev <- doOp op oprev -- Then we separate the values of the maximum depth from @@ -209,28 +221,31 @@ -- Finally, we perform the necessary steps for the given -- type of AD Just (Right (VJP {})) -> - Just . Variable dep . VJP . VJPValue $ vjpHandleOp op (map extractVJP o') vprev + Right . Variable dep . VJP . VJPValue $ + vjpHandleOp op (map extractVJP o') vprev Just (Right (JVP {})) -> - Variable dep . JVP . JVPValue vprev <$> jvpHandleFn op (map extractJVP o') + Variable dep . JVP . JVPValue vprev + <$> jvpHandleOp op (map extractJVP o') _ -> -- Since the maximum depth is non-zero, there must be at -- least one variable of depth > 0 error "find isRight" calculatePDs :: Op -> [ADValue] -> [ADValue] -calculatePDs op p = +calculatePDs op args = -- Create a unique VName for each operand - let n = map (\i -> VName (nameFromString $ "x" ++ show i) i) [1 .. length p] + let n = map (\i -> VName (nameFromString $ "x" ++ show i) i) [1 .. length args] -- Put the operands in the environment - m = M.fromList $ zip n p + m = M.fromList $ zip n args -- Look up, and calculate the partial derivative -- of the operation with respect to each operand pde = fromMaybe (error "lookupPDs failed") $ lookupPDs op $ - map (`LeafExp` opReturnType op) n - in map (fromMaybe (error "evalPrimExp failed") . evalPrimExp m) pde + zipWith (\v val -> LeafExp v $ primValueType $ primitive val) n args + res = map (either (error . ("evalPrimExp failed: " <>)) id . evalPrimExp m) pde + in res -- VJP / Reverse mode automatic differentiation-- -- In reverse mode AD, the entire computation @@ -244,7 +259,7 @@ data Tape = -- | This represents a variable. Each variable is given a unique ID, -- and has an initial value - TapeID Int ADValue + TapeID Depth ADValue | -- | This represents a constant. TapeConst ADValue | -- | This represents the application of a mathematical operation. @@ -279,7 +294,7 @@ let s'' = case op of OpConv op' -> -- In case of type conversion, simply convert the sensitivity - [ fromMaybe (error "deriveTape: doOp failed") $ + [ fromRight (error "deriveTape: doOp failed") $ doOp (OpConv $ flipConvOp op') [s] ] _ -> @@ -291,10 +306,10 @@ foldl' (M.unionWith add) M.empty pd where add x y = - fromMaybe (error "deriveTape: add failed") $ + fromRight (error "deriveTape: add failed") $ doOp (OpBin $ addFor $ opReturnType op) [x, y] mul x y = - fromMaybe (error "deriveTape: mul failed") $ + fromRight (error "deriveTape: mul failed") $ doOp (OpBin $ mulFor $ opReturnType op) [x, y] -- JVP / Forward mode automatic differentiation-- @@ -304,27 +319,26 @@ data JVPValue = JVPValue ADValue ADValue deriving (Show) --- | This calculates the derivative part of the JVPValue resulting +-- | This calculates the tangent part of the JVPValue resulting -- from the application of a mathematical operation on one or more -- JVPValues. -jvpHandleFn :: Op -> [Either ADValue JVPValue] -> Maybe ADValue -jvpHandleFn op p = do +jvpHandleOp :: Op -> [Either ADValue JVPValue] -> Either String ADValue +jvpHandleOp op p = do case op of OpConv _ -> -- In case of type conversion, simply convert - -- the old derivative - doOp op [derivative $ head p] + -- the old tangent + doOp op [tangent $ head p] _ -> do - -- Calculate the new derivative using the chain - -- rule + -- Calculate the new tangent using the chain rule let pds = calculatePDs op $ map primal' p - vs <- zipWithM mul pds $ map derivative p - foldM add (Constant $ blankPrimValue $ opReturnType op) vs + vs <- zipWithM mul pds $ map tangent p + foldM add (Constant $ blankPrimValue op_t) vs where + op_t = opReturnType op primal' (Left v) = v primal' (Right (JVPValue v _)) = v - derivative (Left v) = Constant $ blankPrimValue $ primValueType $ primitive v - derivative (Right (JVPValue _ d)) = d - + tangent (Left _) = Constant $ blankPrimValue $ opReturnType op + tangent (Right (JVPValue _ d)) = d add x y = doOp (OpBin $ addFor $ opReturnType op) [x, y] mul x y = doOp (OpBin $ mulFor $ opReturnType op) [x, y] 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 @@ -200,7 +200,8 @@ valueAccum :: (a -> Value m -> (a, Value m)) -> a -> Value m -> (a, Value m) valueAccum f i v@(ValuePrim {}) = f i v valueAccum f i v@(ValueAD {}) = f i v -valueAccum f i (ValueRecord m) = second ValueRecord $ M.mapAccum (valueAccum f) i m +valueAccum f i (ValueRecord m) = + second ValueRecord $ M.mapAccum (valueAccum f) i m valueAccum f i (ValueArray s a) = do -- TODO: This could probably be better -- Transform into a map diff --git a/src/Language/Futhark/TypeChecker/Modules.hs b/src/Language/Futhark/TypeChecker/Modules.hs --- a/src/Language/Futhark/TypeChecker/Modules.hs +++ b/src/Language/Futhark/TypeChecker/Modules.hs @@ -343,12 +343,12 @@ missingType :: (Pretty a) => Loc -> a -> Either TypeError b missingType loc name = Left . TypeError loc mempty $ - "Module does not define a type named" <+> pretty name <> "." + "Module does not define a type named" <+> dquotes (pretty name) <> "." missingVal :: (Pretty a) => Loc -> a -> Either TypeError b missingVal loc name = Left . TypeError loc mempty $ - "Module does not define a value named" <+> pretty name <> "." + "Module does not define a value named" <+> dquotes (pretty name) <> "." topLevelSize :: Loc -> VName -> Either TypeError b topLevelSize loc name = @@ -358,7 +358,7 @@ missingMod :: (Pretty a) => Loc -> a -> Either TypeError b missingMod loc name = Left . TypeError loc mempty $ - "Module does not define a module named" <+> pretty name <> "." + "Module does not define a module named" <+> dquotes (pretty name) <> "." mismatchedType :: Loc ->