prelude.dx

'## Dex prelude

'Runs before every Dex program unless an alternative is provided with `--prelude`.

'Wrappers around primitives

Int  = %Int
Real = %Real
Unit = %UnitType
Type = %TyKind
Effects = %EffKind

def (&) (a:Type) (b:Type) : Type = %PairType a b
def (,) (x:a) (y:b) : (a & b) = %pair x y
def fst (p: (a & b)) : a = %fst p
def snd (p: (a & b)) : b = %snd p

def idiv (x:Int) (y:Int) : Int = %idiv x y
def rem  (x:Int) (y:Int) : Int = %irem x y
def ipow (x:Int) (y:Int) : Int = %ipow x y

def neg  (x:Real)          : Real = %fneg x
def fdiv (x:Real) (y:Real) : Real = %fdiv x y

data Add a:Type =
  MkAdd (a->a->a) (a->a->a) a  -- add, sub, zero

def (+)  (d:Add a) ?=> : a -> a -> a = case d of MkAdd add _   _    -> add
def (-)  (d:Add a) ?=> : a -> a -> a = case d of MkAdd _   sub _    -> sub
def zero (d:Add a) ?=> : a           = case d of MkAdd _   _   zero -> zero

@instance realAdd : Add Real = MkAdd (\x y. %fadd x y) (\x y. %fsub x y) 0.0
@instance intAdd  : Add Int  = MkAdd (\x y. %iadd x y) (\x y. %isub x y) 0
@instance unitAdd : Add Unit = MkAdd (\x y. ())        (\x y. ())        ()

@instance tabAdd : Add a ?=> Add (n=>a) =
  (MkAdd ( \xs ys. for i. xs.i + ys.i )
         ( \xs ys. for i. xs.i - ys.i )
         ( for _. zero ))

data Mul a:Type = MkMul (a->a->a) a  -- multiply, one

def (*) (d:Mul a) ?=> : a -> a -> a = case d of MkMul mul _   -> mul
def one (d:Mul a) ?=> : a           = case d of MkMul _   one -> one

@instance realMul : Mul Real = MkMul (\x y. %fmul x y) 1.0
@instance intMul  : Mul Int  = MkMul (\x y. %imul x y) 1
@instance unitMul : Mul Unit = MkMul (\x y. ()) ()

data VSpace a:Type = MkVSpace (Add a) (Real -> a -> a)

@superclass
def addFromVSpace (d:VSpace a) : Add a = case d of MkVSpace addDict _ -> addDict

flip : (a -> b -> c) -> (b -> a -> c) = \f x y. f y x
uncurry : (a -> b -> c) -> (a & b) -> c = \f (x,y). f x y
const : a -> b -> a = \x _. x

def (.*)  (d:VSpace a) ?=> : Real -> a -> a = case d of MkVSpace _ scale -> scale
(*.)  : VSpace a ?=> a -> Real -> a = flip (.*)
def (/) (_:VSpace a) ?=> (v:a) (s:Real) : a = (fdiv 1.0 s) .* v

@instance realVS : VSpace Real = MkVSpace realAdd (*)
@instance tabVS  : VSpace a ?=> VSpace (n=>a) = MkVSpace tabAdd \s xs. for i. s .* xs.i
@instance unitVS : VSpace Unit = MkVSpace unitAdd \s u. ()


InternalBool = %Bool

-- XXX: False must come before True for the later coercions to be correct.
data Bool =
  False
  True

def unsafeCoerce (b:Type) (x:a) : b = %unsafeCoerce b x

def (&&) (x:Bool) (y:Bool) : Bool =
  x' = unsafeCoerce InternalBool x
  y' = unsafeCoerce InternalBool y
  unsafeCoerce _ $ %and x' y'

def (||) (x:Bool) (y:Bool) : Bool =
  x' = unsafeCoerce InternalBool x
  y' = unsafeCoerce InternalBool y
  unsafeCoerce _ $ %or x' y'

def not  (x:Bool) : Bool =
  x' = unsafeCoerce InternalBool x
  unsafeCoerce _ $ %not x'

'Sum types

data Maybe a:Type =
  Nothing
  Just a

def isNothing (x:Maybe a) : Bool = case x of
  Nothing -> True
  Just _ -> False

data (|) a:Type b:Type =
  Left  a
  Right b

def select (p:Bool) (x:a) (y:a) : a = case p of
  True  -> x
  False -> y

def b2i (x:Bool) : Int  =
  x' = unsafeCoerce InternalBool x
  %booltoint x'

def i2r (x:Int ) : Real = %inttoreal x
def b2r (x:Bool) : Real = i2r (b2i x)
def todo (a:Type) ?-> : a = %throwError a

'Effects

def Ref (r:Type) (a:Type) : Type = %Ref r a
def get  (ref:Ref h s)       : {State h} s    = %get  ref
def (:=) (ref:Ref h s) (x:s) : {State h} Unit = %put  ref x
def ask  (ref:Ref h r)       : {Read  h} r    = %ask  ref
def (+=) (ref:Ref h w) (x:w) : {Accum h} Unit = %tell ref x
def (!)  (ref:Ref h (n=>a)) (i:n) : Ref h a = %indexRef ref i
def fstRef (ref: Ref h (a & b)) : Ref h a = %fstRef ref
def sndRef (ref: Ref h (a & b)) : Ref h b = %sndRef ref

def withReader
      (eff:Effects) ?-> (a:Type) ?-> (r:Type) ?->
      (init:r) (action: (h:Type ?-> Ref h r -> {Read h|eff} a))
      : {|eff} a =
    def explicitAction (h':Type) (ref:Ref h' r) : {Read h'|eff} a = action ref
    %runReader init explicitAction

def withAccum
      (eff:Effects) ?-> (a:Type) ?-> (w:Type) ?->
      (action: (h:Type ?-> Ref h w -> {Accum h|eff} a))
      : {|eff} (a & w) =
    def explicitAction (h':Type) (ref:Ref h' w) : {Accum h'|eff} a = action ref
    %runWriter explicitAction

def withState
      (eff:Effects) ?-> (a:Type) ?-> (s:Type) ?->
      (init:s)
      (action: (h:Type ?-> Ref h s -> {State h |eff} a))
      : {|eff} (a & s) =
    def explicitAction (h':Type) (ref:Ref h' s) : {State h'|eff} a = action ref
    %runState init explicitAction

'Type classes

data Eq  a:Type = MkEq  (a -> a -> Bool)
data Ord a:Type = MkOrd (Eq a) (a -> a -> Bool) (a -> a -> Bool)  -- eq, gt, lt

@superclass
def eqFromOrd (d:Ord a) : Eq a = case d of MkOrd eq _ _ -> eq

def (==) (d:Eq a) ?=> (x:a) (y:a) : Bool = case d of MkEq eq -> eq x y
def (/=) (d:Eq a) ?=> (x:a) (y:a) : Bool = not $ x == y

def (>)  (d:Ord a) ?=> (x:a) (y:a) : Bool = case d of MkOrd _ gt _  -> gt x y
def (<)  (d:Ord a) ?=> (x:a) (y:a) : Bool = case d of MkOrd _ _  lt -> lt x y
def (<=) (d:Ord a) ?=> (x:a) (y:a) : Bool = x<y || x==y
def (>=) (d:Ord a) ?=> (x:a) (y:a) : Bool = x>y || x==y

@instance intEq  : Eq Int  = MkEq \x y. unsafeCoerce _ $ %ieq x y
@instance realEq : Eq Real = MkEq \x y. unsafeCoerce _ $ %feq x y
@instance unitEq : Eq Unit = MkEq \x y. True

@instance intOrd  : Ord Int  = (MkOrd intEq  (\x y. unsafeCoerce _ $ %igt x y)
                                             (\x y. unsafeCoerce _ $ %ilt x y))
@instance realOrd : Ord Real = (MkOrd realEq (\x y. unsafeCoerce _ $ %fgt x y)
                                             (\x y. unsafeCoerce _ $ %flt x y))
@instance unitOrd : Ord Unit = MkOrd unitEq (\x y. False) (\x y. False)

@instance
def pairEq (eqA: Eq a)?=> (eqB: Eq b)?=> : Eq (a & b) = MkEq $
  \(x1,x2) (y1,y2). x1 == y1 && x2 == y2

@instance
def pairOrd (ordA: Ord a)?=> (ordB: Ord b)?=> : Ord (a & b) =
  pairGt = \(x1,x2) (y1,y2). x1 > y1 || (x1 == y1 && x2 > y2)
  pairLt = \(x1,x2) (y1,y2). x1 < y1 || (x1 == y1 && x2 < y2)
  MkOrd pairEq pairGt pairLt

-- TODO: accumulate using the True/&& monoid
@instance
def tabEq (n:Type) ?-> (eqA: Eq a) ?=> : Eq (n=>a) = MkEq $
  \xs ys.
    numDifferent : Real =
      snd $ withAccum \ref. for i.
        ref += (i2r (b2i (xs.i /= ys.i)))
    numDifferent == 0.0

'Wrappers around C library functions

def exp (x:Real) : Real = %exp x
def exp2 (x:Real) : Real = %exp2 x
def log (x:Real) : Real = %log x
def log2 (x:Real) : Real = %log2 x
def log10 (x:Real) : Real = %log10 x

def sin (x:Real) : Real = %sin x
def cos (x:Real) : Real = %cos x
def tan (x:Real) : Real = %tan x

def floor (x:Real) : Int = %floor x
def ceil (x:Real) : Int = %ceil x
def round (x:Real) : Int = %round x

def sqrt (x:Real) : Real = %sqrt x
def pow (x:Real) (y:Real) : Real = %fpow x y

def lgamma (x:Real) : Real = %ffi lgamma Real x
def log1p  (x:Real) : Real = %ffi log1p  Real x
def lbeta (x:Real) (y:Real) : Real = lgamma x + lgamma y - lgamma (x + y)

'Working with index sets

def Range (low:Int) (high:Int) : Type = %IntRange low high
def Fin (n:Int) : Type = Range 0 n
def ordinal (i:a) : Int = %asint i
def size (n:Type) : Int = %idxSetSize n
def fromOrdinal (n:Type) (i:Int) : n = %asidx n i
def asidx (n:Type) ?-> (i:Int) : n = fromOrdinal n i
def (@) (i:Int) (n:Type) : n = fromOrdinal n i
def ixadd (n:Type) ?-> (i:n) (x:Int) : n = fromOrdinal n $ ordinal i + x
def ixsub (n:Type) ?-> (i:n) (x:Int) : n = fromOrdinal n $ ordinal i - x
def iota (n:Type) : n=>Int = for i. ordinal i

-- TODO: we want Eq and Ord for all index sets, not just `Fin n`
@instance
def finEq (n:Int) ?-> : Eq (Fin n) = MkEq \x y. ordinal x == ordinal y

@instance
def finOrd (n:Int) ?-> : Ord (Fin n) =
  MkOrd finEq (\x y. ordinal x > ordinal y) (\x y. ordinal x < ordinal y)

'Misc

pi : Real = 3.141592653589793

def id (x:a) : a = x
def dup (x:a) : (a & a) = (x, x)
-- TODO: unpack pair in args once we fix the bug
def swap (p:(a&b)) : (b&a) = (snd p, fst p)
def map (f:a->{|eff} b) (xs: n=>a) : {|eff} (n=>b) = for i. f xs.i
def zip (xs:n=>a) (ys:n=>b) : (n=>(a&b)) = for i. (xs.i, ys.i)
def unzip (xys:n=>(a&b)) : (n=>a & n=>b) = (map fst xys, map snd xys)
def fanout (n:Type) (x:a) : n=>a = for i. x
def sq (d:Mul a) ?=> (x:a) : a = x * x
def abs (x:Real) : Real = select (x > 0.0) x (-x)
def mod (x:Int) (y:Int) : Int = rem (y + rem x y) y
def compose (f:b->c) (g:a->b) (x:a) : c = f (g x)

def slice (xs:n=>a) (start:Int) (m:Type) : m=>a =
  for i. xs.(fromOrdinal _ (ordinal i + start))

def scan (init:a) (body:n->a->(a&b)) : (a & n=>b) =
  swap $ withState init \s. for i.
    c = get s
    (c', y) = body i c
    s := c'
    y

def fold (init:a) (body:(n->a->a)) : a = fst $ scan init \i x. (body i x, ())
def reduce (identity:a) (binop:(a->a->a)) (xs:n=>a) : a =
  -- `binop` should be a commutative and associative, and form a
  -- commutative monoid with `identity`
  -- TODO: implement with a parallelizable monoid-parameterized writer
  fold identity (\i c. binop c xs.i)

-- TODO: call this `scan` and call the current `scan` something else
def scan' (init:a) (body:n->a->a) : n=>a = snd $ scan init \i x. dup (body i x)
-- TODO: allow tables-via-lambda and get rid of this
def fsum (xs:n->Real) : Real = snd $ withAccum \ref. for i. ref += xs i
def sum  (_: Add v) ?=> (xs:n=>v) : v = reduce zero (+) xs
def prod (_: Mul v) ?=> (xs:n=>v) : v = reduce one  (*) xs
def mean (n:Type) ?-> (xs:n=>Real) : Real = sum xs / i2r (size n)
def std (xs:n=>Real) : Real = sqrt $ mean (map sq xs) - sq (mean xs)
def any (xs:n=>Bool) : Bool = reduce False (||) xs
def all (xs:n=>Bool) : Bool = reduce True  (&&) xs


def applyN (n:Int) (x:a) (f:a -> a) : a =
  snd $ withState x \ref. for _:(Fin n).
    ref := f (get ref)

def linspace (n:Type) (low:Real) (high:Real) : n=>Real =
  dx = (high - low) / i2r (size n)
  for i:n. low + i2r (ordinal i) * dx

def transpose (x:n=>m=>Real) : m=>n=>Real = for i j. x.j.i
def vdot (x:n=>Real) (y:n=>Real) : Real = fsum \i. x.i * y.i

-- matmul. Better symbol to use? `@`?
(**) : (l=>m=>Real) -> (m=>n=>Real) -> (l=>n=>Real) = \x y.
  y' = transpose y
  for i k. fsum \j. x.i.j * y'.k.j

(**.) : (n=>m=>Real) -> (m=>Real) -> (n=>Real) = \mat v. for i. vdot mat.i v
(.**) : (m=>Real) -> (n=>m=>Real) -> (n=>Real) = flip (**.)

def inner (x:n=>Real) (mat:n=>m=>Real) (y:m=>Real) : Real =
  fsum \(i,j). x.i * mat.i.j * y.j

'Functions for working with the pseudorandom number generator

-- TODO: newtype
Key = Int

def hash (x:Key) (y:Int) : Key = %ffi threefry2x32 Int x y
def newKey (x:Int) : Key = hash 0 x
def splitKey (k:Key) : (Key & Key) = (hash k 0, hash k 1)
def splitKey3 (k:Key) : (Key & Key & Key) =
  (k1, k') = splitKey k
  (k2, k3) = splitKey k'
  (k1, k2, k3)

def many (f:Key->a) (k:Key) (i:n) : a = f (hash k (ordinal i))
def ixkey (k:Key) (i:n) : Key = hash k (ordinal i)
def ixkey2 (k:Key) (i:n) (j:m) : Key = hash (hash k (ordinal i)) (ordinal j)
def rand (k:Key) : Real = %ffi randunif Real k
def randVec (n:Int) (f: Key -> a) (k: Key) : Fin n => a =
  for i:(Fin n). f (ixkey k i)

def randn (k:Key) : Real =
  (k1, k2) = splitKey k
  u1 = rand k1
  u2 = rand k2
  sqrt ((-2.0) * log u1) * cos (2.0 * pi * u2)

def randIdx (n:Type) ?-> (k:Key) : n =
  unif = rand k
  fromOrdinal n $ floor $ unif * i2r (size n)

def bern (p:Real) (k:Key) : Bool = rand k < p

def randnVec (n:Type) ?-> (k:Key) : n=>Real =
  for i. randn (ixkey k i)

'min / max etc

def minBy (_:Ord o) ?=> (f:a->o) (x:a) (y:a) : a = select (f x < f y) x y
def maxBy (_:Ord o) ?=> (f:a->o) (x:a) (y:a) : a = select (f x > f y) x y

def min (_:Ord o) ?=> (x1: o) -> (x2: o) : o = minBy id x1 x2
def max (_:Ord o) ?=> (x1: o) -> (x2: o) : o = maxBy id x1 x2

def minimumBy (_:Ord o) ?=> (f:a->o) (xs:n=>a) : a =
  reduce xs.(0@_) (minBy f) xs
def maximumBy (_:Ord o) ?=> (f:a->o) (xs:n=>a) : a =
  reduce xs.(0@_) (maxBy f) xs

def minimum (_:Ord o) ?=> (xs:n=>o) : o = minimumBy id xs
def maximum (_:Ord o) ?=> (xs:n=>o) : o = maximumBy id xs

def argmin (_:Ord o) ?=> (xs:n=>o) : n =
  zeroth = (0@_, xs.(0@_))
  compare = \(idx1, x1) (idx2, x2).
    select (x1 < x2) (idx1, x1) (idx2, x2)
  zipped = for i. (i, xs.i)
  fst $ reduce zeroth compare zipped


'Automatic differentiation

-- TODO: add vector space constraints
def linearize (f:a->b) (x:a) : (b & a --o b) = %linearize f x
def jvp (f:a->b) (x:a) : a --o b = snd (linearize f x)
def transposeLinear (f:a --o b) : b --o a = %linearTranspose f

def vjp (f:a->b) (x:a) : (b & b --o a) =
  (y, df) = linearize f x
  (y, transposeLinear df)

def grad (f:a->Real) (x:a) : a = snd (vjp f x) 1.0

def deriv (f:Real->Real) (x:Real) : Real = jvp f x 1.0

def derivRev (f:Real->Real) (x:Real) : Real = snd (vjp f x) 1.0

def checkDerivBase (f:Real->Real) (x:Real) : Bool =
  -- TODO: parse 1e-5
  eps = 0.00005
  ansFwd  = deriv    f x
  ansRev  = derivRev f x
  ansNumeric = (f (x + eps) - f (x - eps)) / (2. * eps)
  isClose = \a b. abs (a - b) < 0.001
  isClose ansFwd ansNumeric && isClose ansRev ansNumeric

def checkDeriv (f:Real->Real) (x:Real) : Bool =
  checkDerivBase f x && checkDerivBase (deriv f) x

def while
    (eff:Effects) ?->
    (cond: Unit -> {|eff} Bool)
    (body: Unit -> {|eff} Unit)
    : {|eff} Unit =
  cond' : Unit -> {|eff} InternalBool = \_. unsafeCoerce _ $ cond ()
  %while cond' body

'Vector support

def UNSAFEFromOrdinal (n : Type) (i : Int) : n = %unsafeAsIndex n i

VectorWidth = 4  -- XXX: Keep this synced with the constant defined in Array.hs
VectorReal  = %VectorRealType

def packVector (a : Real) (b : Real) (c : Real) (d : Real) : VectorReal = %vectorPack a b c d
def indexVector (v : VectorReal) (i : Fin VectorWidth) : Real = %vectorIndex v i

-- NB: Backends should be smart enough to optimize this to a vector load from v
def loadVector (v : (Fin VectorWidth)=>Real) : VectorReal =
  idx = Fin VectorWidth
  (packVector v.(UNSAFEFromOrdinal idx 0)
              v.(UNSAFEFromOrdinal idx 1)
              v.(UNSAFEFromOrdinal idx 2)
              v.(UNSAFEFromOrdinal idx 3))
def storeVector (v : VectorReal) : (Fin VectorWidth)=>Real =
  idx = Fin VectorWidth
  [ indexVector v (UNSAFEFromOrdinal idx 0)
  , indexVector v (UNSAFEFromOrdinal idx 1)
  , indexVector v (UNSAFEFromOrdinal idx 2)
  , indexVector v (UNSAFEFromOrdinal idx 3) ]

def broadcastVector (v : Real) : VectorReal = packVector v v v v

@instance vectorRealAdd : Add VectorReal =
  (MkAdd ( \x y. %vfadd x y )
         ( \x y. %vfsub x y )
         ( broadcastVector zero ))
@instance vectorRealMul : Mul VectorReal =
  MkMul (\x y. %vfmul x y) $ packVector 1.0 1.0 1.0 1.0
@instance vectorRealVSpace : VSpace VectorReal =
  MkVSpace vectorRealAdd \x v. broadcastVector x * v

'Tiling

def Tile (n : Type) (m : Type) : Type = %IndexSlice n m

-- One can think of instances of `Tile n m` as injective functions `m -> n`,
-- with the special property that consecutive elements of m map to consecutive
-- elements of n. In this view (+>) is just function application, while ++>
-- is currying followed by function application. We cannot represent currying
-- in isolation, because `Tile n (Tile u v)` does not make sense, unlike `Tile n (u & v)`.
def (+>) (l : Type) ?-> (t:Tile n l) (i : l) : n = %sliceOffset t i
def (++>) (t : Tile n (u & v)) (i : u) : Tile n v = %sliceCurry t i

def tile  (l : Type) ?->
          (fTile : (t:(Tile n l) -> {|eff} l=>a))
          (fScalar : n -> {|eff} a) : {|eff} n=>a = %tiled fTile fScalar
def tile1 (n : Type) ?-> (l : Type) ?-> (m : Type) ?->
          (fTile : (t:(Tile n l) -> {|eff} m=>l=>a))
          (fScalar : n -> {|eff} m=>a) : {|eff} m=>n=>a = %tiledd fTile fScalar

-- TODO: This should become just `loadVector $ for i. arr.(t +> i)`
--       once we are able to eliminate temporary arrays. Until then, we inline for performance...
def loadTile (t : Tile n (Fin VectorWidth)) (arr : n=>Real) : VectorReal =
  idx = Fin VectorWidth
  (packVector arr.(t +> UNSAFEFromOrdinal idx 0)
              arr.(t +> UNSAFEFromOrdinal idx 1)
              arr.(t +> UNSAFEFromOrdinal idx 2)
              arr.(t +> UNSAFEFromOrdinal idx 3))

'Numerical utilities

def logsumexp (x: n=>Real) : Real =
  m = maximum x
  m + (log $ sum for i. exp (x.i - m))

def logsoftmax (x: n=>Real) : n=>Real =
  lse = logsumexp x
  for i. x.i - lse

def softmax (x: n=>Real) : n=>Real =
  m = maximum x
  e =  for i. exp (x.i - m)
  s = sum e
  for i. e.i / s

def evalpoly (_:VSpace v) ?=> (coefficients:n=>v) (x:Real) : v =
  -- Evaluate a polynomial at x.  Same as Numpy's polyval.
  fold zero \i c. coefficients.i + x .* c

data List a:Type =
  AsList n:Int foo:(Fin n => a)

def (<>) (x:List a) (y:List a) : List a =
  (AsList nx xs) = x
  (AsList ny ys) = y
  nz = nx + ny
  AsList _ $ for i:(Fin nz).
    i' = ordinal i
    case i' < nx of
      True  -> xs.(fromOrdinal _ i')
      False -> ys.(fromOrdinal _ (i' - nx))