Skip to content

Introduce basicUnsafeIndexM# #489

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

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Introduce basicUnsafeIndexM# #489

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4b682ae
Introduce basicUnsafeIndexM#
Bodigrim Mar 23, 2024
7acc4ab
Use coerce in definition Unbox instances wrapping primitive vector
Shimuuar Apr 4, 2024
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
7 changes: 6 additions & 1 deletion vector/internal/GenUnboxTuple.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,12 @@
#!/usr/bin/env cabal
{- cabal:
build-depends: base, pretty
-}
{-# LANGUAGE ParallelListComp #-}

module Main where

import Prelude hiding ((<>))
import Text.PrettyPrint

import System.Environment ( getArgs )
Expand Down Expand Up @@ -235,6 +240,6 @@ generate n =
,("basicUnsafeThaw", gen_unsafeThaw)
,("basicLength", gen_length "V")
,("basicUnsafeSlice", gen_unsafeSlice "G" "V")
,("basicUnsafeIndexM", gen_basicUnsafeIndexM)
,("basicUnsafeIndexM#", gen_basicUnsafeIndexM)
,("basicUnsafeCopy", gen_unsafeCopy "V" qG)
,("elemseq", gen_elemseq)]
64 changes: 32 additions & 32 deletions vector/internal/unbox-tuple-instances
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,11 +90,11 @@ instance (Unbox a, Unbox b) => G.Vector Vector (a, b) where
basicUnsafeSlice i_ m_ (V_2 _ as bs)
= V_2 m_ (G.basicUnsafeSlice i_ m_ as)
(G.basicUnsafeSlice i_ m_ bs)
{-# INLINE basicUnsafeIndexM #-}
basicUnsafeIndexM (V_2 _ as bs) i_
{-# INLINE basicUnsafeIndexM# #-}
basicUnsafeIndexM# (V_2 _ as bs) i_
= do
a <- G.basicUnsafeIndexM as i_
b <- G.basicUnsafeIndexM bs i_
a <- G.basicUnsafeIndexM# as i_
b <- G.basicUnsafeIndexM# bs i_
return (a, b)
{-# INLINE basicUnsafeCopy #-}
basicUnsafeCopy (MV_2 _ as1 bs1) (V_2 _ as2 bs2)
Expand Down Expand Up @@ -248,12 +248,12 @@ instance (Unbox a,
= V_3 m_ (G.basicUnsafeSlice i_ m_ as)
(G.basicUnsafeSlice i_ m_ bs)
(G.basicUnsafeSlice i_ m_ cs)
{-# INLINE basicUnsafeIndexM #-}
basicUnsafeIndexM (V_3 _ as bs cs) i_
{-# INLINE basicUnsafeIndexM# #-}
basicUnsafeIndexM# (V_3 _ as bs cs) i_
= do
a <- G.basicUnsafeIndexM as i_
b <- G.basicUnsafeIndexM bs i_
c <- G.basicUnsafeIndexM cs i_
a <- G.basicUnsafeIndexM# as i_
b <- G.basicUnsafeIndexM# bs i_
c <- G.basicUnsafeIndexM# cs i_
return (a, b, c)
{-# INLINE basicUnsafeCopy #-}
basicUnsafeCopy (MV_3 _ as1 bs1 cs1) (V_3 _ as2 bs2 cs2)
Expand Down Expand Up @@ -448,13 +448,13 @@ instance (Unbox a,
(G.basicUnsafeSlice i_ m_ bs)
(G.basicUnsafeSlice i_ m_ cs)
(G.basicUnsafeSlice i_ m_ ds)
{-# INLINE basicUnsafeIndexM #-}
basicUnsafeIndexM (V_4 _ as bs cs ds) i_
{-# INLINE basicUnsafeIndexM# #-}
basicUnsafeIndexM# (V_4 _ as bs cs ds) i_
= do
a <- G.basicUnsafeIndexM as i_
b <- G.basicUnsafeIndexM bs i_
c <- G.basicUnsafeIndexM cs i_
d <- G.basicUnsafeIndexM ds i_
a <- G.basicUnsafeIndexM# as i_
b <- G.basicUnsafeIndexM# bs i_
c <- G.basicUnsafeIndexM# cs i_
d <- G.basicUnsafeIndexM# ds i_
return (a, b, c, d)
{-# INLINE basicUnsafeCopy #-}
basicUnsafeCopy (MV_4 _ as1 bs1 cs1 ds1) (V_4 _ as2
Expand Down Expand Up @@ -700,14 +700,14 @@ instance (Unbox a,
(G.basicUnsafeSlice i_ m_ cs)
(G.basicUnsafeSlice i_ m_ ds)
(G.basicUnsafeSlice i_ m_ es)
{-# INLINE basicUnsafeIndexM #-}
basicUnsafeIndexM (V_5 _ as bs cs ds es) i_
= do
a <- G.basicUnsafeIndexM as i_
b <- G.basicUnsafeIndexM bs i_
c <- G.basicUnsafeIndexM cs i_
d <- G.basicUnsafeIndexM ds i_
e <- G.basicUnsafeIndexM es i_
{-# INLINE basicUnsafeIndexM# #-}
basicUnsafeIndexM# (V_5 _ as bs cs ds es) i_
= do
a <- G.basicUnsafeIndexM# as i_
b <- G.basicUnsafeIndexM# bs i_
c <- G.basicUnsafeIndexM# cs i_
d <- G.basicUnsafeIndexM# ds i_
e <- G.basicUnsafeIndexM# es i_
return (a, b, c, d, e)
{-# INLINE basicUnsafeCopy #-}
basicUnsafeCopy (MV_5 _ as1 bs1 cs1 ds1 es1) (V_5 _ as2
Expand Down Expand Up @@ -1002,15 +1002,15 @@ instance (Unbox a,
(G.basicUnsafeSlice i_ m_ ds)
(G.basicUnsafeSlice i_ m_ es)
(G.basicUnsafeSlice i_ m_ fs)
{-# INLINE basicUnsafeIndexM #-}
basicUnsafeIndexM (V_6 _ as bs cs ds es fs) i_
= do
a <- G.basicUnsafeIndexM as i_
b <- G.basicUnsafeIndexM bs i_
c <- G.basicUnsafeIndexM cs i_
d <- G.basicUnsafeIndexM ds i_
e <- G.basicUnsafeIndexM es i_
f <- G.basicUnsafeIndexM fs i_
{-# INLINE basicUnsafeIndexM# #-}
basicUnsafeIndexM# (V_6 _ as bs cs ds es fs) i_
= do
a <- G.basicUnsafeIndexM# as i_
b <- G.basicUnsafeIndexM# bs i_
c <- G.basicUnsafeIndexM# cs i_
d <- G.basicUnsafeIndexM# ds i_
e <- G.basicUnsafeIndexM# es i_
f <- G.basicUnsafeIndexM# fs i_
return (a, b, c, d, e, f)
{-# INLINE basicUnsafeCopy #-}
basicUnsafeCopy (MV_6 _ as1 bs1 cs1 ds1 es1 fs1) (V_6 _ as2
Expand Down
7 changes: 4 additions & 3 deletions vector/src/Data/Vector.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
Expand Down Expand Up @@ -212,7 +213,7 @@ import qualified Control.Applicative as Applicative
import qualified Data.Foldable as Foldable
import qualified Data.Traversable as Traversable

import qualified GHC.Exts as Exts (IsList(..))
import qualified GHC.Exts as Exts (IsList(..), Int(..))


-- | Boxed vectors, supporting efficient slicing.
Expand Down Expand Up @@ -278,8 +279,8 @@ instance G.Vector Vector a where
{-# INLINE basicUnsafeSlice #-}
basicUnsafeSlice j n (Vector i _ arr) = Vector (i+j) n arr

{-# INLINE basicUnsafeIndexM #-}
basicUnsafeIndexM (Vector i _ arr) j = indexArrayM arr (i+j)
{-# INLINE basicUnsafeIndexM# #-}
basicUnsafeIndexM# (Vector i _ arr) j = indexArrayM arr (i + Exts.I# j)

{-# INLINE basicUnsafeCopy #-}
basicUnsafeCopy (MVector i n dst) (Vector j _ src)
Expand Down
28 changes: 15 additions & 13 deletions vector/src/Data/Vector/Fusion/Bundle/Monadic.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
Expand Down Expand Up @@ -98,10 +99,10 @@ import Control.Monad.Primitive

import qualified Data.List as List
import Data.Char ( ord )
import GHC.Base ( unsafeChr )
import GHC.Base ( unsafeChr, Int(..) )
import Control.Monad ( liftM )
import Prelude
( Eq, Ord, Num, Enum, Functor, Monad, Bool(..), Ordering, Char, Int, Word, Integer, Float, Double, Maybe(..), Either(..), Integral, RealFrac
( Eq, Ord, Num, Enum, Functor, Monad, Bool(..), Ordering, Char, Word, Integer, Float, Double, Maybe(..), Either(..), Integral, RealFrac
, return, fmap, otherwise, id, const, seq, max, maxBound, fromIntegral, truncate
, (+), (-), (<), (<=), (>), (>=), (==), (/=), (&&), (.), ($), (<$), (/) )

Expand Down Expand Up @@ -1091,9 +1092,10 @@ fromVector v = v `seq` n `seq` Bundle (Stream step 0)
n = basicLength v

{-# INLINE step #-}
step i | i >= n = return Done
| otherwise = case basicUnsafeIndexM v i of
Box x -> return $ Yield x (i+1)
step (I# i)
| I# i >= n = return Done
Copy link
Contributor

Choose a reason for hiding this comment

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

Here and everywhere in this PR, when we pattern-match on I# but then construct it again in function's body is GHC 100% guaranteed to optimise allocation of fresh Int away and will just reuse whatever we pattern-matched on?

Copy link
Contributor

Choose a reason for hiding this comment

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

GHC should be quite good at this. So it should eliminate Int allocations.

| otherwise = case basicUnsafeIndexM# v i of
Box x -> return $ Yield x (I# i + 1)


{-# INLINE vstep #-}
Expand All @@ -1112,10 +1114,10 @@ fromVectors us = Bundle (Stream pstep (Left us))
pstep (Left []) = return Done
pstep (Left (v:vs)) = basicLength v `seq` return (Skip (Right (v,0,vs)))

pstep (Right (v,i,vs))
| i >= basicLength v = return $ Skip (Left vs)
| otherwise = case basicUnsafeIndexM v i of
Box x -> return $ Yield x (Right (v,i+1,vs))
pstep (Right (v, I# i, vs))
| I# i >= basicLength v = return $ Skip (Left vs)
| otherwise = case basicUnsafeIndexM# v i of
Box x -> return $ Yield x (Right (v, I# i + 1, vs))

-- FIXME: work around bug in GHC 7.6.1
vstep :: HasCallStack => [v a] -> m (Step [v a] (Chunk v a))
Expand Down Expand Up @@ -1143,10 +1145,10 @@ concatVectors Bundle{sElems = Stream step t}
Skip s' -> return (Skip (Left s'))
Done -> return Done

pstep (Right (v,i,s))
| i >= basicLength v = return (Skip (Left s))
| otherwise = case basicUnsafeIndexM v i of
Box x -> return (Yield x (Right (v,i+1,s)))
pstep (Right (v, I# i, s))
| I# i >= basicLength v = return (Skip (Left s))
| otherwise = case basicUnsafeIndexM# v i of
Box x -> return (Yield x (Right (v, I# i + 1,s)))


vstep s = do
Expand Down
26 changes: 14 additions & 12 deletions vector/src/Data/Vector/Generic.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
Expand Down Expand Up @@ -211,6 +212,7 @@ import Data.Typeable ( Typeable, gcast1 )
import Data.Data ( Data, DataType, Constr, Fixity(Prefix),
mkDataType, mkConstr, constrIndex, mkNoRepType )
import qualified Data.Traversable as T (Traversable(mapM))
import GHC.Exts (Int(..))

-- Length information
-- ------------------
Expand Down Expand Up @@ -246,16 +248,16 @@ infixl 9 !
(!) :: (HasCallStack, Vector v a) => v a -> Int -> a
{-# INLINE_FUSED (!) #-}
-- See NOTE: [Strict indexing]
(!) v !i = checkIndex Bounds i (length v) $ unBox (basicUnsafeIndexM v i)
(!) v (I# i) = checkIndex Bounds (I# i) (length v) $ unBox (basicUnsafeIndexM# v i)

infixl 9 !?
-- | O(1) Safe indexing.
(!?) :: Vector v a => v a -> Int -> Maybe a
{-# INLINE_FUSED (!?) #-}
-- See NOTE: [Strict indexing]
-- Use basicUnsafeIndexM @Box to perform the indexing eagerly.
v !? (!i)
| i `inRange` length v = case basicUnsafeIndexM v i of Box a -> Just a
-- Use basicUnsafeIndexM# @Box to perform the indexing eagerly.
v !? (I# i)
| I# i `inRange` length v = case basicUnsafeIndexM# v i of Box a -> Just a
| otherwise = Nothing


Expand All @@ -273,7 +275,7 @@ last v = v ! (length v - 1)
unsafeIndex :: Vector v a => v a -> Int -> a
{-# INLINE_FUSED unsafeIndex #-}
-- See NOTE: [Strict indexing]
unsafeIndex v !i = checkIndex Unsafe i (length v) $ unBox (basicUnsafeIndexM v i)
unsafeIndex v (I# i) = checkIndex Unsafe (I# i) (length v) $ unBox (basicUnsafeIndexM# v i)

-- | /O(1)/ First element, without checking if the vector is empty.
unsafeHead :: Vector v a => v a -> a
Expand Down Expand Up @@ -333,7 +335,7 @@ unsafeLast v = unsafeIndex v (length v - 1)
-- element) is evaluated eagerly.
indexM :: (HasCallStack, Vector v a, Monad m) => v a -> Int -> m a
{-# INLINE_FUSED indexM #-}
indexM v !i = checkIndex Bounds i (length v) $ liftBox $ basicUnsafeIndexM v i
indexM v (I# i) = checkIndex Bounds (I# i) (length v) $ liftBox $ basicUnsafeIndexM# v i

-- | /O(1)/ First element of a vector in a monad. See 'indexM' for an
-- explanation of why this is useful.
Expand All @@ -351,9 +353,9 @@ lastM v = indexM v (length v - 1)
-- explanation of why this is useful.
unsafeIndexM :: (Vector v a, Monad m) => v a -> Int -> m a
{-# INLINE_FUSED unsafeIndexM #-}
unsafeIndexM v !i = checkIndex Unsafe i (length v)
unsafeIndexM v (I# i) = checkIndex Unsafe (I# i) (length v)
$ liftBox
$ basicUnsafeIndexM v i
$ basicUnsafeIndexM# v i

-- | /O(1)/ First element in a monad, without checking for empty vectors.
-- See 'indexM' for an explanation of why this is useful.
Expand Down Expand Up @@ -1011,7 +1013,7 @@ backpermute v is = seq v
-- NOTE: we do it this way to avoid triggering LiberateCase on n in
-- polymorphic code
index :: HasCallStack => Int -> Box a
index !i = checkIndex Bounds i n $ basicUnsafeIndexM v i
index (I# i) = checkIndex Bounds (I# i) n $ basicUnsafeIndexM# v i

-- | Same as 'backpermute', but without bounds checking.
unsafeBackpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
Expand All @@ -1028,7 +1030,7 @@ unsafeBackpermute v is = seq v
{-# INLINE index #-}
-- NOTE: we do it this way to avoid triggering LiberateCase on n in
-- polymorphic code
index !i = checkIndex Unsafe i n $ basicUnsafeIndexM v i
index (I# i) = checkIndex Unsafe (I# i) n $ basicUnsafeIndexM# v i

-- Safe destructive updates
-- ------------------------
Expand Down Expand Up @@ -2552,9 +2554,9 @@ streamR v = v `seq` n `seq` (Bundle.unfoldr get n `Bundle.sized` Exact n)

{-# INLINE get #-}
get 0 = Nothing
get i = let !i' = i-1
get i = let !(I# i') = i-1
in
case basicUnsafeIndexM v i' of Box x -> Just (x, i')
case basicUnsafeIndexM# v i' of Box x -> Just (x, I# i')

-- | /O(n)/ Construct a vector from a 'Bundle', proceeding from right to left.
unstreamR :: Vector v a => Bundle v a -> v a
Expand Down
23 changes: 16 additions & 7 deletions vector/src/Data/Vector/Generic/Base.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
Expand Down Expand Up @@ -34,6 +35,7 @@ import qualified Data.Primitive.PrimArray as Prim

import Control.Monad.ST
import Data.Kind (Type)
import GHC.Exts (Int(..), Int#)

-- | @Mutable v s a@ is the mutable version of the immutable vector type @v a@ with
-- the state token @s@. It is injective on GHC 8 and newer.
Expand Down Expand Up @@ -112,7 +114,13 @@ class MVector (Mutable v) a => Vector v a where
--
-- which does not have this problem, because indexing (but not the returned
-- element!) is evaluated immediately.
basicUnsafeIndexM :: v a -> Int -> Box a
basicUnsafeIndexM :: v a -> Int -> Box a
basicUnsafeIndexM xs (I# i) = basicUnsafeIndexM# xs i
{-# INLINE basicUnsafeIndexM #-}

basicUnsafeIndexM# :: v a -> Int# -> Box a
basicUnsafeIndexM# xs i = basicUnsafeIndexM xs (I# i)
{-# INLINE basicUnsafeIndexM# #-}

-- | /Assumed complexity: O(n)/
--
Expand All @@ -131,11 +139,12 @@ class MVector (Mutable v) a => Vector v a where
where
!n = basicLength src

do_copy i | i < n = do
x <- liftBox $ basicUnsafeIndexM src i
M.basicUnsafeWrite dst i x
do_copy (i+1)
| otherwise = return ()
do_copy (I# i)
| I# i < n = do
x <- liftBox $ basicUnsafeIndexM# src i
M.basicUnsafeWrite dst (I# i) x
do_copy (I# i + 1)
| otherwise = return ()

-- | Evaluate @a@ as far as storing it in a vector would and yield @b@.
-- The @v a@ argument only fixes the type and is not touched. This method is
Expand All @@ -152,4 +161,4 @@ class MVector (Mutable v) a => Vector v a where
elemseq _ = \_ x -> x

{-# MINIMAL basicUnsafeFreeze, basicUnsafeThaw, basicLength,
basicUnsafeSlice, basicUnsafeIndexM #-}
basicUnsafeSlice, (basicUnsafeIndexM | basicUnsafeIndexM#) #-}
5 changes: 3 additions & 2 deletions vector/src/Data/Vector/Primitive.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RoleAnnotations #-}
Expand Down Expand Up @@ -249,8 +250,8 @@ instance Prim a => G.Vector Vector a where
{-# INLINE basicUnsafeSlice #-}
basicUnsafeSlice j n (Vector i _ arr) = Vector (i+j) n arr

{-# INLINE basicUnsafeIndexM #-}
basicUnsafeIndexM (Vector i _ arr) j = return $! indexByteArray arr (i+j)
{-# INLINE basicUnsafeIndexM# #-}
basicUnsafeIndexM# (Vector i _ arr) j = return $! indexByteArray arr (i + Exts.I# j)

{-# INLINE basicUnsafeCopy #-}
basicUnsafeCopy (MVector i n dst) (Vector j _ src)
Expand Down
Loading
Loading