atrophy (empty) → 0.1.0.0
raw patch · 10 files changed
+663/−0 lines, 10 filesdep +HUnitdep +QuickCheckdep +atrophy
Dependencies added: HUnit, QuickCheck, atrophy, base, contiguous, deepseq, quickcheck-classes, random, tasty, tasty-bench, tasty-hunit, tasty-quickcheck, wide-word
Files
- CHANGELOG.md +5/−0
- LICENSE +20/−0
- atrophy.cabal +68/−0
- bench/Main.hs +87/−0
- src/Atrophy.hs +12/−0
- src/Atrophy/Internal.hs +173/−0
- src/Atrophy/Internal/LongDivision.hs +84/−0
- src/Atrophy/LongDivision.hs +46/−0
- src/Atrophy/LongMultiplication.hs +90/−0
- tests/Main.hs +78/−0
+ CHANGELOG.md view
@@ -0,0 +1,5 @@+# Revision history for atrophy++## 0.1.0.0 -- YYYY-mm-dd++* First version. Released on an unsuspecting world.
+ LICENSE view
@@ -0,0 +1,20 @@+Copyright (c) 2022 Zachary Churchill++Permission is hereby granted, free of charge, to any person obtaining+a copy of this software and associated documentation files (the+"Software"), to deal in the Software without restriction, including+without limitation the rights to use, copy, modify, merge, publish,+distribute, sublicense, and/or sell copies of the Software, and to+permit persons to whom the Software is furnished to do so, subject to+the following conditions:++The above copyright notice and this permission notice shall be included+in all copies or substantial portions of the Software.++THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.+IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY+CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,+TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE+SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ atrophy.cabal view
@@ -0,0 +1,68 @@+cabal-version: 2.4+name: atrophy+version: 0.1.0.0+synopsis: Faster integer division and modulus operations +description:+ Fast div/mod via arithmetic strength reduction.++ Good compilers already perform this optimization for divisors that are known at compile time; this library enables this optimization for divisors that are only known at runtime.+license: MIT+license-file: LICENSE+author: Zachary Churchill <zacharyachurchill@gmail.com>+maintainer: Zachary Churchill <zacharyachurchill@gmail.com>+copyright: 2022 Zachary Churchill+category: Math+extra-source-files: CHANGELOG.md++source-repository head+ type: git+ location: https://github.com/goolord/atrophy/++library+ exposed-modules:+ Atrophy+ , Atrophy.LongDivision+ , Atrophy.LongMultiplication+ other-modules:+ Atrophy.Internal+ Atrophy.Internal.LongDivision+ -- other-extensions:+ ghc-options: -Wall -O2 -Wredundant-constraints+ build-depends:+ base >=4.10.0.0 && <5+ , wide-word+ , contiguous >= 0.6.0.0+ hs-source-dirs: src+ default-language: Haskell2010++test-suite atrophy-test+ default-language: Haskell2010+ type: exitcode-stdio-1.0+ hs-source-dirs: tests+ main-is: Main.hs+ build-depends:+ base >=4.10.0.0 && <5+ , wide-word+ , contiguous >= 0.6.0.0+ , HUnit+ , tasty+ , QuickCheck+ , quickcheck-classes+ , tasty+ , tasty-hunit+ , tasty-quickcheck+ , atrophy++benchmark atrophy-bench+ type: exitcode-stdio-1.0+ build-depends:+ base >=4.10.0.0 && <5+ , tasty+ , tasty-bench+ , deepseq+ , atrophy+ , random+ ghc-options: -Wall -O2+ default-language: Haskell2010+ hs-source-dirs: bench+ main-is: Main.hs
+ bench/Main.hs view
@@ -0,0 +1,87 @@+{-# LANGUAGE+ BangPatterns+ , GADTs+ , DeriveGeneric+ , StandaloneDeriving+ , MagicHash+ , DataKinds+ , GeneralizedNewtypeDeriving+ , TypeApplications+ , ScopedTypeVariables+ , NumericUnderscores+#-}++{-# OPTIONS_GHC+ -fno-warn-orphans+#-}+{-# LANGUAGE DerivingStrategies #-}++module Main where++import Test.Tasty.Bench (bench, bgroup, defaultMain, nf, Benchmark, envWithCleanup, nfIO)+import Control.DeepSeq (NFData, force)+import Atrophy+import GHC.Generics+import Data.Word (Word64, Word32)+import Control.Exception (evaluate)+import Test.Tasty (withResource)+import System.Mem (performMajorGC)+import System.Random.Stateful+import Data.Proxy (Proxy (Proxy))++deriving instance Generic StrengthReducedW64+instance NFData StrengthReducedW64++deriving instance Generic StrengthReducedW32+instance NFData StrengthReducedW32++instance (Bounded a, Num a, UniformRange a) => Uniform (NonZero a) where+ uniformM g = NonZero <$> uniformRM (1, maxBound) g++deriving newtype instance NFData a => NFData (NonZero a)++manyRandom :: forall a. (Uniform a, NFData a) => IO [a]+manyRandom = uniformListM 10_000 globalStdGen++randomEnv :: NFData b => IO b+ -> (b -> Benchmark)+ -> Benchmark+randomEnv a = envWithCleanup (a >>= evaluate . force) (const performMajorGC)++randomEnv' :: NFData b => IO b+ -> (IO b -> Benchmark)+ -> Benchmark+randomEnv' a = withResource (a >>= evaluate . force) (const performMajorGC)++main :: IO ()+main = do+ defaultMain $+ [ bgroup "atrophy"+ [ bgroup "Word64" $ atrophyBench div64 new64+ , bgroup "Word32" $ atrophyBench div' (new StrengthReducedW32)+ ]+ , bgroup "ghc"+ [ bgroup "Word64" $ ghcBench (Proxy @Word64)+ , bgroup "Word32" $ ghcBench (Proxy @Word32)+ ]+ ]++ghcBench :: forall a. (NFData a, Uniform a, Integral a, Random a, Bounded a) => Proxy a -> [Benchmark]+ghcBench _ =+ [ randomEnv (manyRandom @(a, a)) $ \somePairs ->+ bench "div 10000 uniques" $ nf (\xs -> fmap (\(x, y) -> x `div` y) xs) somePairs+ , randomEnv' ((,) <$> randomRIO (1, maxBound) <*> manyRandom @a) $ \x -> bench "div 10000, 1 divisor" $ nfIO $ do + (divisor', dividends) <- x+ pure $ fmap (\dividend' -> dividend' `div` divisor') dividends+ ]++atrophyBench :: forall base sr out. (NFData base, Uniform base, UniformRange base, NFData sr, NFData out, Random base, Num base, Bounded base) => (base -> sr -> out) -> (NonZero base -> sr) -> [Benchmark]+atrophyBench divF newF =+ [ randomEnv (uniformM globalStdGen) $ \divisor' ->+ bench "new" $ nf newF divisor'+ , randomEnv (manyRandom @(base, base)) $ \somePairs ->+ bench "div 10000 uniques" $ nf (fmap (\(x, y) -> x `divF` newF (NonZero y))) somePairs+ , randomEnv' ((,) <$> (newF <$> (NonZero <$> randomRIO (1, maxBound))) <*> manyRandom @base) $ \x -> bench "div 10000, 1 divisor" $ nfIO $ do + (divisor', dividends) <- x+ pure $ fmap (\dividend' -> dividend' `divF` divisor') dividends+ ]
+ src/Atrophy.hs view
@@ -0,0 +1,12 @@+module Atrophy+ ( module REXPORT+ , StrengthReducedW128 (..)+ , StrengthReducedW64 (..)+ , StrengthReducedW32 (..)+ , StrengthReducedW16 (..)+ , StrengthReducedW8 (..)+ ) where++import Atrophy.LongDivision as REXPORT+import Atrophy.LongMultiplication as REXPORT+import Atrophy.Internal as REXPORT
+ src/Atrophy/Internal.hs view
@@ -0,0 +1,173 @@+{-# LANGUAGE+ CPP+ , TypeApplications+ , DataKinds+ , FlexibleContexts+ , DuplicateRecordFields+ , TypeFamilies+ , BangPatterns+ , NumericUnderscores+ , ScopedTypeVariables+ , DerivingStrategies+ , GeneralizedNewtypeDeriving+#-}++module Atrophy.Internal where++import Data.WideWord.Word128+import Data.Bits+import Atrophy.Internal.LongDivision+import GHC.Records+import Data.Word++newtype NonZero a = NonZero a+ deriving newtype (Num, Show)++instance (Bounded a, Num a) => Bounded (NonZero a) where+ minBound = 1+ maxBound = NonZero maxBound++{-# INLINE isPowerOf2 #-}+isPowerOf2 :: (Bits a, Num a) => a -> Bool+isPowerOf2 x = (x .&. (x - 1)) == 0++{-# INLINE new64 #-}+new64 :: NonZero Word64 -> StrengthReducedW64+new64 (NonZero divi) =+ if isPowerOf2 divi+ then StrengthReducedW64 0 divi+ else+ let quotient = divide128MaxBy64 $ fromIntegral divi+ in StrengthReducedW64 (quotient + 1) divi++{-# INLINE divRem64 #-}+divRem64 ::+ ( HasField "divisor" strRed a+ , HasField "multiplier" strRed Word128+ , Integral a+ , FiniteBits a+ ) => a -> strRed -> (a, a)+divRem64 dividend divis =+ case getField @"multiplier" divis of+ 0 ->+ let+ quotient = dividend `unsafeShiftR` (countTrailingZeros $ getField @"divisor" divis)+ remainder = dividend .&. (getField @"divisor" divis - 1)+ in (quotient, remainder)+ multiplier' ->+ let+ numerator128 = fromIntegral @_ @Word128 dividend+ multipliedHi = numerator128 * (upper128 multiplier')+ multipliedLo = upper128 (numerator128 * (lower128 multiplier'))++ quotient = fromIntegral (upper128 (multipliedHi + multipliedLo))+ remainder = dividend - quotient * getField @"divisor" divis+ in (quotient, remainder)++{-# INLINE divRem #-}+{-# SPECIALIZE divRem :: Word32 -> StrengthReducedW32 -> (Word32, Word32) #-}+divRem :: forall strRed a b.+ ( HasField "divisor" strRed a+ , HasField "multiplier" strRed b+ , Integral a+ , FiniteBits a, Integral b, FiniteBits (Half b), Bits b) => a -> strRed -> (a, a)+divRem dividend divis =+ case getField @"multiplier" divis of+ 0 ->+ let+ quotient = dividend `unsafeShiftR` (countTrailingZeros $ getField @"divisor" divis)+ remainder = dividend .&. (getField @"divisor" divis - 1)+ in (quotient, remainder)+ multiplier' ->+ let+ numerator64 = fromIntegral @_ @b dividend+ multipliedHi = numerator64 * (upperHalf multiplier')+ multipliedLo = upperHalf (numerator64 * (lowerHalf multiplier'))++ quotient = fromIntegral (upperHalf (multipliedHi + multipliedLo))+ remainder = dividend - quotient * getField @"divisor" divis+ in (quotient, remainder)++{-# INLINE new #-}+{-# SPECIALIZE new :: (Word64 -> Word32 -> StrengthReducedW32) -> NonZero Word32 -> StrengthReducedW32 #-}+new :: (Bits t, Integral t, Bounded (Multiplier t), Integral (Multiplier t)) =>((Multiplier t) -> t -> a) -> (NonZero t) -> a+new con (NonZero divi) =+ if isPowerOf2 divi+ then con 0 divi+ else+ let quotient = maxBound `div` fromIntegral divi+ in con (quotient + 1) divi++{-# INLINE div64 #-}+{-# SPECIALIZE div64 :: Word64 -> StrengthReducedW64 -> Word64 #-}+div64 :: (HasField "divisor" r b, HasField "multiplier" r Word128,+ Integral b, FiniteBits b) =>+ b -> r -> b+div64 a rhs = fst $ divRem64 a rhs++{-# INLINE rem64 #-}+{-# SPECIALIZE rem64 :: Word64 -> StrengthReducedW64 -> Word64 #-}+rem64 :: (HasField "divisor" r b, HasField "multiplier" r Word128,+ Integral b, FiniteBits b) =>+ b -> r -> b+rem64 a rhs = snd $ divRem64 a rhs++{-# INLINE div' #-}+{-# SPECIALIZE div' :: Word32 -> StrengthReducedW32 -> Word32 #-}+div' ::+ ( HasField "divisor" strRed b+ , HasField "multiplier" strRed w+ , Integral b, FiniteBits b, Integral w, FiniteBits (Half w), Bits w) => b -> strRed -> b+div' a rhs = fst $ divRem a rhs++{-# INLINE rem' #-}+{-# SPECIALIZE rem' :: Word32 -> StrengthReducedW32 -> Word32 #-}+rem' ::+ ( HasField "divisor" strRed b+ , HasField "multiplier" strRed w+ , Integral b, FiniteBits b, Integral w, FiniteBits (Half w), Bits w+ ) => b -> strRed -> b+rem' a rhs = snd $ divRem a rhs++{-# INLINE lower128 #-}+lower128 :: Word128 -> Word128+lower128 (Word128 _hi low) = Word128 0 low++{-# INLINE upper128 #-}+upper128 :: Word128 -> Word128+upper128 (Word128 hi _low) = Word128 0 hi++{-# INLINE lowerHalf #-}+lowerHalf :: forall w. ( FiniteBits (Half w), Bits w) =>w -> w+lowerHalf w = (w `unsafeShiftL` halfSize) `unsafeShiftR` halfSize+ where+ halfSize = finiteBitSize @(Half w) zeroBits++{-# INLINE upperHalf #-}+upperHalf :: forall w. ( Bits w, FiniteBits (Half w)) =>w -> w+upperHalf w = w `unsafeShiftR` halfSize+ where+ halfSize = finiteBitSize @(Half w) zeroBits++type family Multiplier a where+ Multiplier Word64 = Word128+ Multiplier Word32 = Word64+ Multiplier Word16 = Word32+ Multiplier Word8 = Word16++type family Half a where+ Half Word128 = Word64+ Half Word64 = Word32+ Half Word32 = Word16+ Half Word16 = Word8++data StrengthReducedW64 = StrengthReducedW64 { multiplier :: {-# UNPACK #-} !Word128, divisor :: {-# UNPACK #-} !Word64 }+data StrengthReducedW32 = StrengthReducedW32 { multiplier :: {-# UNPACK #-} !Word64, divisor :: {-# UNPACK #-} !Word32 }+data StrengthReducedW16 = StrengthReducedW16 { multiplier :: {-# UNPACK #-} !Word32, divisor :: {-# UNPACK #-} !Word16 }+data StrengthReducedW8 = StrengthReducedW7 { multiplier :: {-# UNPACK #-} !Word16, divisor :: {-# UNPACK #-} !Word8 }++data StrengthReducedW128 = StrengthReducedW128+ { multiplierHi :: {-#UNPACK #-} !Word128+ , multiplierLo :: {-#UNPACK #-} !Word128+ , divisor :: {-#UNPACK #-} !Word128+ }
+ src/Atrophy/Internal/LongDivision.hs view
@@ -0,0 +1,84 @@+{-# LANGUAGE+ TypeApplications+ , ScopedTypeVariables+ , LambdaCase+ , NumericUnderscores+#-}++module Atrophy.Internal.LongDivision where++import Data.WideWord.Word128+import Data.Word+import Data.Bits++-- divides a 128-bit number by a 64-bit divisor, returning the quotient as a 64-bit number+-- assumes that the divisor and numerator have both already been bit-shifted so that countLeadingZeros divisor == 0+{-# INLINE divide128By64Preshifted #-}+divide128By64Preshifted :: Word64 -> Word64 -> Word64 -> Word64+divide128By64Preshifted numeratorHi numeratorLo' divisor =+ let+ numeratorMid = fromIntegral @Word64 @Word128 (numeratorLo' `unsafeShiftR` 32)+ numeratorLo = fromIntegral @Word32 @Word128 (fromIntegral @Word64 @Word32 numeratorLo')+ divisorFull128 = fromIntegral @Word64 @Word128 divisor+ divisorHi = divisor `unsafeShiftR` 32++ -- To get the upper 32 bits of the quotient, we want to divide 'fullUpperNumerator' by 'divisor'+ -- but the problem is, fullUpperNumerator is a 96-bit number, meaning we would need to use u128 to do the division all at once, and the whole point of this is that we don't want to do 128 bit divison because it's slow+ -- so instead, we'll shift both the numerator and divisor right by 32, giving us a 64 bit / 32 bit division. This won't give us the exact quotient -- but it will be close.+ fullUpperNumerator = (Word128 0 numeratorHi `unsafeShiftL` 32) .|. numeratorMid+ quotientHi :: Word64+ quotientHi = min (numeratorHi `div` divisorHi) (fromIntegral $ maxBound @Word32)+ productHi = (Word128 0 quotientHi) * divisorFull128++ -- quotientHi contains our guess at what the quotient is! the problem is that we got this by ignoring the lower 32 bits of the divisor. when we account for that, the quotient might be slightly lower+ -- we will know our quotient is too high if quotient * divisor > numerator. if it is, decrement until it's in range+ (productHi', quotientHi') = clampToFull productHi quotientHi divisorFull128 fullUpperNumerator++ remainderHi = fullUpperNumerator - productHi'++ -- repeat the process using the lower half of the numerator+ fullLowerNumerator = (remainderHi `unsafeShiftL` 32) .|. numeratorLo++ quotientLo = min ((fromIntegral @_ @Word64 remainderHi) `div` divisorHi) (fromIntegral $ maxBound @Word32)+ productLo = (Word128 0 quotientLo) * divisorFull128++ -- again, quotientLo is just a guess at this point, it might be slightly too large+ (_, quotientLo') = clampToFull productLo quotientLo divisorFull128 fullLowerNumerator++ -- We now have our separate quotients, now we just have to add them together+ in (quotientHi' `unsafeShiftL` 32) .|. quotientLo'++divide128MaxBy64 :: Word64 -> Word128+divide128MaxBy64 divisor =+ let+ quotientHi = maxBound @Word64 `div` divisor;+ remainderHi = maxBound @Word64 - quotientHi * divisor;++ leadingZeros = countLeadingZeros divisor+ quotientLo = if leadingZeros >= 32+ then+ let+ numeratorMid = (remainderHi `unsafeShiftL` 32) .|. (fromIntegral (maxBound @Word32))+ quotientMid = numeratorMid `div` divisor;+ remainderMid = numeratorMid - quotientMid * divisor;++ numeratorLo = (remainderMid `unsafeShiftL` 32) .|. (fromIntegral (maxBound @Word32))+ quotientLo' = numeratorLo `div` divisor++ in (quotientMid `unsafeShiftL` 32) .|. quotientLo'+ else+ let+ numeratorHi = if leadingZeros > 0+ then (remainderHi `unsafeShiftL` leadingZeros) .|. (maxBound @Word64 `unsafeShiftR` (64 - leadingZeros))+ else remainderHi+ numeratorLo = maxBound @Word64 `unsafeShiftL` leadingZeros;+ in divide128By64Preshifted numeratorHi numeratorLo (divisor `unsafeShiftL` leadingZeros)+ in ((fromIntegral quotientHi) `unsafeShiftL` 64) .|. (fromIntegral quotientLo)++clampToFull :: Word128 -> Word64 -> Word128 -> Word128 -> (Word128, Word64)+clampToFull product' quotient' divisorFull128 fullUpperNumerator = go product' quotient'+ where+ go prod quotient =+ if prod > fullUpperNumerator+ then go (prod - divisorFull128) (quotient - 1)+ else (prod, quotient)
+ src/Atrophy/LongDivision.hs view
@@ -0,0 +1,46 @@+{-# LANGUAGE+ TypeApplications+ , ScopedTypeVariables+ , LambdaCase+ , NumericUnderscores+#-}++module Atrophy.LongDivision+ ( module X+ , module Atrophy.LongDivision+ )+ where++import Data.Word+import Atrophy.Internal.LongDivision as X+import Atrophy.Internal+import qualified Data.Primitive.Contiguous as Contiguous+import Data.Primitive.Contiguous (PrimArray, Mutable, Sliced)+import Control.Monad.ST.Strict (ST)+import Data.STRef.Strict (newSTRef, readSTRef, writeSTRef)+import Data.Bits++{-# NOINLINE longDivision #-}+longDivision :: forall s. Sliced PrimArray Word64 -> StrengthReducedW64 -> Mutable PrimArray s Word64 -> ST s ()+longDivision numeratorSlice reducedDivisor quotient = do+ remainder <- newSTRef 0+ (flip Contiguous.itraverse_) numeratorSlice $ \i numerator -> do+ readSTRef remainder >>= \case+ 0 -> do+ -- The remainder is zero, which means we can take a shortcut and only do a single division!+ let (digitQuotient, digitRemainder) = divRem numerator reducedDivisor++ Contiguous.write quotient i digitQuotient+ writeSTRef remainder digitRemainder++ remainder' -> do+ -- Do one division that includes the running remainder and the upper half of this numerator element,+ -- then a second division for the first division's remainder combinedwith the lower half+ let upperNumerator = (remainder' `unsafeShiftL` 32) .|. (numerator `unsafeShiftR` 32)+ let (upperQuotient, upperRemainder) = divRem upperNumerator reducedDivisor++ let lowerNumerator = (upperRemainder `unsafeShiftL` 32) .|. (0x00000000_ffffffff .&. numerator)+ let (lowerQuotient, lowerRemainder) = divRem lowerNumerator reducedDivisor++ Contiguous.write quotient i $ (upperQuotient `unsafeShiftL` 32) .|. lowerQuotient+ writeSTRef remainder lowerRemainder
+ src/Atrophy/LongMultiplication.hs view
@@ -0,0 +1,90 @@+{-# LANGUAGE+ TypeApplications+ , ScopedTypeVariables+ , LambdaCase+#-}++module Atrophy.LongMultiplication where++import Data.WideWord.Word128+import Data.Word+import qualified Data.Primitive.Contiguous as Contiguous+import Data.Primitive.Contiguous (PrimArray, MutableSliced, Mutable)+import Control.Monad.ST.Strict (ST)+import Data.STRef.Strict (newSTRef, modifySTRef, readSTRef)+import Data.Bits+import Data.Foldable (for_)++{-# INLINE multiply256By128UpperBits #-}+multiply256By128UpperBits :: Word128 -> Word128 -> Word128 -> Word128+multiply256By128UpperBits aHi aLo b =+ let+ -- Break a and b into little-endian 64-bit chunks+ aChunks :: PrimArray Word64+ aChunks = Contiguous.quadrupleton+ (word128Lo64 aLo)+ (word128Hi64 aLo)+ (word128Lo64 aHi)+ (word128Hi64 aHi)+ bChunks :: PrimArray Word64+ bChunks = Contiguous.doubleton+ (word128Lo64 b)+ (word128Hi64 b)++ -- Multiply b by a, one chunk of b at a time+ prod :: PrimArray Word64+ prod = Contiguous.create $ do+ prod' <- Contiguous.replicateMut 6 0+ flip Contiguous.itraverse_ bChunks $ \bIndex bDigit -> do+ pSize <- Contiguous.sizeMut prod'+ multiply256By64Helper+ (Contiguous.sliceMut prod' bIndex (pSize - bIndex))+ aChunks+ bDigit+ pure prod'++ in Word128+ { word128Hi64 = Contiguous.index prod 5+ , word128Lo64 = Contiguous.index prod 4+ }++{-# INLINE multiply256By64Helper #-}+multiply256By64Helper :: forall s. MutableSliced PrimArray s Word64 -> PrimArray Word64 -> Word64 -> ST s ()+multiply256By64Helper _ _ 0 = pure ()+multiply256By64Helper prod a b = do+ carry <- newSTRef 0+ productSize <- Contiguous.sizeMut prod+ let+ aSize = Contiguous.size a+ productLo :: MutableSliced PrimArray s Word64+ productLo = Contiguous.sliceMut prod 0 aSize+ productHi :: MutableSliced PrimArray s Word64+ productHi = Contiguous.sliceMut prod aSize (productSize - aSize)+ -- Multiply each of the digits in a by b, adding them into the 'prod' value.+ -- We don't zero out prod, because we this will be called multiple times, so it probably contains a previous iteration's partial prod, and we're adding + carrying on top of it+ for_ [0..aSize - 1] $ \i -> do+ p <- Contiguous.read productLo i+ let aDigit = Contiguous.index a i+ modifySTRef carry $ \x -> x+ + Word128 0 p+ + (Word128 0 aDigit * Word128 0 b)+ Contiguous.write prod i . word128Lo64 =<< readSTRef carry+ modifySTRef carry (`unsafeShiftR` 64)++ let productHiSize = productSize - aSize+ for_ [0..productHiSize - 1] $ \i -> do+ p <- Contiguous.read productHi i+ modifySTRef carry (+ Word128 0 p)+ Contiguous.write productHi i . word128Lo64 =<< readSTRef carry+ modifySTRef carry (`unsafeShiftR` 64)++ readSTRef carry >>= \case+ 0 -> pure ()+ _ -> error "carry overflow during multiplication!"++-- compute prod += a * b+{-# INLINE longMultiply #-}+longMultiply :: forall s. PrimArray Word64 -> Word64 -> Mutable PrimArray s Word64 -> ST s ()+longMultiply a b prod = do+ prod' <- Contiguous.toSliceMut prod+ multiply256By64Helper prod' a b
+ tests/Main.hs view
@@ -0,0 +1,78 @@+{-# LANGUAGE TypeOperators #-}++{-# OPTIONS_GHC -Wall #-}+{-# OPTIONS_GHC -fno-warn-orphans #-}+{-# OPTIONS_GHC -Wno-unused-top-binds #-}+{-# LANGUAGE StandaloneDeriving #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE DerivingVia #-}++module Main (main) where++import Test.QuickCheck hiding (NonZero)+import Test.Tasty+import Test.Tasty.QuickCheck hiding (NonZero)+import qualified Test.Tasty.QuickCheck as QC (NonZero(..))+import Data.WideWord.Word128+import Data.WideWord.Word256+import Atrophy+import Data.Word++main :: IO ()+main = do+ defaultMain tests++tests :: TestTree+tests =+ testGroup "Tests" [unitTests]++instance Arbitrary Word128 where+ arbitrary = Word128 <$> arbitrary <*> arbitrary++deriving via (QC.NonZero a) instance (Num a, Eq a, Arbitrary a) => Arbitrary (NonZero a)++-- wrong+naiveMultiply256By128UpperBits :: Word128 -> Word128 -> Word128 -> Word128+naiveMultiply256By128UpperBits aHi aLo b =+ let+ a' = Word256+ { word256hi = word128Hi64 aHi+ , word256m1 = word128Lo64 aHi+ , word256m0 = word128Hi64 aLo+ , word256lo = word128Lo64 aLo+ }+ b' = Word256+ { word256hi = word128Hi64 b+ , word256m1 = word128Lo64 b+ , word256m0 = 0+ , word256lo = 0+ }+ Word256 h l _ _ = a' * b'+ in Word128 h l++unitTests :: TestTree+unitTests = testGroup "Unit tests"+ [ testGroup "Long multiplication"+ [ -- testProperty "multiply256By128UpperBits" $ equivalentOnArbitrary3 multiply256By128UpperBits naiveMultiply256By128UpperBits+ ]+ , testGroup "Long division"+ [ testGroup "StrengthReducedW64" + [ testProperty "div64" $ \(a, b) -> ourDiv64 a b === theirDiv a b+ , testProperty "div32" $ \(a, b) -> ourDiv32 a b === theirDiv a b+ ]+ ]+ ]++ourDiv64 :: NonZero Word64 -> NonZero Word64 -> Word64+ourDiv64 (NonZero dividend) divi =+ let sr = new64 divi+ in div64 dividend sr++ourDiv32 :: NonZero Word32 -> NonZero Word32 -> Word32+ourDiv32 (NonZero dividend) divi =+ let sr = new StrengthReducedW32 divi+ in div' dividend sr++theirDiv :: Integral a => NonZero a -> NonZero a -> a+theirDiv (NonZero dividend) (NonZero divi) =+ dividend `div` divi