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