From e80b2593edc3d216905279ebcfa797593a1efbfc Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Sun, 26 May 2024 14:57:34 +0200
Subject: Fast Fractional ops via C code

---
 src/Data/Array/Nested/Internal/Arith/Foreign.hs  | 18 ++++++++++++++
 src/Data/Array/Nested/Internal/Arith/Lists.hs    | 31 ++++++++++++++++++++++--
 src/Data/Array/Nested/Internal/Arith/Lists/TH.hs |  4 ++-
 3 files changed, 50 insertions(+), 3 deletions(-)

(limited to 'src/Data/Array/Nested/Internal/Arith')

diff --git a/src/Data/Array/Nested/Internal/Arith/Foreign.hs b/src/Data/Array/Nested/Internal/Arith/Foreign.hs
index 49effa1..ac83188 100644
--- a/src/Data/Array/Nested/Internal/Arith/Foreign.hs
+++ b/src/Data/Array/Nested/Internal/Arith/Foreign.hs
@@ -24,12 +24,30 @@ $(fmap concat . forM typesList $ \arithtype -> do
                [t| CInt -> Int64 -> Ptr $ttyp -> Ptr $ttyp -> $ttyp -> IO () |])
       ])
 
+$(fmap concat . forM floatTypesList $ \arithtype -> do
+    let ttyp = conT (atType arithtype)
+    let base = "fbinary_" ++ atCName arithtype
+    sequence $ catMaybes
+      [Just (ForeignD . ImportF CCall Unsafe ("oxarop_" ++ base ++ "_sv") (mkName ("c_" ++ base ++ "_sv")) <$>
+               [t| CInt -> Int64 -> Ptr $ttyp -> $ttyp -> Ptr $ttyp -> IO () |])
+      ,Just (ForeignD . ImportF CCall Unsafe ("oxarop_" ++ base ++ "_vv") (mkName ("c_" ++ base ++ "_vv")) <$>
+               [t| CInt -> Int64 -> Ptr $ttyp -> Ptr $ttyp -> Ptr $ttyp -> IO () |])
+      ,Just (ForeignD . ImportF CCall Unsafe ("oxarop_" ++ base ++ "_vs") (mkName ("c_" ++ base ++ "_vs")) <$>
+               [t| CInt -> Int64 -> Ptr $ttyp -> Ptr $ttyp -> $ttyp -> IO () |])
+      ])
+
 $(fmap concat . forM typesList $ \arithtype -> do
     let ttyp = conT (atType arithtype)
     let base = "unary_" ++ atCName arithtype
     pure . ForeignD . ImportF CCall Unsafe ("oxarop_" ++ base) (mkName ("c_" ++ base)) <$>
       [t| CInt -> Int64 -> Ptr $ttyp -> Ptr $ttyp -> IO () |])
 
+$(fmap concat . forM floatTypesList $ \arithtype -> do
+    let ttyp = conT (atType arithtype)
+    let base = "funary_" ++ atCName arithtype
+    pure . ForeignD . ImportF CCall Unsafe ("oxarop_" ++ base) (mkName ("c_" ++ base)) <$>
+      [t| CInt -> Int64 -> Ptr $ttyp -> Ptr $ttyp -> IO () |])
+
 $(fmap concat . forM typesList $ \arithtype -> do
     let ttyp = conT (atType arithtype)
     let base = "reduce_" ++ atCName arithtype
diff --git a/src/Data/Array/Nested/Internal/Arith/Lists.hs b/src/Data/Array/Nested/Internal/Arith/Lists.hs
index 91e50ad..ce2836d 100644
--- a/src/Data/Array/Nested/Internal/Arith/Lists.hs
+++ b/src/Data/Array/Nested/Internal/Arith/Lists.hs
@@ -14,13 +14,18 @@ data ArithType = ArithType
   , atCName :: String  -- "i32"
   }
 
+floatTypesList :: [ArithType]
+floatTypesList =
+  [ArithType ''Float "float"
+  ,ArithType ''Double "double"
+  ]
+
 typesList :: [ArithType]
 typesList =
   [ArithType ''Int32 "i32"
   ,ArithType ''Int64 "i64"
-  ,ArithType ''Float "float"
-  ,ArithType ''Double "double"
   ]
+  ++ floatTypesList
 
 -- data ArithBOp = BO_ADD | BO_SUB | BO_MUL deriving (Show, Enum, Bounded)
 $(genArithDataType Binop "ArithBOp")
@@ -37,6 +42,21 @@ $(do clauses <- readArithLists Binop
               ,return $ FunD (mkName "aboNumOp") clauses])
 
 
+-- data ArithFBOp = FB_DIV deriving (Show, Enum, Bounded)
+$(genArithDataType FBinop "ArithFBOp")
+
+$(genArithNameFun FBinop ''ArithFBOp "afboName" (map toLower . drop 3))
+$(genArithEnumFun FBinop ''ArithFBOp "afboEnum")
+
+$(do clauses <- readArithLists FBinop
+                  (\name _num hsop -> return (Clause [ConP (mkName name) [] []]
+                                                     (NormalB (VarE 'mkName `AppE` LitE (StringL hsop)))
+                                                     []))
+                  return
+     sequence [SigD (mkName "afboNumOp") <$> [t| ArithFBOp -> Name |]
+              ,return $ FunD (mkName "afboNumOp") clauses])
+
+
 -- data ArithUOp = UO_NEG | UO_ABS | UO_SIGNUM | ... deriving (Show, Enum, Bounded)
 $(genArithDataType Unop "ArithUOp")
 
@@ -44,6 +64,13 @@ $(genArithNameFun Unop ''ArithUOp "auoName" (map toLower . drop 3))
 $(genArithEnumFun Unop ''ArithUOp "auoEnum")
 
 
+-- data ArithFUOp = FU_RECIP | ... deriving (Show, Enum, Bounded)
+$(genArithDataType FUnop "ArithFUOp")
+
+$(genArithNameFun FUnop ''ArithFUOp "afuoName" (map toLower . drop 3))
+$(genArithEnumFun FUnop ''ArithFUOp "afuoEnum")
+
+
 -- data ArithRedOp = RO_SUM1 | RO_PRODUCT1 deriving (Show, Enum, Bounded)
 $(genArithDataType Redop "ArithRedOp")
 
diff --git a/src/Data/Array/Nested/Internal/Arith/Lists/TH.hs b/src/Data/Array/Nested/Internal/Arith/Lists/TH.hs
index b748b97..b40a066 100644
--- a/src/Data/Array/Nested/Internal/Arith/Lists/TH.hs
+++ b/src/Data/Array/Nested/Internal/Arith/Lists/TH.hs
@@ -9,7 +9,7 @@ import Language.Haskell.TH
 import Text.Read
 
 
-data OpKind = Binop | Unop | Redop
+data OpKind = Binop | FBinop | Unop | FUnop | Redop
   deriving (Show, Eq)
 
 readArithLists :: OpKind
@@ -46,7 +46,9 @@ readArithLists targetkind fop fcombine = do
     parseField s = break (`elem` ",)") (dropWhile (== ' ') s)
 
     parseKind "BINOP" = Just Binop
+    parseKind "FBINOP" = Just FBinop
     parseKind "UNOP" = Just Unop
+    parseKind "FUNOP" = Just FUnop
     parseKind "REDOP" = Just Redop
     parseKind _ = Nothing
 
-- 
cgit v1.2.3-70-g09d2