Compile count{Leading,Trailing}Zeros to corresponding x86_64 instructions under -mbmi2

This works similarly to existing implementation for popCount.
Trac ticket: #16086.

3 files changed, 34 insertions, 0 deletions

diff --git a/libraries/base/Data/Bits.hs b/libraries/base/Data/Bits.hs
index 000e663b83..69446f9adc 100644
--- a/libraries/base/Data/Bits.hs
+++ b/libraries/base/Data/Bits.hs
@@ -438,6 +438,9 @@ instance Bits Int where
     {-# INLINE shift #-}
     {-# INLINE bit #-}
     {-# INLINE testBit #-}
+    -- We want popCnt# to be inlined in user code so that `ghc -msse4.2`
+    -- can compile it down to a popcnt instruction without an extra function call
+    {-# INLINE popCount #-}
 
     zeroBits = 0
 
@@ -478,13 +481,16 @@ instance Bits Int where
 instance FiniteBits Int where
     finiteBitSize _ = WORD_SIZE_IN_BITS
     countLeadingZeros  (I# x#) = I# (word2Int# (clz# (int2Word# x#)))
+    {-# INLINE countLeadingZeros #-}
     countTrailingZeros (I# x#) = I# (word2Int# (ctz# (int2Word# x#)))
+    {-# INLINE countTrailingZeros #-}
 
 -- | @since 2.01
 instance Bits Word where
     {-# INLINE shift #-}
     {-# INLINE bit #-}
     {-# INLINE testBit #-}
+    {-# INLINE popCount #-}
 
     (W# x#) .&.   (W# y#)    = W# (x# `and#` y#)
     (W# x#) .|.   (W# y#)    = W# (x# `or#`  y#)
@@ -519,7 +525,9 @@ instance Bits Word where
 instance FiniteBits Word where
     finiteBitSize _ = WORD_SIZE_IN_BITS
     countLeadingZeros  (W# x#) = I# (word2Int# (clz# x#))
+    {-# INLINE countLeadingZeros #-}
     countTrailingZeros (W# x#) = I# (word2Int# (ctz# x#))
+    {-# INLINE countTrailingZeros #-}
 
 -- | @since 2.01
 instance Bits Integer where
diff --git a/libraries/base/GHC/Int.hs b/libraries/base/GHC/Int.hs
index 2c5ca9d5a8..d87d352cb7 100644
--- a/libraries/base/GHC/Int.hs
+++ b/libraries/base/GHC/Int.hs
@@ -177,6 +177,7 @@ instance Bits Int8 where
     {-# INLINE shift #-}
     {-# INLINE bit #-}
     {-# INLINE testBit #-}
+    {-# INLINE popCount #-}
 
     (I8# x#) .&.   (I8# y#)   = I8# (word2Int# (int2Word# x# `and#` int2Word# y#))
     (I8# x#) .|.   (I8# y#)   = I8# (word2Int# (int2Word# x# `or#`  int2Word# y#))
@@ -211,6 +212,8 @@ instance Bits Int8 where
 
 -- | @since 4.6.0.0
 instance FiniteBits Int8 where
+    {-# INLINE countLeadingZeros #-}
+    {-# INLINE countTrailingZeros #-}
     finiteBitSize _ = 8
     countLeadingZeros  (I8# x#) = I# (word2Int# (clz8# (int2Word# x#)))
     countTrailingZeros (I8# x#) = I# (word2Int# (ctz8# (int2Word# x#)))
@@ -381,6 +384,7 @@ instance Bits Int16 where
     {-# INLINE shift #-}
     {-# INLINE bit #-}
     {-# INLINE testBit #-}
+    {-# INLINE popCount #-}
 
     (I16# x#) .&.   (I16# y#)  = I16# (word2Int# (int2Word# x# `and#` int2Word# y#))
     (I16# x#) .|.   (I16# y#)  = I16# (word2Int# (int2Word# x# `or#`  int2Word# y#))
@@ -415,6 +419,8 @@ instance Bits Int16 where
 
 -- | @since 4.6.0.0
 instance FiniteBits Int16 where
+    {-# INLINE countLeadingZeros #-}
+    {-# INLINE countTrailingZeros #-}
     finiteBitSize _ = 16
     countLeadingZeros  (I16# x#) = I# (word2Int# (clz16# (int2Word# x#)))
     countTrailingZeros (I16# x#) = I# (word2Int# (ctz16# (int2Word# x#)))
@@ -587,6 +593,7 @@ instance Bits Int32 where
     {-# INLINE shift #-}
     {-# INLINE bit #-}
     {-# INLINE testBit #-}
+    {-# INLINE popCount #-}
 
     (I32# x#) .&.   (I32# y#)  = I32# (word2Int# (int2Word# x# `and#` int2Word# y#))
     (I32# x#) .|.   (I32# y#)  = I32# (word2Int# (int2Word# x# `or#`  int2Word# y#))
@@ -622,6 +629,8 @@ instance Bits Int32 where
 
 -- | @since 4.6.0.0
 instance FiniteBits Int32 where
+    {-# INLINE countLeadingZeros #-}
+    {-# INLINE countTrailingZeros #-}
     finiteBitSize _ = 32
     countLeadingZeros  (I32# x#) = I# (word2Int# (clz32# (int2Word# x#)))
     countTrailingZeros (I32# x#) = I# (word2Int# (ctz32# (int2Word# x#)))
@@ -825,6 +834,7 @@ instance Bits Int64 where
     {-# INLINE shift #-}
     {-# INLINE bit #-}
     {-# INLINE testBit #-}
+    {-# INLINE popCount #-}
 
     (I64# x#) .&.   (I64# y#)  = I64# (word64ToInt64# (int64ToWord64# x# `and64#` int64ToWord64# y#))
     (I64# x#) .|.   (I64# y#)  = I64# (word64ToInt64# (int64ToWord64# x# `or64#`  int64ToWord64# y#))
@@ -1002,6 +1012,7 @@ instance Bits Int64 where
     {-# INLINE shift #-}
     {-# INLINE bit #-}
     {-# INLINE testBit #-}
+    {-# INLINE popCount #-}
 
     (I64# x#) .&.   (I64# y#)  = I64# (word2Int# (int2Word# x# `and#` int2Word# y#))
     (I64# x#) .|.   (I64# y#)  = I64# (word2Int# (int2Word# x# `or#`  int2Word# y#))
@@ -1078,6 +1089,8 @@ uncheckedIShiftRA64# = uncheckedIShiftRA#
 
 -- | @since 4.6.0.0
 instance FiniteBits Int64 where
+    {-# INLINE countLeadingZeros #-}
+    {-# INLINE countTrailingZeros #-}
     finiteBitSize _ = 64
 #if WORD_SIZE_IN_BITS < 64
     countLeadingZeros  (I64# x#) = I# (word2Int# (clz64# (int64ToWord64# x#)))
diff --git a/libraries/base/GHC/Word.hs b/libraries/base/GHC/Word.hs
index d19a31dfb2..e714392e9c 100644
--- a/libraries/base/GHC/Word.hs
+++ b/libraries/base/GHC/Word.hs
@@ -168,6 +168,7 @@ instance Bits Word8 where
     {-# INLINE shift #-}
     {-# INLINE bit #-}
     {-# INLINE testBit #-}
+    {-# INLINE popCount #-}
 
     (W8# x#) .&.   (W8# y#)   = W8# (x# `and#` y#)
     (W8# x#) .|.   (W8# y#)   = W8# (x# `or#`  y#)
@@ -201,6 +202,8 @@ instance Bits Word8 where
 
 -- | @since 4.6.0.0
 instance FiniteBits Word8 where
+    {-# INLINE countLeadingZeros #-}
+    {-# INLINE countTrailingZeros #-}
     finiteBitSize _ = 8
     countLeadingZeros  (W8# x#) = I# (word2Int# (clz8# x#))
     countTrailingZeros (W8# x#) = I# (word2Int# (ctz8# x#))
@@ -356,6 +359,7 @@ instance Bits Word16 where
     {-# INLINE shift #-}
     {-# INLINE bit #-}
     {-# INLINE testBit #-}
+    {-# INLINE popCount #-}
 
     (W16# x#) .&.   (W16# y#)  = W16# (x# `and#` y#)
     (W16# x#) .|.   (W16# y#)  = W16# (x# `or#`  y#)
@@ -389,6 +393,8 @@ instance Bits Word16 where
 
 -- | @since 4.6.0.0
 instance FiniteBits Word16 where
+    {-# INLINE countLeadingZeros #-}
+    {-# INLINE countTrailingZeros #-}
     finiteBitSize _ = 16
     countLeadingZeros  (W16# x#) = I# (word2Int# (clz16# x#))
     countTrailingZeros (W16# x#) = I# (word2Int# (ctz16# x#))
@@ -590,6 +596,7 @@ instance Bits Word32 where
     {-# INLINE shift #-}
     {-# INLINE bit #-}
     {-# INLINE testBit #-}
+    {-# INLINE popCount #-}
 
     (W32# x#) .&.   (W32# y#)  = W32# (x# `and#` y#)
     (W32# x#) .|.   (W32# y#)  = W32# (x# `or#`  y#)
@@ -623,6 +630,8 @@ instance Bits Word32 where
 
 -- | @since 4.6.0.0
 instance FiniteBits Word32 where
+    {-# INLINE countLeadingZeros #-}
+    {-# INLINE countTrailingZeros #-}
     finiteBitSize _ = 32
     countLeadingZeros  (W32# x#) = I# (word2Int# (clz32# x#))
     countTrailingZeros (W32# x#) = I# (word2Int# (ctz32# x#))
@@ -762,6 +771,7 @@ instance Bits Word64 where
     {-# INLINE shift #-}
     {-# INLINE bit #-}
     {-# INLINE testBit #-}
+    {-# INLINE popCount #-}
 
     (W64# x#) .&.   (W64# y#)  = W64# (x# `and64#` y#)
     (W64# x#) .|.   (W64# y#)  = W64# (x# `or64#`  y#)
@@ -914,6 +924,7 @@ instance Bits Word64 where
     {-# INLINE shift #-}
     {-# INLINE bit #-}
     {-# INLINE testBit #-}
+    {-# INLINE popCount #-}
 
     (W64# x#) .&.   (W64# y#)  = W64# (x# `and#` y#)
     (W64# x#) .|.   (W64# y#)  = W64# (x# `or#`  y#)
@@ -959,6 +970,8 @@ uncheckedShiftRL64# = uncheckedShiftRL#
 
 -- | @since 4.6.0.0
 instance FiniteBits Word64 where
+    {-# INLINE countLeadingZeros #-}
+    {-# INLINE countTrailingZeros #-}
     finiteBitSize _ = 64
     countLeadingZeros  (W64# x#) = I# (word2Int# (clz64# x#))
     countTrailingZeros (W64# x#) = I# (word2Int# (ctz64# x#))