Changing Haskell implementation of k-nucleotide to vector-hashtables.

Main.hs

{-# LANGUAGE ScopedTypeVariables #-}

-- The Computer Language Benchmarks Game
-- https://salsa.debian.org/benchmarksgame-team/benchmarksgame/
--
-- Contributed by cahu ette

module Main where

import Data.Bits
import Data.List
import Data.Word
import Data.Hashable
import Data.Traversable
import Text.Printf
import Data.Maybe

import Control.Monad
import Control.Monad.ST
import Control.Parallel.Strategies

import qualified Data.Map.Strict                 as M
import qualified Data.Vector.Hashtables.Internal as H
import qualified Data.Vector.Unboxed.Mutable     as UV
import qualified Data.ByteString.Char8           as B


type HashTable s k v = H.Dictionary s UV.MVector k UV.MVector v


{- By using only 2 bits to encode keys, it's important to use a different table
 - for different key sizes. Otherwise, if we encode 'A' as 0x00, "AT" and
 - "AAT" would map to the same bucket in the table.
 -
 - We could use 3 bits per letter to avoid this problem if needed.
 -}
bitsForChar :: Char -> Word64
bitsForChar 'a' = 0
bitsForChar 'A' = 0
bitsForChar 'c' = 1
bitsForChar 'C' = 1
bitsForChar 'g' = 2
bitsForChar 'G' = 2
bitsForChar 't' = 3
bitsForChar 'T' = 3
bitsForChar _   = error "Ay, Caramba!"


charForBits :: Word64 -> Char
charForBits 0 = 'A'
charForBits 1 = 'C'
charForBits 2 = 'G'
charForBits 3 = 'T'
charForBits _ = error "Ay, Caramba!"


packKey :: B.ByteString -> Word64
packKey = go zeroBits
  where
    go k bs = case B.uncons bs of
        Nothing      -> k
        Just (c, cs) -> go (unsafeShiftL k 2 .|. bitsForChar c) cs
{-# INLINE packKey #-}

unpackKey :: Int -> Word64 -> B.ByteString
unpackKey = go []
  where
    go s 0 _ = B.pack s
    go s l i = go (charForBits (i .&. 3) : s) (l-1) (unsafeShiftR i 2)
{-# INLINE unpackKey #-}


countOccurrences :: Int -> Int -> B.ByteString -> ST s (HashTable s Word64 Int)
countOccurrences jumpSize frameSize input = do
    t <- H.initialize 1024

    let go bs = unless (B.length bs < frameSize) $ do
            let k = takeFrame bs
            H.alter t (Just . maybe 1 (+1)) (packKey k)
            go (nextFrame bs)

    go input
    return t

  where
    takeFrame = B.take frameSize
    nextFrame = B.drop jumpSize


extractSequence :: String -> B.ByteString -> B.ByteString
extractSequence s = findSeq
  where
    prefix = B.pack ('>' : s)
    skipSeq =
          B.dropWhile (/= '>')
        . B.drop 1
    takeSeq =
          B.filter    (/= '\n')
        . B.takeWhile (/=  '>') -- extract until next header
        . B.dropWhile (/= '\n') -- skip header
    findSeq str
        | prefix `B.isPrefixOf` str  =  takeSeq str
        | otherwise                  =  findSeq (skipSeq str)



main :: IO ()
main = do
    s <- extractSequence "THREE" <$> B.getContents

    let keys    = ["GGT","GGTA","GGTATT","GGTATTTTAATT","GGTATTTTAATTTATAGT"]
    let threads = [0 .. 63]

    let threadWorkOcc key tid = runST $ do
            t <- countOccurrences (length threads) (B.length key) (B.drop tid s)
            fromMaybe 0 <$> H.lookup t (packKey key)

    let calcOcc key = sum $ runEval $
            mapM (rpar . threadWorkOcc (B.pack key)) threads

    let threadWorkFreq len tid = runST $ do
            t  <- countOccurrences (length threads) len (B.drop tid s)
            vs <- H.toList t
            return $ map (\(k, v) -> (B.unpack (unpackKey len k), freq v)) vs
          where
            freq v = 100 * fromIntegral v / fromIntegral (B.length s - len + 1)

    let calcFreq len =
            let l = concat $ runEval $ mapM (rpar . threadWorkFreq len) threads
                m = foldr (uncurry $ M.insertWith (+)) M.empty l
            in
                M.toList m

    let resultsOcc = map (\k -> (k, calcOcc k)) keys

    printFreq (calcFreq 1)
    putStrLn ""
    printFreq (calcFreq 2)
    putStrLn ""
    forM_ resultsOcc $ \(k, r) -> printf "%d\t%s\n" r k

  where

    sortFreq = sortBy
        (\ (k :: String, v :: Double) (k', v') ->
            (compare v' v) `mappend` (compare k k'))

    printFreq :: [(String, Double)] -> IO ()
    printFreq l = forM_ (sortFreq l) $ uncurry (printf "%s %.3f\n")

Try H.initialize 0 instead of 1024. You might not need such extreme capacity.

That doesn't make much difference, I actually think 1024 is quite low, I expect the benchmark to store millions of distinct elements in the hash table.

There is an example input file here that you can test it with: https://benchmarksgame-team.pages.debian.net/benchmarksgame/download/knucleotide-input.txt. On that file the hashtables implementation takes 0.013 seconds and this new implementation takes 0.051 seconds, of which almost all of the time is spent in countOccurrences in this new version.

Replacing alter by a combination of lookup and insert improves the performance significantly.

I think, alter might be used in combination with $!.
Update: I will add alter to Comparison benchmark.

Adding $! like this: H.alter t (\x -> Just $! maybe 1 (+1) x) (packKey k) doesn't help. I also don't think that strictness is an issue because I'm using unboxed vectors.

Looking at the core it seems one difference is that insertWithIndex is not inlined in the alter version and it is also not specialized, so that will lead to a lot of unknown function calls.

Oh, I see insertWithIndex is recursive so it can't be inlined, here I think a manual SAT transformation would help.

This is what I mean by SAT:

insertWithIndex
  :: (MVector ks k, MVector vs v, PrimMonad m, Hashable k, Eq k)
  => Int -> Int -> k -> v -> MutVar (PrimState m) (Dictionary_ (PrimState m) ks k vs v) -> Dictionary_ (PrimState m) ks k vs v -> Int -> m ()
insertWithIndex !targetBucket !hashCode' key' value' getDRef d@Dictionary{..} = go
  where
    go i
      | i >= 0 = do
           hc <- hashCode ! i
           if hc == hashCode'
               then do
                   k  <- key !~ i
                   if k == key'
                       then value <~~ i $ value'
                       else go =<< next ! i
               else go =<< next ! i
      | otherwise = addOrResize targetBucket hashCode' key' value' getDRef d
{-# INLINE insertWithIndex #-}

Could you please share the input file on which program runs 24 sec?

Here: https://surfdrive.surf.nl/files/index.php/s/tohcryAeBL3dPHp (66.6MB)

On my machine it runs around 48s.

With +RTS -N4 and compiled with -O2 -threaded? It takes 81 seconds for me without multithreading.

Yes.

• -O2 -threaded compilation flags.
• runtime:

time knucleotids +RTS -N4 -A64m -n4m -qb0 -K2048M -RTS < fasta25000000.txt
A 30.295
T 30.151
C 19.800
G 19.754

AA 9.177
TA 9.132
AT 9.131
TT 9.091
CA 6.002
AC 6.001
AG 5.987
GA 5.984
CT 5.971
TC 5.971
GT 5.957
TG 5.956
CC 3.917
GC 3.911
CG 3.909
GG 3.902

1471758 GGT
446535 GGTA
47336 GGTATT
893 GGTATTTTAATT
893 GGTATTTTAATTTATAGT
knucleotids +RTS -N4 -A64m -n4m -qb0 -K2048M -RTS < fasta25000000.txt  171.68s user 0.73s system 364% cpu 47.367 total