search

microsoft / knossos-ksc

Compiler with automatic differentiation
Other
45 stars 10 forks source link

Convert "if" to masking in some instances #891

Open awf opened 3 years ago

awf commented 3 years ago

Today, on CPU, vrelu3 is faster using masking than "if" (#877).

Can we automatically convert all "if"s to masking? It would seem not.

So can we indicate which ifs to convert. By analogy with checkpoint and inline,

($inline ((\ x . body) a b c))  -> let x = (a,b,c) in body
($checkpoint e) -> e in most passes, (e, freeVars{e}) in suffwdpass

We could perhaps

(def $if_to_mask_and_on_cpu Float (x : Float) x)
(rule "if_to_mask" ((a : Bool) (b : Float) (c : Float))
    ($if_to_mask_and_on_cpu (if a b c)) ;->
    (add (mul (to_float a) b) (mul (sub 1.0 (to_float a)) c)))