store primop eval functions in hashmap for faster lookup

grin-compiler · Apr 13, 2022 · f5e6806 · f5e6806
1 parent d3ed07d
commit f5e6806
Show file tree

Hide file tree

Showing 34 changed files with 1,272 additions and 1,296 deletions.
diff --git a/external-stg-interpreter/external-stg-interpreter.cabal b/external-stg-interpreter/external-stg-interpreter.cabal
@@ -74,6 +74,7 @@ library
                        text,
                        bytestring,
                        containers,
+                       unordered-containers,
                        primitive,
                        vector,
                        mtl,

diff --git a/external-stg-interpreter/lib/Stg/Interpreter.hs b/external-stg-interpreter/lib/Stg/Interpreter.hs
@@ -18,6 +18,7 @@ import Data.Maybe
 import Data.List (partition, isSuffixOf)
 import Data.Set (Set)
 import Data.Map (Map)
+import qualified Data.HashMap.Lazy as HashMap
 import qualified Data.Map.Strict as StrictMap
 import qualified Data.Map as Map
 import qualified Data.Set as Set
@@ -44,7 +45,6 @@ import qualified Stg.Interpreter.ThreadScheduler as Scheduler
 import qualified Stg.Interpreter.Debugger        as Debugger
 import qualified Stg.Interpreter.Debugger.Region as Debugger
 import qualified Stg.Interpreter.GC as GC
-
 import qualified Stg.Interpreter.PrimOp.Addr          as PrimAddr
 import qualified Stg.Interpreter.PrimOp.Array         as PrimArray
 import qualified Stg.Interpreter.PrimOp.SmallArray    as PrimSmallArray
@@ -852,34 +852,45 @@ scheduleWaitThread (StgTSO* tso, /*[out]*/HaskellObj* ret, Capability **pcap)
 -}
 ---------------------- primops
 
+primOpMap :: HashMap.HashMap Name PrimOpFunDef
+primOpMap = HashMap.fromList primOpList
+
+primOpList :: [(Name, PrimOpFunDef)]
+primOpList = concat
+  [ PrimAddr.primOps
+  , PrimArray.primOps
+  , PrimSmallArray.primOps
+  , PrimArrayArray.primOps
+  , PrimByteArray.primOps
+  , PrimChar.primOps
+  , PrimConcurrency.primOps
+  , PrimDelayWait.primOps
+  , PrimParallelism.primOps
+  , PrimExceptions.primOps
+  , PrimFloat.primOps
+  , PrimDouble.primOps
+  , PrimInt16.primOps
+  , PrimInt8.primOps
+  , PrimInt.primOps
+  , PrimMutVar.primOps
+  , PrimMVar.primOps
+  , PrimNarrowings.primOps
+  , PrimPrefetch.primOps
+  , PrimStablePointer.primOps
+  , PrimWeakPointer.primOps
+  , PrimWord16.primOps
+  , PrimWord8.primOps
+  , PrimWord.primOps
+  , PrimTagToEnum.primOps
+  , PrimUnsafe.primOps
+  , PrimMiscEtc.primOps
+  ]
+
+dupPrimOps :: [(Name, Int)]
+dupPrimOps = Map.toList $ Map.filter (> 1) $ foldl (\m n-> Map.insertWith (+) n 1 m) Map.empty $ map fst primOpList
+
 evalPrimOp :: HasCallStack => Name -> [Atom] -> Type -> Maybe TyCon -> M [Atom]
-evalPrimOp =
-  PrimAddr.evalPrimOp $
-  PrimArray.evalPrimOp $
-  PrimSmallArray.evalPrimOp $
-  PrimArrayArray.evalPrimOp $
-  PrimByteArray.evalPrimOp $
-  PrimChar.evalPrimOp $
-  PrimConcurrency.evalPrimOp $
-  PrimDelayWait.evalPrimOp $
-  PrimParallelism.evalPrimOp $
-  PrimExceptions.evalPrimOp $
-  PrimFloat.evalPrimOp $
-  PrimDouble.evalPrimOp $
-  PrimInt16.evalPrimOp $
-  PrimInt8.evalPrimOp $
-  PrimInt.evalPrimOp $
-  PrimMutVar.evalPrimOp $
-  PrimMVar.evalPrimOp $
-  PrimNarrowings.evalPrimOp $
-  PrimPrefetch.evalPrimOp $
-  PrimStablePointer.evalPrimOp $
-  PrimWeakPointer.evalPrimOp $
-  PrimWord16.evalPrimOp $
-  PrimWord8.evalPrimOp $
-  PrimWord.evalPrimOp $
-  PrimTagToEnum.evalPrimOp $
-  PrimUnsafe.evalPrimOp $
-  PrimMiscEtc.evalPrimOp $
-  unsupported where
-    unsupported op args _t _tc = stgErrorM $ "unsupported StgPrimOp: " ++ show op ++ " args: " ++ show args
+evalPrimOp op args t tc = case HashMap.lookup op primOpMap of
+  Just (NormalOp opEval)  -> opEval args
+  Just (SpecOp opEval)    -> opEval args t tc
+  Nothing                 -> stgErrorM $ "unsupported StgPrimOp: " ++ show op ++ " args: " ++ show args
diff --git a/external-stg-interpreter/lib/Stg/Interpreter/Base.hs b/external-stg-interpreter/lib/Stg/Interpreter/Base.hs
@@ -3,6 +3,7 @@ module Stg.Interpreter.Base where
 
 import Data.Word
 import Foreign.Ptr
+import Control.Monad.Writer (Writer, execWriter, tell)
 import Control.Monad.State.Strict
 import Data.List (foldl')
 import Data.Set (Set)
@@ -670,7 +671,23 @@ readHeapClosure a = readHeap a >>= \o -> case o of
 
 -- primop related
 
+data PrimOpFunDef
+  = NormalOp  ([Atom] -> M [Atom])
+  | SpecOp    ([Atom] -> Type -> Maybe TyCon -> M [Atom])
+
+type OpM = Writer [(Name, PrimOpFunDef)]
+
+getPrimOpList :: OpM () -> [(Name, PrimOpFunDef)]
+getPrimOpList = execWriter
+
+defOp :: Name -> ([Atom] -> M [Atom]) -> OpM ()
+defOp name eval = tell [(name, NormalOp eval)]
+
+defSpecOp :: Name -> ([Atom] -> Type -> Maybe TyCon -> M [Atom]) -> OpM ()
+defSpecOp name eval = tell [(name, SpecOp eval)]
+
 type PrimOpEval = Name -> [Atom] -> Type -> Maybe TyCon -> M [Atom]
+type PrimOpEvalFun = [Atom] -> Type -> Maybe TyCon -> M [Atom]
 
 --type BuiltinStgEval = Atom -> M [Atom]
 --type BuiltinStgApply = Atom -> [Atom] -> M [Atom]