GHC can't constant fold even basic power (^) applications for Int (and others?)

[nix-shell:/tmp]$ ghc -O2 -fforce-recomp -ddump-simpl S.hs 2>&1 | tail
[GblId,
 Unf=Unf{Src=<vanilla>, TopLvl=True, Value=False, ConLike=False,
         WorkFree=False, Expandable=False, Guidance=IF_ARGS [] 50 20}]
foo
  = case GHC.Real.$wf1 2# 1# of ww2_a2zt { __DEFAULT ->
    GHC.Types.I# ww2_a2zt
    }

[nix-shell:/tmp]$ cat S.hs
module S (foo) where
foo :: Int
foo = (2 :: Int) ^ (1 :: Int)

This seems like a fairly strange thing to not optimise when constants are known on both sides. I'm complaining about Int for this particular ticket.

changed weight to 5

added Tbug Trac import labels

changed the description

For the record, currently (^) is defined as:

-- | raise a number to a non-negative integral power
{-# SPECIALISE [1] (^) ::
        Integer -> Integer -> Integer,
        Integer -> Int -> Integer,
        Int -> Int -> Int #-}
{-# INLINABLE [1] (^) #-}    -- See Note [Inlining (^)]
(^) :: (Num a, Integral b) => a -> b -> a
x0 ^ y0 | y0 < 0    = errorWithoutStackTrace "Negative exponent"
        | y0 == 0   = 1
        | otherwise = f x0 y0
    where -- f : x0 ^ y0 = x ^ y
          f x y | even y    = f (x * x) (y `quot` 2)
                | y == 1    = x
                | otherwise = g (x * x) (y `quot` 2) x         -- See Note [Half of y - 1]
          -- g : x0 ^ y0 = (x ^ y) * z
          g x y z | even y = g (x * x) (y `quot` 2) z
                  | y == 1 = x * z
                  | otherwise = g (x * x) (y `quot` 2) (x * z) -- See Note [Half of y - 1]

{- Note [Half of y - 1]
   ~~~~~~~~~~~~~~~~~~~~~
   Since y is guaranteed to be odd and positive here,
   half of y - 1 can be computed as y `quot` 2, optimising subtraction away.
-}

To perform constant folding, it would be better to have primitives such as:

ipowInt :: Int# -> Int# -> Int#
ipowWord :: Word# -> Word# -> Word#

that we can match on in Core.

Then we could add (^) as a method of Num a, change its type to be (^) :: a -> a -> a and use the appropriate primitives (or fall back to the generic implementation otherwise). Exactly like we do for other primitives.

Changing the type of (^) is a breaking change but it shouldn't harm much. It needs the approval of the CLC though.

---

By the way, the generic implementation isn't very efficient for Int/Word. The following one that I've just adapted from [1] performs at least twice as fast in my tests:

ipowInt :: Int -> Int -> Int
ipowInt x y
   | y < 0     = errorWithoutStackTrace "Negative exponent"
   | otherwise = go 1 x y
   where
      go r b e =
         let
            e1 = e .&. 1
            r' = r * (b * e1 + (e1 `xor` 1)) -- branchless
            e' = e `unsafeShiftR` 1
         in case e' of
            0 -> r'
            _ -> go r' (b*b) e'

This is another pretty compelling argument in favor of performing the change mentioned above.

[1] https://stackoverflow.com/questions/101439/the-most-efficient-way-to-implement-an-integer-based-power-function-powint-int

Then we could add (^) as a method of Num a, change its type to be (^) :: a -> a -> a and use the appropriate primitives (or fall back to the generic implementation otherwise). Exactly like we do for other primitives.

Hmm in fact it won't work for non integral Num types (e.g., Double)... Maybe we could add another function (ipow?) to Integral instead and add rules to convert (^) into ipow for known cases.

added Pnormal label

I wonder if expression specialisation (cf !12319) could help here. Similarly to SpecConstr, we could specialise calls with literal arguments.

foo :: Word
foo = (2 :: Word) ^ (1 :: Word)

Gives the following Core:

foo
  = case Test.$w$spowImpl 2## 1## of ww_s14k { __DEFAULT ->
    GHC.Types.W# ww_s14k
    }

If we could specialise $w$spowImpl to its 2## 1## arguments, it would work.

We would get something like this (replacing powImpl with only powImplAcc):

{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE UnboxedTuples #-}

module Test (foo) where

import GHC.Exts

foo :: Word
-- foo = (2 :: Word) ^ (1 :: Word)
-- call our wrapper instead
foo = powImplAccW 2 1 1

-- powImplAcc wrapper specialised to Word
{-# INLINE powImplAccW #-}
powImplAccW :: Word -> Word -> Word -> Word
powImplAccW (W# x) (W# y) (W# z) = W# (powImplAcc x y z)


-- powImplAcc worker specialised to Word
{-# INLINABLE powImplAcc #-}
powImplAcc :: Word# -> Word# -> Word# -> Word#
powImplAcc (W# -> x) (W# -> y) (W# -> z)
  | even y    = case powImplAccW (x * x) (y `quot` 2) z of W# r -> r
  | y == 1    = case x * z of W# r -> r
  | otherwise = case powImplAccW (x * x) (y `quot` 2) (x * z) of W# r -> r

-- powImplAcc specialised to specific Word arguments: 2, 1, 1
powImplAcc_spec :: (# #) -> Word#
powImplAcc_spec _ =
  let x = 2
      y = 1
      z = 1
  in if
    | even y    -> case powImplAccW (x * x) (y `quot` 2) z of W# r -> r
    | y == 1    -> case x * z of W# r -> r
    | otherwise -> case powImplAccW (x * x) (y `quot` 2) (x * z) of W# r -> r

-- future auto-generated rule
{-# RULES "foo" powImplAcc 2## 1## 1## = powImplAcc_spec (# #) #-}   

which correctly produces:

- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
foo :: Word
[GblId,
 Unf=Unf{Src=<vanilla>, TopLvl=True,
         Value=True, ConLike=True, WorkFree=True, Expandable=True,
         Guidance=IF_ARGS [] 10 10}]
foo = GHC.Types.W# 2##

added specialisation label

Yes, with !12319 I think we could do this:

(^) :: (Num a, Integral b) => a -> b -> a
x0 ^ y0 | y0 < 0    = errorWithoutStackTrace "Negative exponent"
        | y0 == 0   = 1
        | otherwise = pow_f x0 y0

pow_f :: (Num a, Integral b) => a -> b -> a
pow_f x y | even y    = pow_f (x * x) (y `quot` 2)
          | y == 1    = x
          | otherwise = pow_g (x * x) (y `quot` 2) x

pow_g :: (Num a, Integral b) => a -> b -> a -> a
pow_g x y z | even y = pow_g (x * x) (y `quot` 2) z
            | y == 1 = x * z
            | otherwise = pow_g (x * x) (y `quot` 2) (x * z)

{-# SPECIALISE           (^) @Int @Int #-}
{-# SPECIALISE forall x. (^) @Int @Int x 0 #-}
{-# SPECIALISE forall x. (^) @Int @Int x 1 #-}
{-# SPECIALISE forall x. (^) @Int @Int x 2 #-}

{-# SPECIALISE forall x. pow_f @Int @Int x 0 #-}
{-# SPECIALISE forall x. pow_f @Int @Int x 1 #-}
{-# SPECIALISE forall x. pow_f @Int @Int x 2 #-}

{-# SPECIALISE forall x. pow_g @Int @Int x 0 #-}
{-# SPECIALISE forall x. pow_g @Int @Int x 1 #-}
{-# SPECIALISE forall x. pow_g @Int @Int x 2 #-}

That should unroll the loop twice. It's be nicer to be able to say "k where k < 3" but that's a lot harder.

Would you like to try it?

Why not, but I'll wait for !12319 to be merged first :)

Also instead of writing explicitly the SPECIALISE rules, could we automate this?

• When:
  • we see a call site foo lit1 lit2 ... litn, where the lit* are literals
  • and foo is saturated
  • and foo has an unfolding foo_rhs
• Then generate:
  • a binding foo_spec (##) = foo_rhs lit1 lit2 ... litn
  • a rule: {-# RULES "foo_spec" foo lit1 lit2 ... litn = foo_spec (##) #-}

We could also handle case foo arg1 arg2 ... argn, where at least some of the args are literals (instead of all of them).

Also instead of writing explicitly the SPECIALISE rules, could we automate this?

I am doubtful. If we see (^) x 1000 do we really want to unroll a 1000-iteration loop?

mentioned in issue #24359