Median.hs

module LLVM.DSL.Example.Median where

import qualified LLVM.DSL.Expression as Expr
import LLVM.DSL.Expression (Exp, (==*), (<=*), (&&*), (||*))

import qualified LLVM.Extra.Nice.Vector as NVec
import qualified LLVM.Extra.Nice.Value as NV

import qualified LLVM.Core as LLVM

import Data.Tuple.HT (uncurry3)
import Data.Word (Word8)



median3IfThen :: (NV.Comparison a) => Exp a -> Exp a -> Exp a -> Exp a
median3IfThen a b c =
   Expr.ifThenElse (a<=*b)
      (Expr.ifThenElse (b<=*c)
         b
         (Expr.ifThenElse (a<=*c) c a))
      (Expr.ifThenElse (a<=*c)
         a
         (Expr.ifThenElse (b<=*c) c b))

median3Select ::
   (NV.Comparison a, NV.Select a) => Exp a -> Exp a -> Exp a -> Exp a
median3Select a b c =
   Expr.select (a<=*b)
      (Expr.select (b<=*c)
         b
         (Expr.select (a<=*c) c a))
      (Expr.select (a<=*c)
         a
         (Expr.select (b<=*c) c b))

median3SelectShared ::
   (NV.Comparison a, NV.Select a) => Exp a -> Exp a -> Exp a -> Exp a
median3SelectShared a b c =
   Expr.with (a<=*b) $ \a_le_b ->
   Expr.with (b<=*c) $ \b_le_c ->
   Expr.with (a<=*c) $ \a_le_c ->
   Expr.select a_le_b
      (Expr.select b_le_c b (Expr.select a_le_c c a))
      (Expr.select a_le_c a (Expr.select b_le_c c b))

median3MinMax ::
   (NV.Comparison a, NV.Select a) => Exp a -> Exp a -> Exp a -> Exp a
median3MinMax a b c =
   let minab = Expr.min a b in
   let maxab = Expr.max a b in
   Expr.select (maxab <=* c) maxab $ Expr.select (minab <=* c) c minab

median3MinMaxVector ::
   (LLVM.Positive n, NVec.C a) =>
   (NV.Comparison a, NV.Select a) =>
   NVec.T n a -> NVec.T n a -> NVec.T n a ->
   LLVM.CodeGenFunction r (NVec.T n a)
median3MinMaxVector a b c =
   NVec.map (uncurry3 (Expr.unliftM3 median3MinMax) . NV.unzip3) $
   NVec.zip3 a b c


type NV = NV.T

median3Case ::
   (NV.Comparison a, NV.Select a) =>
   NV a -> NV a -> NV a -> LLVM.CodeGenFunction r (NV a)
median3Case a b c = do
   a_le_b <- NV.cmp LLVM.CmpLE a b
   a_le_c <- NV.cmp LLVM.CmpLE a c
   b_le_c <- NV.cmp LLVM.CmpLE b c
   let mask = NV.fromInteger'
   a_le_b_mask <- NV.select a_le_b (mask 1) (mask 0)
   a_le_c_mask <- NV.select a_le_c (mask 2) (mask 0)
   b_le_c_mask <- NV.select b_le_c (mask 4) (mask 0)
   maskNV <- NV.or a_le_b_mask =<< NV.or a_le_c_mask b_le_c_mask
   let maskE = Expr.lift0 (maskNV :: NV Word8)
   selectB <- Expr.unExp (maskE ==* 0 ||* maskE ==* 7)
   selectA <- Expr.unExp (maskE ==* 1 ||* maskE ==* 6)
   NV.select selectA a =<< NV.select selectB b c

median3CaseVec ::
   (NV.Comparison a, NV.Select a) =>
   NV a -> NV a -> NV a -> LLVM.CodeGenFunction r (NV a)
median3CaseVec a b c = do
   a_le_b <- NV.cmp LLVM.CmpLE a b
   a_le_c <- NV.cmp LLVM.CmpLE a c
   b_le_c <- NV.cmp LLVM.CmpLE b c
   let check ab ac bc =
         Expr.select (Expr.lift0 a_le_b) 1 0 ==* (ab :: Exp Word8)
         &&*
         Expr.select (Expr.lift0 a_le_c) 1 0 ==* (ac :: Exp Word8)
         &&*
         Expr.select (Expr.lift0 b_le_c) 1 0 ==* (bc :: Exp Word8)
   selectB <- Expr.unExp (check 0 0 0 ||* check 1 1 1)
   selectA <- Expr.unExp (check 1 0 0 ||* check 0 1 1)
   NV.select selectA a =<< NV.select selectB b c