Add (^), (^^)

Fixes #51
AccelerateHS · Nov 8, 2016 · df2ce87 · df2ce87
1 parent 362f7b8
commit df2ce87
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 7 deletions.
diff --git a/Data/Array/Accelerate.hs b/Data/Array/Accelerate.hs
@@ -276,7 +276,7 @@ module Data.Array.Accelerate (
   (&&*), (||*), not,
 
   -- *** Numeric operations
-  subtract, even, odd, gcd, lcm,
+  subtract, even, odd, gcd, lcm, (^), (^^),
 
   -- *** Shape manipulation
   index0, index1, unindex1, index2, unindex2, index3, unindex3,

diff --git a/Data/Array/Accelerate/Language.hs b/Data/Array/Accelerate/Language.hs
@@ -1,8 +1,9 @@
-{-# LANGUAGE ConstraintKinds  #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE TypeFamilies     #-}
-{-# LANGUAGE TypeOperators    #-}
-{-# LANGUAGE ViewPatterns     #-}
+{-# LANGUAGE ConstraintKinds     #-}
+{-# LANGUAGE FlexibleContexts    #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeFamilies        #-}
+{-# LANGUAGE TypeOperators       #-}
+{-# LANGUAGE ViewPatterns        #-}
 -- |
 -- Module      : Data.Array.Accelerate.Language
 -- Copyright   : [2008..2016] Manuel M T Chakravarty, Gabriele Keller
@@ -94,7 +95,7 @@ module Data.Array.Accelerate.Language (
   (!), (!!), shape, size, shapeSize,
 
   -- * Numeric functions
-  subtract, even, odd, gcd, lcm,
+  subtract, even, odd, gcd, lcm, (^), (^^),
 
   -- * Conversions
   ord, chr, boolToInt, bitcast,
@@ -778,6 +779,42 @@ lcm x y
   $ abs ((x `quot` (gcd x y)) * y)
 
 
+-- | Raise a number to a non-negative integral power
+--
+infixr 8 ^
+(^) :: forall a b. (Num a, Integral b) => Exp a -> Exp b -> Exp a
+x0 ^ y0 = cond (y0 <=* 0) 1 (f x0 y0)
+  where
+    f :: Exp a -> Exp b -> Exp a
+    f x y =
+      let (x',y') = untup2
+                  $ while (\(untup2 -> (_,v)) -> even v)
+                          (\(untup2 -> (u,v)) -> tup2 (u * u, v `quot` 2))
+                          (tup2 (x, y))
+      in
+      cond (y' ==* 1) x' (g (x'*x') ((y'-1) `quot` 2) x')
+
+    g :: Exp a -> Exp b -> Exp a -> Exp a
+    g x y z =
+      let (x',_,z') = untup3
+                    $ while (\(untup3 -> (_,v,_)) -> v /=* 1)
+                            (\(untup3 -> (u,v,w)) ->
+                              cond (even v) (tup3 (u*u, v     `quot` 2, w))
+                                            (tup3 (u*u, (v-1) `quot` 2, w*u)))
+                            (tup3 (x,y,z))
+      in
+      x' * z'
+
+-- | Raise a number to an integral power
+--
+infixr 8 ^^
+(^^) :: (Fractional a, Integral b) => Exp a -> Exp b -> Exp a
+x ^^ n
+  = cond (n >=* 0)
+  {- then -} (x ^ n)
+  {- else -} (recip (x ^ (negate n)))
+
+
 -- Conversions
 -- -----------