Merge pull request unisonweb#5447 from unisonweb/cp/more-instructions

develop7 · Dec 3, 2024 · 5542282 · 5542282
2 parents ce0e9af + 4ac40d4
commit 5542282
Show file tree

Hide file tree

Showing 11 changed files with 222 additions and 90 deletions.
diff --git a/unison-runtime/src/Unison/Runtime/ANF.hs b/unison-runtime/src/Unison/Runtime/ANF.hs
@@ -1291,9 +1291,11 @@ data POp
   | DECI -- dec
   | LEQI -- <=
   | EQLI -- ==
+  | TRNC -- truncate0
   -- Nat
   | ADDN -- +
   | SUBN -- -
+  | DRPN -- drop
   | MULN
   | DIVN -- /
   | MODN -- mod
@@ -1414,6 +1416,13 @@ data POp
   | TFRC -- try force
   | SDBL -- sandbox link list
   | SDBV -- sandbox check for Values
+  -- Refs
+  | REFN -- Ref.new
+  | REFR -- Ref.read
+  | REFW -- Ref.write
+  | RCAS -- Ref.cas
+  | RRFC -- Ref.readForCas
+  | TIKR -- Ref.Ticket.read
   deriving (Show, Eq, Ord, Enum, Bounded)
 
 type ANormal = ABTN.Term ANormalF

diff --git a/unison-runtime/src/Unison/Runtime/ANF/Serialize.hs b/unison-runtime/src/Unison/Runtime/ANF/Serialize.hs
@@ -654,6 +654,14 @@ pOpCode op = case op of
   IORI -> 126
   XORI -> 127
   COMI -> 128
+  DRPN -> 129
+  TRNC -> 130
+  REFN -> 131
+  REFR -> 132
+  REFW -> 133
+  RCAS -> 134
+  RRFC -> 135
+  TIKR -> 136
 
 pOpAssoc :: [(POp, Word16)]
 pOpAssoc = map (\op -> (op, pOpCode op)) [minBound .. maxBound]

diff --git a/unison-runtime/src/Unison/Runtime/Builtin.hs b/unison-runtime/src/Unison/Runtime/Builtin.hs
@@ -49,12 +49,6 @@ import Data.ByteString (hGet, hGetSome, hPut)
 import Data.ByteString.Lazy qualified as L
 import Data.Default (def)
 import Data.Digest.Murmur64 (asWord64, hash64)
-import Data.IORef as SYS
-  ( IORef,
-    newIORef,
-    readIORef,
-    writeIORef,
-  )
 import Data.IP (IP)
 import Data.Map qualified as Map
 import Data.PEM (PEM, pemContent, pemParseLBS)
@@ -182,11 +176,7 @@ import Unison.Util.Bytes qualified as Bytes
 import Unison.Util.EnumContainers as EC
 import Unison.Util.RefPromise
   ( Promise,
-    Ticket,
-    casIORef,
     newPromise,
-    peekTicket,
-    readForCAS,
     readPromise,
     tryReadPromise,
     writePromise,
@@ -373,7 +363,7 @@ shli = binop SHLI
 shri = binop SHRI
 powi = binop POWI
 
-addn, subn, muln, divn, modn, shln, shrn, pown :: (Var v) => SuperNormal v
+addn, subn, muln, divn, modn, shln, shrn, pown, dropn :: (Var v) => SuperNormal v
 addn = binop ADDN
 subn = binop SUBN
 muln = binop MULN
@@ -382,6 +372,7 @@ modn = binop MODN
 shln = binop SHLN
 shrn = binop SHRN
 pown = binop POWN
+dropn = binop DRPN
 
 eqi, eqn, lti, ltn, lei, len :: (Var v) => SuperNormal v
 eqi = cmpop EQLI
@@ -490,17 +481,7 @@ i2f = unop ITOF
 n2f = unop NTOF
 
 trni :: (Var v) => SuperNormal v
-trni = unop0 4 $ \[x, z, b, tag, n] ->
-  -- TODO: Do we need to do all calculations _before_ the branch?
-  -- Should probably just replace this with an instruction.
-  TLetD z UN (TLit $ N 0)
-    . TLetD b UN (TPrm LEQI [x, z])
-    . TLetD tag UN (TLit $ I $ fromIntegral $ unboxedTypeTagToInt NatTag)
-    . TLetD n UN (TPrm CAST [x, tag])
-    . TMatch b
-    $ MatchIntegral
-      (mapSingleton 1 $ TVar z)
-      (Just $ TVar n)
+trni = unop TRNC
 
 modular :: (Var v) => POp -> (Bool -> ANormal v) -> SuperNormal v
 modular pop ret =
@@ -518,20 +499,6 @@ oddi = modular MODI (\b -> if b then tru else fls)
 evnn = modular MODN (\b -> if b then fls else tru)
 oddn = modular MODN (\b -> if b then tru else fls)
 
-dropn :: (Var v) => SuperNormal v
-dropn = binop0 4 $ \[x, y, b, r, tag, n] ->
-  TLetD b UN (TPrm LEQN [x, y])
-    -- TODO: Can we avoid this work until after the branch?
-    -- Should probably just replace this with an instruction.
-    . TLetD tag UN (TLit $ I $ fromIntegral $ unboxedTypeTagToInt NatTag)
-    . TLetD r UN (TPrm SUBN [x, y])
-    . TLetD n UN (TPrm CAST [r, tag])
-    $ ( TMatch b $
-          MatchIntegral
-            (mapSingleton 1 $ TLit $ N 0)
-            (Just $ TVar n)
-      )
-
 appendt, taket, dropt, indext, indexb, sizet, unconst, unsnoct :: (Var v) => SuperNormal v
 appendt = binop0 0 $ \[x, y] -> TPrm CATT [x, y]
 taket = binop0 0 $ \[x, y] ->
@@ -1019,6 +986,54 @@ any'extract =
       TMatch v $
         MatchData Ty.anyRef (mapSingleton 0 $ ([BX], TAbs v1 (TVar v1))) Nothing
 
+-- Refs
+
+ -- The docs for IORef state that IORef operations can be observed
+  -- out of order ([1]) but actually GHC does emit the appropriate
+  -- load and store barriers nowadays ([2], [3]).
+  --
+  -- [1] https://hackage.haskell.org/package/base-4.17.0.0/docs/Data-IORef.html#g:2
+  -- [2] https://github.com/ghc/ghc/blob/master/compiler/GHC/StgToCmm/Prim.hs#L286
+  -- [3] https://github.com/ghc/ghc/blob/master/compiler/GHC/StgToCmm/Prim.hs#L298
+ref'read :: SuperNormal Symbol
+ref'read =
+  unop0 0 $ \[ref] -> (TPrm REFR [ref])
+
+ref'write :: SuperNormal Symbol
+ref'write =
+  binop0 0 $ \[ref, val] -> (TPrm REFW [ref, val])
+
+-- In GHC, CAS returns both a Boolean and the current value of the
+-- IORef, which can be used to retry a failed CAS.
+-- This strategy is more efficient than returning a Boolean only
+-- because it uses a single call to cmpxchg in assembly (see [1]) to
+-- avoid an extra read per CAS iteration, however it's not supported
+-- in Scheme.
+-- Therefore, we adopt the more common signature that only returns a
+-- Boolean, which doesn't even suffer from spurious failures because
+-- GHC issues loads of mutable variables with memory_order_acquire
+-- (see [2])
+--
+-- [1]: https://github.com/ghc/ghc/blob/master/rts/PrimOps.cmm#L697
+-- [2]: https://github.com/ghc/ghc/blob/master/compiler/GHC/StgToCmm/Prim.hs#L285
+ref'cas :: SuperNormal Symbol
+ref'cas =
+  Lambda [BX, BX, BX]
+    . TAbss [x, y, z]
+    . TLetD b UN (TPrm RCAS [x, y, z])
+    $ boolift b
+  where
+    (x, y, z, b) = fresh
+
+ref'ticket'read :: SuperNormal Symbol
+ref'ticket'read = unop0 0 $ TPrm TIKR
+
+ref'readForCas :: SuperNormal Symbol
+ref'readForCas = unop0 0 $ TPrm RRFC
+
+ref'new :: SuperNormal Symbol
+ref'new = unop0 0 $ TPrm REFN
+
 seek'handle :: ForeignOp
 seek'handle instr =
   ([BX, BX, BX],)
@@ -1895,7 +1910,14 @@ builtinLookup =
         ("validateSandboxed", (Untracked, check'sandbox)),
         ("Value.validateSandboxed", (Tracked, value'sandbox)),
         ("sandboxLinks", (Tracked, sandbox'links)),
-        ("IO.tryEval", (Tracked, try'eval))
+        ("IO.tryEval", (Tracked, try'eval)),
+        ("Ref.read", (Tracked, ref'read)),
+        ("Ref.write", (Tracked, ref'write)),
+        ("Ref.cas", (Tracked, ref'cas)),
+        ("Ref.Ticket.read", (Tracked, ref'ticket'read)),
+        ("Ref.readForCas", (Tracked, ref'readForCas)),
+        ("Scope.ref", (Untracked, ref'new)),
+        ("IO.ref", (Tracked, ref'new))
       ]
       ++ foreignWrappers
 
@@ -2394,53 +2416,6 @@ declareForeigns = do
   declareForeign Tracked "STM.retry" unitDirect . mkForeign $
     \() -> unsafeSTMToIO STM.retry :: IO Val
 
-  -- Scope and Ref stuff
-  declareForeign Untracked "Scope.ref" (argNDirect 1)
-    . mkForeign
-    $ \(c :: Val) -> newIORef c
-
-  declareForeign Tracked "IO.ref" (argNDirect 1)
-    . mkForeign
-    $ \(c :: Val) -> evaluate c >>= newIORef
-
-  -- The docs for IORef state that IORef operations can be observed
-  -- out of order ([1]) but actually GHC does emit the appropriate
-  -- load and store barriers nowadays ([2], [3]).
-  --
-  -- [1] https://hackage.haskell.org/package/base-4.17.0.0/docs/Data-IORef.html#g:2
-  -- [2] https://github.com/ghc/ghc/blob/master/compiler/GHC/StgToCmm/Prim.hs#L286
-  -- [3] https://github.com/ghc/ghc/blob/master/compiler/GHC/StgToCmm/Prim.hs#L298
-  declareForeign Untracked "Ref.read" (argNDirect 1) . mkForeign $
-    \(r :: IORef Val) -> readIORef r
-
-  declareForeign Untracked "Ref.write" arg2To0 . mkForeign $
-    \(r :: IORef Val, c :: Val) -> evaluate c >>= writeIORef r
-
-  declareForeign Tracked "Ref.readForCas" (argNDirect 1) . mkForeign $
-    \(r :: IORef Val) -> readForCAS r
-
-  declareForeign Tracked "Ref.Ticket.read" (argNDirect 1) . mkForeign $
-    \(t :: Ticket Val) -> pure $ peekTicket t
-
-  -- In GHC, CAS returns both a Boolean and the current value of the
-  -- IORef, which can be used to retry a failed CAS.
-  -- This strategy is more efficient than returning a Boolean only
-  -- because it uses a single call to cmpxchg in assembly (see [1]) to
-  -- avoid an extra read per CAS iteration, however it's not supported
-  -- in Scheme.
-  -- Therefore, we adopt the more common signature that only returns a
-  -- Boolean, which doesn't even suffer from spurious failures because
-  -- GHC issues loads of mutable variables with memory_order_acquire
-  -- (see [2])
-  --
-  -- [1]: https://github.com/ghc/ghc/blob/master/rts/PrimOps.cmm#L697
-  -- [2]: https://github.com/ghc/ghc/blob/master/compiler/GHC/StgToCmm/Prim.hs#L285
-  declareForeign Tracked "Ref.cas" (argNToBool 3) . mkForeign $
-    \(r :: IORef Val, t :: Ticket Val, v :: Val) -> fmap fst $
-      do
-        t <- evaluate t
-        casIORef r t v
-
   declareForeign Tracked "Promise.new" unitDirect . mkForeign $
     \() -> newPromise @Val
 

diff --git a/unison-runtime/src/Unison/Runtime/Foreign.hs b/unison-runtime/src/Unison/Runtime/Foreign.hs
@@ -22,6 +22,7 @@ where
 import Control.Concurrent (MVar, ThreadId)
 import Control.Concurrent.STM (TVar)
 import Crypto.Hash qualified as Hash
+import Data.Atomics qualified as Atomic
 import Data.IORef (IORef)
 import Data.Tagged (Tagged (..))
 import Data.X509 qualified as X509
@@ -261,6 +262,8 @@ instance BuiltinForeign Value where foreignRef = Tagged Ty.valueRef
 
 instance BuiltinForeign TimeSpec where foreignRef = Tagged Ty.timeSpecRef
 
+instance BuiltinForeign (Atomic.Ticket a) where foreignRef = Tagged Ty.ticketRef
+
 data HashAlgorithm where
   -- Reference is a reference to the hash algorithm
   HashAlgorithm :: (Hash.HashAlgorithm a) => Reference -> a -> HashAlgorithm

diff --git a/unison-runtime/src/Unison/Runtime/MCode.hs b/unison-runtime/src/Unison/Runtime/MCode.hs
@@ -323,6 +323,7 @@ data UPrim1
   | FLOR -- floor
   | TRNF -- truncate
   | RNDF -- round
+  | TRNC -- truncate
   deriving (Show, Eq, Ord, Enum, Bounded)
 
 data UPrim2
@@ -366,6 +367,7 @@ data UPrim2
   | MAXF -- max
   | MINF -- min
   | CAST -- unboxed runtime type cast (int to nat, etc.)
+  | DRPN -- dropn
   deriving (Show, Eq, Ord, Enum, Bounded)
 
 data BPrim1
@@ -400,6 +402,11 @@ data BPrim1
   -- debug
   | DBTX -- debug text
   | SDBL -- sandbox link list
+  | -- Refs
+    REFN -- Ref.new
+  | REFR -- Ref.read
+  | RRFC
+  | TIKR
   deriving (Show, Eq, Ord, Enum, Bounded)
 
 data BPrim2
@@ -435,6 +442,8 @@ data BPrim2
   -- code
   | SDBX -- sandbox
   | SDBV -- sandbox Value
+  -- Refs
+  | REFW -- Ref.write
   deriving (Show, Eq, Ord, Enum, Bounded)
 
 data MLit
@@ -472,6 +481,9 @@ data GInstr comb
       !BPrim2
       !Int
       !Int
+  | -- Use a check-and-set ticket to update a reference
+    -- (ref stack index, ticket stack index, new value stack index)
+    RefCAS !Int !Int !Int
   | -- Call out to a Haskell function. This is considerably slower
     -- for very simple operations, hence the primops.
     ForeignCall
@@ -1196,6 +1208,7 @@ emitPOp ANF.ADDI = emitP2 ADDI
 emitPOp ANF.ADDN = emitP2 ADDN
 emitPOp ANF.SUBI = emitP2 SUBI
 emitPOp ANF.SUBN = emitP2 SUBN
+emitPOp ANF.DRPN = emitP2 DRPN
 emitPOp ANF.MULI = emitP2 MULI
 emitPOp ANF.MULN = emitP2 MULN
 emitPOp ANF.DIVI = emitP2 DIVI
@@ -1218,6 +1231,7 @@ emitPOp ANF.INCI = emitP1 INCI
 emitPOp ANF.INCN = emitP1 INCN
 emitPOp ANF.DECI = emitP1 DECI
 emitPOp ANF.DECN = emitP1 DECN
+emitPOp ANF.TRNC = emitP1 TRNC
 emitPOp ANF.TZRO = emitP1 TZRO
 emitPOp ANF.LZRO = emitP1 LZRO
 emitPOp ANF.POPC = emitP1 POPC
@@ -1323,6 +1337,13 @@ emitPOp ANF.SDBV = emitBP2 SDBV
 emitPOp ANF.EROR = emitBP2 THRO
 emitPOp ANF.TRCE = emitBP2 TRCE
 emitPOp ANF.DBTX = emitBP1 DBTX
+-- Refs
+emitPOp ANF.REFN = emitBP1 REFN
+emitPOp ANF.REFR = emitBP1 REFR
+emitPOp ANF.REFW = emitBP2 REFW
+emitPOp ANF.RCAS = refCAS
+emitPOp ANF.RRFC = emitBP1 RRFC
+emitPOp ANF.TIKR = emitBP1 TIKR
 -- non-prim translations
 emitPOp ANF.BLDS = Seq
 emitPOp ANF.FORK = \case
@@ -1380,6 +1401,13 @@ emitBP2 p a =
     "wrong number of args for binary boxed primop: "
       ++ show (p, a)
 
+refCAS :: Args -> Instr
+refCAS (VArgN (primArrayToList -> [i, j, k])) = RefCAS i j k
+refCAS a =
+  internalBug $
+    "wrong number of args for refCAS: "
+      ++ show a
+
 emitDataMatching ::
   (Var v) =>
   Reference ->

diff --git a/unison-runtime/src/Unison/Runtime/MCode/Serialize.hs b/unison-runtime/src/Unison/Runtime/MCode/Serialize.hs
@@ -160,6 +160,7 @@ data InstrT
   | AtomicallyT
   | SeqT
   | TryForceT
+  | RefCAST
 
 instance Tag InstrT where
   tag2word UPrim1T = 0
@@ -179,6 +180,7 @@ instance Tag InstrT where
   tag2word AtomicallyT = 14
   tag2word SeqT = 15
   tag2word TryForceT = 16
+  tag2word RefCAST = 17
 
   word2tag 0 = pure UPrim1T
   word2tag 1 = pure UPrim2T
@@ -197,6 +199,7 @@ instance Tag InstrT where
   word2tag 14 = pure AtomicallyT
   word2tag 15 = pure SeqT
   word2tag 16 = pure TryForceT
+  word2tag 17 = pure RefCAST
   word2tag n = unknownTag "InstrT" n
 
 putInstr :: (MonadPut m) => GInstr cix -> m ()
@@ -205,6 +208,7 @@ putInstr = \case
   (UPrim2 up i j) -> putTag UPrim2T *> putTag up *> pInt i *> pInt j
   (BPrim1 bp i) -> putTag BPrim1T *> putTag bp *> pInt i
   (BPrim2 bp i j) -> putTag BPrim2T *> putTag bp *> pInt i *> pInt j
+  (RefCAS i j k) -> putTag RefCAST *> pInt i *> pInt j *> pInt k
   (ForeignCall b w a) -> putTag ForeignCallT *> serialize b *> pWord w *> putArgs a
   (SetDyn w i) -> putTag SetDynT *> pWord w *> pInt i
   (Capture w) -> putTag CaptureT *> pWord w
@@ -226,6 +230,7 @@ getInstr =
     UPrim2T -> UPrim2 <$> getTag <*> gInt <*> gInt
     BPrim1T -> BPrim1 <$> getTag <*> gInt
     BPrim2T -> BPrim2 <$> getTag <*> gInt <*> gInt
+    RefCAST -> RefCAS <$> gInt <*> gInt <*> gInt
     ForeignCallT -> ForeignCall <$> deserialize <*> gWord <*> getArgs
     SetDynT -> SetDyn <$> gWord <*> gInt
     CaptureT -> Capture <$> gWord