diff --git a/cbits/simd.h b/cbits/simd.h --- a/cbits/simd.h +++ b/cbits/simd.h @@ -15,6 +15,13 @@ size_t target_length, uint8_t *source); +void avx2_cmpeq8_para( + uint8_t *bytes, + size_t bytes_length, + uint8_t **targets, + size_t targets_length, + uint8_t *source); + void avx2_and_bits( uint8_t *target, size_t target_length, diff --git a/cbits/simd_avx2.c b/cbits/simd_avx2.c --- a/cbits/simd_avx2.c +++ b/cbits/simd_avx2.c @@ -54,6 +54,33 @@ #endif } +void avx2_cmpeq8_para( + uint8_t *bytes, + size_t bytes_length, + uint8_t **targets, + size_t targets_length, + uint8_t *source) { +#if defined(AVX2_ENABLED) + size_t i; + + for (i = 0; i < targets_length * 2; ++i) { + size_t j; + + __m256i v_data_a = *(__m256i *)(source + (i * 32)); + + for (j = 0; j < bytes_length; ++j) { + uint8_t *target = targets[j]; + uint32_t *target32 = (uint32_t *)target; + __m256i v_comparand = _mm256_set1_epi8(bytes[j]); + uint32_t *out_mask = (uint32_t*)target; + __m256i v_results_a = _mm256_cmpeq_epi8(v_data_a, v_comparand); + uint32_t mask = (uint32_t)_mm256_movemask_epi8(v_results_a); + target32[i] = mask; + } + } +#endif +} + void avx2_and_bits( uint8_t *target, size_t target_length, diff --git a/hw-simd.cabal b/hw-simd.cabal --- a/hw-simd.cabal +++ b/hw-simd.cabal @@ -3,7 +3,7 @@ -- see: https://github.com/sol/hpack name: hw-simd -version: 0.1.1.2 +version: 0.1.1.3 synopsis: SIMD library description: Please see the README on Github at category: Data, Bit, SIMD diff --git a/src/HaskellWorks/Data/Simd/Comparison/Avx2.hs b/src/HaskellWorks/Data/Simd/Comparison/Avx2.hs --- a/src/HaskellWorks/Data/Simd/Comparison/Avx2.hs +++ b/src/HaskellWorks/Data/Simd/Comparison/Avx2.hs @@ -1,6 +1,7 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} module HaskellWorks.Data.Simd.Comparison.Avx2 where @@ -9,6 +10,7 @@ import Data.Word import qualified Data.ByteString as BS +import qualified Data.Vector as DV import qualified Data.Vector.Storable as DVS import qualified Foreign.ForeignPtr as F import qualified Foreign.Marshal.Unsafe as F @@ -87,4 +89,32 @@ cmpEqWord8s w8 = CS.toChunkString . cmpEqWord8s w8 . BS.toByteStrings {-# INLINE cmpEqWord8s #-} +class CmpEqWord8sPara a where + type CmpEqWord8sParaTarget a + cmpEqWord8sPara :: DVS.Vector Word8 -> a -> CmpEqWord8sParaTarget a + +instance CmpEqWord8sPara (DVS.Vector Word64) where + type CmpEqWord8sParaTarget (DVS.Vector Word64) = DV.Vector (DVS.Vector Word64) + cmpEqWord8sPara w8s v = case DVS.unsafeCast v :: DVS.Vector Word8 of + u -> case DVS.unsafeToForeignPtr u of + (srcFptr, srcOffset, srcLength) -> if disalignment == 0 + then F.unsafeLocalState $ do + tgtFptr <- F.mallocForeignPtrBytes (srcLength * DVS.length w8s) + F.withForeignPtr srcFptr $ \srcPtr -> do + F.withForeignPtr tgtFptr $ \tgtPtr -> do + let tgtsPtrsV :: DVS.Vector (F.Ptr Word8) = DVS.constructN (DVS.length w8s) $ \t -> + tgtPtr `F.plusPtr` (DVS.length t * DVS.length v) + let (w8sFptr, _, w8sLen) = DVS.unsafeToForeignPtr w8s + let (tgtsPtrsFptr, _, _) = DVS.unsafeToForeignPtr tgtsPtrsV + F.withForeignPtr w8sFptr $ \w8sPtr -> do + F.withForeignPtr tgtsPtrsFptr $ \tgtsPtrsPtr -> do + F.avx2Cmpeq8Para (F.castPtr w8sPtr) (fromIntegral w8sLen) (F.castPtr tgtsPtrsPtr) (fromIntegral w64sLen) (srcPtr `F.plusPtr` srcOffset) + + let tgtV = DVS.unsafeFromForeignPtr tgtFptr 0 (w64sLen * DVS.length w8s) + return $ DV.constructN (DVS.length w8s) $ \t -> DVS.take w64sLen (DVS.drop (DV.length t * w64sLen) tgtV) + + else error $ "Unaligned byte string: " <> show disalignment + where w64sLen = srcLength `div` 64 + disalignment = srcLength - w64sLen * 64 + {-# INLINE cmpEqWord8sPara #-} diff --git a/src/HaskellWorks/Data/Simd/Internal/Foreign.chs b/src/HaskellWorks/Data/Simd/Internal/Foreign.chs --- a/src/HaskellWorks/Data/Simd/Internal/Foreign.chs +++ b/src/HaskellWorks/Data/Simd/Internal/Foreign.chs @@ -21,6 +21,11 @@ {#call unsafe avx2_cmpeq8 as c_cmpeq8#} byte target targetLength source {-# INLINE avx2Cmpeq8 #-} +avx2Cmpeq8Para :: Ptr UInt8 -> Size -> Ptr (Ptr UInt8) -> Size -> Ptr UInt8 -> IO () +avx2Cmpeq8Para bytes bytes_length target targetLength source = requireAvx2 $ do + {#call unsafe avx2_cmpeq8_para as c_cmpeq8_para#} bytes bytes_length target targetLength source +{-# INLINE avx2Cmpeq8Para #-} + avx2AndBits :: Ptr UInt8 -> Size -> Ptr UInt8 -> Ptr UInt8 -> IO () avx2AndBits target targetLength source_a source_b = requireAvx2 $ do {#call unsafe avx2_and_bits as c_avx2_and_bits#} target targetLength source_a source_b diff --git a/test/HaskellWorks/Data/Simd/ComparisonSpec.hs b/test/HaskellWorks/Data/Simd/ComparisonSpec.hs --- a/test/HaskellWorks/Data/Simd/ComparisonSpec.hs +++ b/test/HaskellWorks/Data/Simd/ComparisonSpec.hs @@ -11,6 +11,7 @@ import Hedgehog import Test.Hspec +import qualified Data.Vector as DV import qualified Data.Vector.Storable as DVS import qualified HaskellWorks.Data.Simd.Comparison.Avx2 as AVX import qualified HaskellWorks.Data.Simd.Comparison.Stock as STK @@ -35,3 +36,20 @@ $ G.word8x8 (G.choice [pure w, G.word8 R.constantBounded]) AVX.cmpEqWord8s w vs === STK.cmpEqWord8s w vs vs === vs + it "AVX2 Para 1" $ requireProperty $ do + w <- forAll $ G.word8 R.constantBounded + vs :: DVS.Vector Word64 <- forAll + $ G.storableVector (R.singleton 8) + $ G.word8x8 (G.choice [pure w, G.word8 R.constantBounded]) + DV.head (AVX.cmpEqWord8sPara (DVS.singleton w) vs) === AVX.cmpEqWord8s w vs + it "AVX2 Para 2" $ requireProperty $ do + len <- forAll $ (* 8) <$> G.int (R.linear 0 8) + w0 <- forAll $ G.word8 R.constantBounded + w1 <- forAll $ G.word8 R.constantBounded + vs :: DVS.Vector Word64 <- forAll + $ G.storableVector (R.singleton len) + $ G.word8x8 (G.choice [pure w0, pure w1, G.word8 R.constantBounded]) + let actual = AVX.cmpEqWord8sPara (DVS.fromList [w0, w1]) vs + (actual DV.! 0) === AVX.cmpEqWord8s w0 vs + (actual DV.! 1) === AVX.cmpEqWord8s w1 vs +