futhark-0.22.3: src/Futhark/AD/Rev/Scan.hs
module Futhark.AD.Rev.Scan (diffScan) where
import Control.Monad
import Futhark.AD.Rev.Monad
import Futhark.Analysis.PrimExp.Convert
import Futhark.Builder
import Futhark.IR.SOACS
import Futhark.Tools
import Futhark.Transform.Rename
import Futhark.Util (pairs, unpairs)
data FirstOrSecond = WrtFirst | WrtSecond
-- computes `d(x op y)/dx` or d(x op y)/dy
mkScanAdjointLam :: VjpOps -> Lambda SOACS -> FirstOrSecond -> ADM (Lambda SOACS)
mkScanAdjointLam ops lam0 which = do
let len = length $ lambdaReturnType lam0
lam <- renameLambda lam0
let p2diff =
case which of
WrtFirst -> take len $ lambdaParams lam
WrtSecond -> drop len $ lambdaParams lam
p_adjs <- mapM unitAdjOfType (lambdaReturnType lam)
vjpLambda ops p_adjs (map paramName p2diff) lam
-- Should generate something like:
-- `\ j -> let i = n - 1 - j
-- if i < n-1 then ( ys_adj[i], df2dx ys[i] xs[i+1]) else (0,1) )`
-- where `ys` is the result of scan
-- `xs` is the input of scan
-- `ys_adj` is the known adjoint of ys
-- `j` draw values from `iota n`
mkScanFusedMapLam :: VjpOps -> SubExp -> Lambda SOACS -> [VName] -> [VName] -> [VName] -> ADM (Lambda SOACS)
mkScanFusedMapLam ops w scn_lam xs ys ys_adj = do
lam <- mkScanAdjointLam ops scn_lam WrtFirst
ys_ts <- mapM lookupType ys
par_i <- newParam "i" $ Prim int64
let i = paramName par_i
mkLambda [par_i] $
fmap varsRes . letTupExp "x"
=<< eIf
(toExp $ le64 i .==. 0)
( buildBody_ $ do
zs <- mapM (letSubExp "ct_zero" . zeroExp . rowType) ys_ts
os <- mapM (letSubExp "ct_one" . oneExp . rowType) ys_ts
pure $ subExpsRes $ unpairs $ zip zs os
)
( buildBody_ $ do
j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))
j1 <- letSubExp "j1" =<< toExp (pe64 w - le64 i)
let index idx arr t = BasicOp $ Index arr $ fullSlice t [DimFix idx]
y_s <- forM (zip ys_adj ys_ts) $ \(y_, t) ->
letSubExp (baseString y_ ++ "_j") $ index j y_ t
lam_rs <-
eLambda lam . map pure $
zipWith (index j) ys ys_ts ++ zipWith (index j1) xs ys_ts
pure $ unpairs $ zip (subExpsRes y_s) lam_rs
)
-- \(a1, b1) (a2, b2) -> (a2 + b2 * a1, b1 * b2)
mkScanLinFunO :: Type -> ADM (Scan SOACS)
mkScanLinFunO t = do
let pt = elemType t
zero <- letSubExp "zeros" $ zeroExp t
one <- letSubExp "ones" $ oneExp t
tmp <- mapM newVName ["a1", "b1", "a2", "b2"]
let [a1, b1, a2, b2] = tmp
pet = primExpFromSubExp pt . Var
lam <- mkLambda (map (\v -> Param mempty v t) [a1, b1, a2, b2]) . fmap varsRes $
tabNest (arrayRank t) [a1, b1, a2, b2] $ \_ [a1', b1', a2', b2'] -> do
x <- letExp "x" <=< toExp $ pet a2' ~+~ pet b2' ~*~ pet a1'
y <- letExp "y" <=< toExp $ pet b1' ~*~ pet b2'
pure [x, y]
pure $ Scan lam [zero, one]
-- build the map following the scan with linear-function-composition:
-- for each (ds,cs) length-n array results of the scan, combine them as:
-- `let rs = map2 (\ d_i c_i -> d_i + c_i * y_adj[n-1]) d c |> reverse`
-- but insert explicit indexing to reverse inside the map.
mkScan2ndMaps :: SubExp -> (Type, VName, (VName, VName)) -> ADM VName
mkScan2ndMaps w (arr_tp, y_adj, (ds, cs)) = do
y_adj_last <- letExp (baseString y_adj <> "_last") =<< eLast y_adj
par_i <- newParam "i" $ Prim int64
lam <- mkLambda [par_i] $ do
let i = paramName par_i
j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))
dj <- letExp (baseString ds ++ "_dj") $ BasicOp $ Index ds $ fullSlice arr_tp [DimFix j]
cj <- letExp (baseString cs ++ "_cj") $ BasicOp $ Index cs $ fullSlice arr_tp [DimFix j]
let pet = primExpFromSubExp (elemType arr_tp) . Var
fmap varsRes . tabNest (arrayRank (rowType arr_tp)) [y_adj_last, dj, cj] $ \_ [y_adj_last', dj', cj'] ->
letTupExp "res" <=< toExp $ pet dj' ~+~ pet cj' ~*~ pet y_adj_last'
iota <- letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64
letExp "after_scan" $ Op (Screma w [iota] (ScremaForm [] [] lam))
-- perform the final map, which is fusable with the maps obtained from `mkScan2ndMaps`
-- let xs_contribs =
-- map3 (\ i a r -> if i==0 then r else (df2dy (ys[i-1]) a) \bar{*} r)
-- (iota n) xs rs
mkScanFinalMap :: VjpOps -> SubExp -> Lambda SOACS -> [VName] -> [VName] -> [VName] -> ADM [VName]
mkScanFinalMap ops w scan_lam xs ys rs = do
let eltps = lambdaReturnType scan_lam
lam <- mkScanAdjointLam ops scan_lam WrtSecond
par_i <- newParam "i" $ Prim int64
let i = paramName par_i
par_x <- mapM (\(x, t) -> newParam (baseString x ++ "_par_x") t) $ zip xs eltps
par_r <- mapM (\(r, t) -> newParam (baseString r ++ "_par_r") t) $ zip rs eltps
map_lam <-
mkLambda (par_i : par_x ++ par_r) $
fmap varsRes . letTupExp "scan_contribs"
=<< eIf
(toExp $ le64 i .==. 0)
(resultBodyM $ map (Var . paramName) par_r)
( buildBody_ $ do
im1 <- letSubExp "im1" =<< toExp (le64 i - 1)
ys_im1 <- forM ys $ \y ->
letSubExp (baseString y <> "_im1") =<< eIndex y (eSubExp im1)
lam_res <-
mapM (letExp "const" . BasicOp . SubExp . resSubExp)
=<< eLambda lam (map eSubExp $ ys_im1 ++ map (Var . paramName) par_x)
fmap (varsRes . mconcat) . forM (zip3 lam_res (map paramName par_r) eltps) $
\(lam_r, r, eltp) -> do
let pet = primExpFromSubExp (elemType eltp) . Var
tabNest (arrayRank eltp) [lam_r, r] $ \_ [lam_r', r'] ->
letTupExp "res" <=< toExp $ pet lam_r' ~*~ pet r'
)
iota <- letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64
letTupExp "scan_contribs" $ Op (Screma w (iota : xs ++ rs) (ScremaForm [] [] map_lam))
diffScan :: VjpOps -> [VName] -> SubExp -> [VName] -> Scan SOACS -> ADM ()
diffScan ops ys w as scan = do
ys_adj <- mapM lookupAdjVal ys
as_ts <- mapM lookupType as
map1_lam <- mkScanFusedMapLam ops w (scanLambda scan) as ys ys_adj
scans_lin_fun_o <- mapM mkScanLinFunO $ lambdaReturnType $ scanLambda scan
iota <-
letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64
r_scan <-
letTupExp "adj_ctrb_scan" $
Op (Screma w [iota] (ScremaForm scans_lin_fun_o [] map1_lam))
red_nms <- mapM (mkScan2ndMaps w) (zip3 as_ts ys_adj (pairs r_scan))
as_contribs <- mkScanFinalMap ops w (scanLambda scan) as ys red_nms
zipWithM_ updateAdj as as_contribs