Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Evaluation] [Names] Define all lookups in terms of 'contIndexZero' #6702

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
+65 −88
Open

[Evaluation] [Names] Define all lookups in terms of 'contIndexZero' #6702

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions ...ore/changelog.d/20241126_020822_effectfully_define_1_indexing_via_0_indexing.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
### Changed

- In #6702 made variable lookup faster increasing overall performance of the evaluator by 1%.
148 changes: 61 additions & 87 deletions plutus-core/index-envs/src/Data/RandomAccessList/SkewBinary.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,11 @@
{-# LANGUAGE ViewPatterns #-}
module Data.RandomAccessList.SkewBinary
( RAList(Cons,Nil)
, contIndexZero
, contIndexOne
, safeIndexZero
, unsafeIndexZero
, safeIndexOne
, safeIndexOneCont
, unsafeIndexOne
, Data.RandomAccessList.SkewBinary.null
, uncons
Expand Down Expand Up @@ -48,7 +49,7 @@ data RAList a = BHead
null :: RAList a -> Bool
null Nil = True
null _ = False
{-# INLINABLE null #-}
{-# INLINE null #-}

{-# complete Cons, Nil #-}
{-# complete BHead, Nil #-}
Expand All @@ -63,6 +64,7 @@ cons :: a -> RAList a -> RAList a
cons x = \case
(BHead w1 t1 (BHead w2 t2 ts')) | w1 == w2 -> BHead (2*w1+1) (Node x t1 t2) ts'
ts -> BHead 1 (Leaf x) ts
{-# INLINE cons #-}

-- /O(1)/
uncons :: RAList a -> Maybe (a, RAList a)
Expand All @@ -74,122 +76,94 @@ uncons = \case
-- split the node in two)
in Just (x, BHead halfSize t1 $ BHead halfSize t2 ts)
Nil -> Nothing
{-# INLINE uncons #-}

-- 0-based
unsafeIndexZero :: RAList a -> Word64 -> a
unsafeIndexZero Nil _ = error "out of bounds"
unsafeIndexZero (BHead w t ts) !i =
if i < w
then indexTree w i t
else unsafeIndexZero ts (i-w)
where
indexTree :: Word64 -> Word64 -> Tree a -> a
indexTree 1 0 (Leaf x) = x
indexTree _ _ (Leaf _) = error "out of bounds"
indexTree _ 0 (Node x _ _) = x
indexTree treeSize offset (Node _ t1 t2) =
let halfSize = unsafeShiftR treeSize 1 -- probably faster than `div w 2`
in if offset <= halfSize
then indexTree halfSize (offset - 1) t1
else indexTree halfSize (offset - 1 - halfSize) t2

-- 0-based
safeIndexZero :: RAList a -> Word64 -> Maybe a
safeIndexZero Nil _ = Nothing
safeIndexZero (BHead w t ts) !i =
if i < w
then indexTree w i t
else safeIndexZero ts (i-w)
where
indexTree :: Word64 -> Word64 -> Tree a -> Maybe a
indexTree 1 0 (Leaf x) = Just x
indexTree _ _ (Leaf _) = Nothing
indexTree _ 0 (Node x _ _) = Just x
indexTree treeSize offset (Node _ t1 t2) =
let halfSize = unsafeShiftR treeSize 1 -- probably faster than `div w 2`
in if offset <= halfSize
then indexTree halfSize (offset - 1) t1
else indexTree halfSize (offset - 1 - halfSize) t2

-- 1-based
unsafeIndexOne :: RAList a -> Word64 -> a
unsafeIndexOne Nil _ = error "out of bounds"
unsafeIndexOne (BHead w t ts) !i =
if i <= w
then indexTree w i t
else unsafeIndexOne ts (i-w)
where
indexTree :: Word64 -> Word64 -> Tree a -> a
indexTree _ 0 _ = error "index zero"
indexTree 1 1 (Leaf x) = x
indexTree _ _ (Leaf _) = error "out of bounds"
indexTree _ 1 (Node x _ _) = x
indexTree treeSize offset (Node _ t1 t2) =
let halfSize = unsafeShiftR treeSize 1 -- probably faster than `div w 2`
offset' = offset - 1
in if offset' <= halfSize
then indexTree halfSize offset' t1
else indexTree halfSize (offset' - halfSize) t2

{- Note [Optimizations of safeIndexOneCont]
Bangs in the local definitions of 'safeIndexOneCont' are needed to tell GHC that the functions are
{- Note [Optimizations of contIndexZero]
Bangs in the local definitions of 'contIndexZero' are needed to tell GHC that the functions are
strict in the 'Word64' argument, so that GHC produces workers operating on @Word64#@.

The function itself is CPS-ed, so that the arguments force the local definitions to be retained
within 'safeIndexOneCont' instead of being pulled out via full-laziness or some other optimization
pass. This ensures that when 'safeIndexOneCont' gets inlined, the local definitions appear directly
in the GHC Core, allowing GHC to inline the arguments of 'safeIndexOneCont' and transform the whole
within 'contIndexZero' instead of being pulled out via full-laziness or some other optimization
pass. This ensures that when 'contIndexZero' gets inlined, the local definitions appear directly
in the GHC Core, allowing GHC to inline the arguments of 'contIndexZero' and transform the whole
thing into a beautiful recursive join point full of @Word64#@s, i.e. allocating very little if
anything at all.
-}

-- See Note [Optimizations of safeIndexOneCont].
safeIndexOneCont :: forall a b. b -> (a -> b) -> RAList a -> Word64 -> b
safeIndexOneCont z f = findTree where
-- See Note [Optimizations of contIndexZero].
contIndexZero :: forall a b. b -> (a -> b) -> RAList a -> Word64 -> b
contIndexZero z f = findTree where
findTree :: RAList a -> Word64 -> b
-- See Note [Optimizations of safeIndexOneCont].
-- See Note [Optimizations of contIndexZero].
findTree Nil !_ = z
findTree (BHead w t ts) i =
if i <= w
if i < w
then indexTree w i t
else findTree ts (i-w)

indexTree :: Word64 -> Word64 -> Tree a -> b
-- See Note [Optimizations of safeIndexOneCont].
indexTree !w 1 t = case t of
-- See Note [Optimizations of contIndexZero].
indexTree !w 0 t = case t of
Node x _ _ -> f x
Leaf x -> if w == 1 then f x else z
indexTree _ 0 _ = z -- "index zero"
indexTree _ _ (Leaf _) = z
indexTree treeSize offset (Node _ t1 t2) =
let halfSize = unsafeShiftR treeSize 1 -- probably faster than `div w 2`
offset' = offset - 1
in if offset' <= halfSize
then indexTree halfSize offset' t1
else indexTree halfSize (offset' - halfSize) t2
{-# INLINE safeIndexOneCont #-}
in if offset <= halfSize
then indexTree halfSize (offset - 1) t1
else indexTree halfSize (offset - 1 - halfSize) t2
{-# INLINE contIndexZero #-}

contIndexOne :: forall a b. b -> (a -> b) -> RAList a -> Word64 -> b
contIndexOne z _ _ 0 = z
contIndexOne z f t n = contIndexZero z f t (n - 1)
{-# INLINE contIndexOne #-}

-- 0-based
unsafeIndexZero :: RAList a -> Word64 -> a
unsafeIndexZero = contIndexZero (error "out of bounds") id
{-# INLINE unsafeIndexZero #-}

-- 0-based
safeIndexZero :: RAList a -> Word64 -> Maybe a
safeIndexZero = contIndexZero Nothing Just
{-# INLINE safeIndexZero #-}

-- 1-based
unsafeIndexOne :: RAList a -> Word64 -> a
unsafeIndexOne = contIndexOne (error "out of bounds") id
{-# INLINE unsafeIndexOne #-}

-- 1-based
safeIndexOne :: RAList a -> Word64 -> Maybe a
safeIndexOne = safeIndexOneCont Nothing Just
safeIndexOne = contIndexOne Nothing Just
{-# INLINE safeIndexOne #-}

instance RAL.RandomAccessList (RAList a) where
type Element (RAList a) = a

{-# INLINABLE empty #-}
empty = Nil
{-# INLINABLE cons #-}
{-# INLINE empty #-}
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We really want all of that to inline, no point is chickening out and saying INLINABLE.


cons = Cons
{-# INLINABLE uncons #-}
{-# INLINE cons #-}

uncons = uncons
{-# INLINABLE length #-}
length Nil = 0
length (BHead sz _ tl) = sz + RAL.length tl
{-# INLINABLE indexZero #-}
{-# INLINE uncons #-}

length = go 0 where
go !acc Nil = acc
go !acc (BHead sz _ tl) = go (acc + sz) tl
{-# INLINE length #-}

indexZero = safeIndexZero
{-# INLINABLE indexOne #-}
{-# INLINE indexZero #-}

indexOne = safeIndexOne
{-# INLINABLE unsafeIndexZero #-}
{-# INLINE indexOne #-}

unsafeIndexZero = unsafeIndexZero
{-# INLINABLE unsafeIndexOne #-}
{-# INLINE unsafeIndexZero #-}

unsafeIndexOne = unsafeIndexOne
{-# INLINE unsafeIndexOne #-}
2 changes: 1 addition & 1 deletion plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Evaluation/Machine/Cek/Internal.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -878,7 +878,7 @@ enterComputeCek = computeCek
-- | Look up a variable name in the environment.
lookupVarName :: NamedDeBruijn -> CekValEnv uni fun ann -> CekM uni fun s (CekValue uni fun ann)
lookupVarName varName@(NamedDeBruijn _ varIx) varEnv =
Env.safeIndexOneCont
Env.contIndexOne
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This names better matches the other names in the module, should've gone with it originally.

(throwingWithCause _MachineError OpenTermEvaluatedMachineError . Just $ Var () varName)
pure
varEnv
Expand Down