Bit twiddling in haskell vs. C

BitTwiddle.hs

{-# LANGUAGE BangPatterns             #-}
{-# LANGUAGE MagicHash                #-}
{-# LANGUAGE ForeignFunctionInterface #-}

module Main where

import Criterion
import Criterion.Main
import Data.Bits
import Foreign.C
import GHC.Exts
import GHC.Prim

{-

C code:

int mask(int a, int b) { return -(a==b); }

int cmaskTest(int num) {
    int acc = num;
    for(; num; --num) {
        int a = mask(num & 0x1,   num & 0x3  );
        int b = mask(num & 0x4,   num & 0xc  );
        int c = mask(num & 0x10,  num & 0x30 );
        int d = mask(num & 0x40,  num & 0xc0 );
        int e = mask(num & 0x100, num & 0x300);

        acc = acc + a + b + c + d + e;
    }

    return acc;
}


gives:

	xorl	%eax, %eax
	cmpl	%esi, %edi
	sete	%al
	decl	%eax

-}



mask :: Int# -> Int# -> Int#
mask !a# !b# = let !(I# z#) = fromEnum (a# ==# b#)
                   !q#      = negateInt# z#
               in q#

{-
mask:
	cmpq %rsi,%r14
	je LcyL
	xorl %ebx,%ebx
	jmp *0(%rbp)
LcyL:
	movq $-1,%rbx
	jmp *0(%rbp)

-}


mask2 :: Int# -> Int# -> Int#
mask2 !a# !b# = dest#
  where
    !d# = a# -# b#

    !(I# ws#) = bitSize (0::Int) - 1

    !m# = uncheckedIShiftRA# d# ws#

    !r# = word2Int# (xor# (int2Word# d#) (int2Word# m#)) -# m# -# 1#

    !dest# = uncheckedIShiftRA# r# ws#

{-
mask2 gives straight-line code:

	movq %r14,%rax
	subq %rsi,%rax
	movq %rax,%rcx
	sarq $63,%rcx
	xorq %rcx,%rax
	subq %rcx,%rax
	leaq -1(%rax),%rbx
	sarq $63,%rbx

The hope is that even though this is more complicated, the lack of jumps will
improve performance.

Now for some criterion tests.
-}

foreign import ccall unsafe "cmaskTest" cmaskTest :: CInt -> CInt

numTests :: Int
numTests = 10000000

testCMask :: Pure
testCMask = whnf cmaskTest $ toEnum numTests

testMaskOne :: Pure
testMaskOne = whnf f numTests
  where
    f !(I# num#) = I# (f' num# num#)

    f' !num# !acc# = if num# ==# 0#
                       then acc#
                       else f' nextNum# dest#
      where
        !nextNum# = num# -# 1#

        !wnum# = int2Word# num#

        !num1#  = word2Int# (and# wnum# (int2Word# 0x1#  ))
        !num2#  = word2Int# (and# wnum# (int2Word# 0x3#  ))
        !num3#  = word2Int# (and# wnum# (int2Word# 0x4#  ))
        !num4#  = word2Int# (and# wnum# (int2Word# 0xc#  ))
        !num5#  = word2Int# (and# wnum# (int2Word# 0x10# ))
        !num6#  = word2Int# (and# wnum# (int2Word# 0x30# ))
        !num7#  = word2Int# (and# wnum# (int2Word# 0x40# ))
        !num8#  = word2Int# (and# wnum# (int2Word# 0xc0# ))
        !num9#  = word2Int# (and# wnum# (int2Word# 0x100#))
        !num10# = word2Int# (and# wnum# (int2Word# 0x300#))

        !a#       = mask num1# num2#
        !b#       = mask num3# num4#
        !c#       = mask num5# num6#
        !d#       = mask num7# num8#
        !e#       = mask num9# num10#

        !dest#    = acc# +# a# +# b# +# c# +# d# +# e#


testMaskTwo :: Pure
testMaskTwo = whnf f numTests
  where
    f !(I# num#) = I# (f' num# num#)

    f' !num# !acc# = if num# ==# 0#
                       then acc#
                       else f' nextNum# dest#
      where
        !nextNum# = num# -# 1#

        !wnum# = int2Word# num#

        !num1#  = word2Int# (and# wnum# (int2Word# 0x1#  ))
        !num2#  = word2Int# (and# wnum# (int2Word# 0x3#  ))
        !num3#  = word2Int# (and# wnum# (int2Word# 0x4#  ))
        !num4#  = word2Int# (and# wnum# (int2Word# 0xc#  ))
        !num5#  = word2Int# (and# wnum# (int2Word# 0x10# ))
        !num6#  = word2Int# (and# wnum# (int2Word# 0x30# ))
        !num7#  = word2Int# (and# wnum# (int2Word# 0x40# ))
        !num8#  = word2Int# (and# wnum# (int2Word# 0xc0# ))
        !num9#  = word2Int# (and# wnum# (int2Word# 0x100#))
        !num10# = word2Int# (and# wnum# (int2Word# 0x300#))

        !a#       = mask2 num1# num2#
        !b#       = mask2 num3# num4#
        !c#       = mask2 num5# num6#
        !d#       = mask2 num7# num8#
        !e#       = mask2 num9# num10#

        !dest#    = acc# +# a# +# b# +# c# +# d# +# e#


main = defaultMain [ bench "mask1" testMaskOne
                   , bench "mask2" testMaskTwo
                   , bench "cmask" testCMask ]


{- on my laptop, timings are:

mask1: 177.9ms
mask2: 87.9ms
cmask: 56.42ms

-}