Opened 2 months ago

Closed 2 months ago

#13614 closed bug (worksforme)

Rewrite rules not applied exhaustively when simplifying from plugin

Reported by: nomeata Owned by:
Priority: normal Milestone:
Component: GHC API Version: 8.1
Keywords: Cc:
Operating System: Unknown/Multiple Architecture: Unknown/Multiple
Type of failure: None/Unknown Test Case:
Blocked By: Blocking:
Related Tickets: Differential Rev(s):
Wiki Page:

Description

Consider this program:

{-# OPTIONS_GHC -O -fplugin TestPlugin #-}
module Test where

foo :: Int -> Int
foo = id
{-# INLINE [0] foo #-}

{-# RULES
"foo1" [1] foo 1 = foo 2
"foo2" [1] foo 2 = foo 3
 #-}

fun :: Int -> Int -> Int
fun = (+)
{-# NOINLINE fun #-}

test = foo 1 `fun` foo 2

I would expect that one run of the simplifier in phase 1 will turn this into

test = foo 3 `fun` foo 3

I am using this plugin to test this:

module TestPlugin where

import System.Exit
import Control.Monad

import GhcPlugins
import Simplify
import CoreStats
import SimplMonad
import FamInstEnv
import SimplEnv

-- Plugin boiler plate

plugin :: Plugin
plugin = defaultPlugin { installCoreToDos = install }

install :: [CommandLineOption] -> [CoreToDo] -> CoreM [CoreToDo]
install _ (simpl:xs) = return $ simpl : pass : xs
  where pass = CoreDoPluginPass "Test" testPass

-- The plugin

testPass :: ModGuts -> CoreM ModGuts
testPass guts = do
    let [expr] = [ e | NonRec v e <- mg_binds guts
                     , occNameString (occName v) == "test" ]
    simplified_expression <- simplify guts expr

    putMsg $
        text "Test" $$
        nest 4 (hang (text "Before" <> colon) 4 (ppr expr)) $$
        nest 4 (hang (text "After" <> colon) 4 (ppr simplified_expression))

    liftIO $ exitFailure

-- A simplifier

simplify :: ModGuts -> CoreExpr -> CoreM CoreExpr
simplify guts expr = do
    dflags <- getDynFlags

    let dflags' = dflags { ufVeryAggressive = True }
    us <- liftIO $ mkSplitUniqSupply 's'
    let sz = exprSize expr

    rule_base <- getRuleBase
    vis_orphs <- getVisibleOrphanMods
    let rule_base2 = extendRuleBaseList rule_base (mg_rules guts)
    let rule_env = RuleEnv rule_base2 vis_orphs

    (expr', _) <- liftIO $ initSmpl dflags' rule_env emptyFamInstEnvs us sz $
        simplExpr (simplEnv 1) >=> simplExpr (simplEnv 1) $
        expr
    return expr'

simplEnv :: Int -> SimplEnv
simplEnv p = mkSimplEnv $ SimplMode { sm_names = ["Test"]
                                    , sm_phase = Phase p
                                    , sm_rules = True
                                    , sm_inline = True
                                    , sm_eta_expand = True
                                    , sm_case_case = True }

But I get:

$ ghc-head -package ghc Test.hs
[1 of 2] Compiling TestPlugin       ( TestPlugin.hs, TestPlugin.o )
[2 of 2] Compiling Test             ( Test.hs, Test.o )
Test
    Before: fun (foo (GHC.Types.I# 1#)) (foo (GHC.Types.I# 2#))
    After: fun (foo (GHC.Types.I# 2#)) (foo (GHC.Types.I# 3#))

If I however compile this without the plugin, and look at what’s happening with -dverbose-core2core, I observe this:

…
test :: Int
test = fun (foo (GHC.Types.I# 1#)) (foo (GHC.Types.I# 2#))
…
==================== Simplifier ====================
  Max iterations = 4
  SimplMode {Phase = 1 [main],
             inline,
             rules,
             eta-expand,
             case-of-case}
…
test :: Int
test = fun (foo (GHC.Types.I# 3#)) (foo (GHC.Types.I# 3#))
…

So what am I doing wrong in my plugin? Any help is appreciated.

Change History (10)

comment:1 Changed 2 months ago by simonpj

It's not just that they aren't getting applied; somehow foo 1 is getting rewritten to foo 2 which is deeply strange. I have no hypothesis here.

I'm not sure why you are simplifying twice.

You don't seem to be doing occurrence analysis, although I can't see how that will lead to this behaviour.

Try -ddump-rule-rewrites.

comment:2 in reply to:  1 Changed 2 months ago by nomeata

Replying to simonpj:

It's not just that they aren't getting applied; somehow foo 1 is getting rewritten to foo 2 which is deeply strange.

That is not strange, one of the rules is

"foo1" [1] foo 1 = foo 2

after all. But I wonder why it does not then apply the second rule.

I'm not sure why you are simplifying twice.

Just in case it takes multiple simplifier iterations to apply rules exhaustively (although it does not help).

My hypothesis is as follows:

Rewrite rules are attached (via the IdInfo field ruleInfo)to the Id of the outermost function on the LHS of the rule; in this case foo (and not looked up in some mapping, as one might naively expect). As far as I can tell, the rule foo1 and foo2 are attached to all occurrences of foo in the module – excluding the ones on the RHS of the rules themselves!

This would explain this behaviour. It would rewrite

foo[ruleInfo=foo1,foo2] 1 `fun` foo[ruleInfo=foo1,foo2] 2

to

foo[ruleInfo=] 2 `fun` foo[ruleInfo=] 3

and then not apply any more rules (because ruleInfo) is empty.

This hypothesis would explain the behavior; what puzzles me is that it works as expected if I compile the module as normal. And, pragmaticall, what I wave to do so that my plugin acts as expected.

Last edited 2 months ago by nomeata (previous) (diff)

comment:3 Changed 2 months ago by simonpj

Ah now I get it

As far as I can tell, the rule foo1 and foo2 are attached to all occurrences of foo in the module – excluding the ones on the RHS of the rules themselves!

That should not matter: at every occurrence of a variable it is looked up in the in-scope set, to get the "master copy", so after simplification all occurrences point to the binder. This is done by SimplEnv.substId.

There is a debug WARN if the Id is not in the in-scope set.

But I bet you have a non-debug compiler; and that you have an in-scope set that does not include foo. Might that be it?

comment:4 Changed 2 months ago by nomeata

Resolution: worksforme
Status: newclosed

But I bet you have a non-debug compiler; and that you have an in-scope set that does not include foo. Might that be it?

That seems to be it! (Although I had a hard time fixing it because of a bug(?) in the recompilation checker: If you used -dynamic-too before, but forget it later, GHC will happily use old TestPlugin.dyn_o/TestPlugin.o files, so I kept seeing unchanged behaviour… will file a bug).

Incidentially, why is this not an issue for simplifyExpr in SimplCore – I don’t see code there adding anything to the in-scope set? Or is it ok because that function is not going to do such elaborate things as applying rules anyways?

comment:5 Changed 2 months ago by nomeata

Resolution: worksforme
Status: closednew

Ok, so this works fine for locally defined things. But I run into the same (or a similar) problem where rules attached to globally defined things do not fire.

Here is my test file (with commented-out rule):

{-# OPTIONS_GHC -O -fplugin TestPlugin #-}
module Test where
import GHC.Base (foldr)
{- # RULES "foldr/id_mine" GHC.Base.foldr (:) [] = id #-}
test :: [a] -> [a]
test xs = map id xs

and here the plugin, with the fix from earlier:

module TestPlugin where

import System.Exit
import Control.Monad

import GhcPlugins
import Simplify
import CoreStats
import SimplMonad
import FamInstEnv
import SimplEnv
import OccurAnal

-- Plugin boiler plate

plugin :: Plugin
plugin = defaultPlugin { installCoreToDos = install }

install :: [CommandLineOption] -> [CoreToDo] -> CoreM [CoreToDo]
install _ (simpl:xs) = return $ simpl : pass : xs
  where pass = CoreDoPluginPass "Test" testPass

-- The plugin

testPass :: ModGuts -> CoreM ModGuts
testPass guts = do
    let [expr] = [ e | NonRec v e <- mg_binds guts
                     , occNameString (occName v) == "test" ]
    simplified_expression <- simplify guts expr

    putMsg $
        text "Test" $$
        nest 4 (hang (text "Before" <> colon) 4 (ppr expr)) $$
        nest 4 (hang (text "After" <> colon) 4 (ppr simplified_expression))

    liftIO $ exitFailure

-- A simplifier

simplify :: ModGuts -> CoreExpr -> CoreM CoreExpr
simplify guts expr = do
    dflags <- getDynFlags

    let dflags' = dflags { ufVeryAggressive = True }

    us <- liftIO $ mkSplitUniqSupply 's'
    let sz = exprSize expr

    rule_base <- getRuleBase
    vis_orphs <- getVisibleOrphanMods
    let rule_base2 = extendRuleBaseList rule_base (mg_rules guts)
    let rule_env = RuleEnv rule_base2 vis_orphs
    let top_lvls = bindersOfBinds (mg_binds guts)


    (expr', _) <- liftIO $ initSmpl dflags' rule_env emptyFamInstEnvs us sz $ do
        simplExpr (simplEnv top_lvls 1) (occurAnalyseExpr expr)
    return expr'



simplEnv :: [Var] -> Int -> SimplEnv
simplEnv vars p = env1
  where
    env1 = addNewInScopeIds env0 vars
    env0 =  mkSimplEnv $ SimplMode { sm_names = ["Test"]
                                   , sm_phase = Phase p
                                   , sm_rules = True
                                   , sm_inline = True
                                   , sm_eta_expand = True
                                   , sm_case_case = True }

If I run this I get:

$ ghc-head -O  -dynamic-too -package ghc Test.hs -fforce-recomp
[1 of 2] Compiling TestPlugin       ( TestPlugin.hs, TestPlugin.o )
[2 of 2] Compiling Test             ( Test.hs, Test.o )
Test
    Before:
        \ (@ a) (xs :: [a]) ->
          GHC.Base.build
            @ a
            (\ (@ b1)
               (c [OS=OneShot] :: a -> b1 -> b1)
               (n [OS=OneShot] :: b1) ->
               GHC.Base.foldr @ a @ b1 c n xs)
    After:
        \ (@ a) (xs :: [a]) ->
          GHC.Base.foldr @ a @ [a] (GHC.Types.: @ a) (GHC.Types.[] @ a) xs

Note that GHC.Base.foldr is still there in the After:: expression, despite the foldr/id rule in GHC.Base, which should simplify this code!

If I add that rule to my module (as hinted at above), it does fire:

$ ghc-head -O  -dynamic-too -package ghc Test.hs -fforce-recomp
[1 of 2] Compiling TestPlugin       ( TestPlugin.hs, TestPlugin.o )
[2 of 2] Compiling Test             ( Test.hs, Test.o )
Test
    Before:
        \ (@ a) (xs :: [a]) ->
          GHC.Base.build
            @ a
            (\ (@ b1)
               (c [OS=OneShot] :: a -> b1 -> b1)
               (n [OS=OneShot] :: b1) ->
               GHC.Base.foldr @ a @ b1 c n xs)
    After: \ (@ a) (xs :: [a]) -> xs

comment:6 Changed 2 months ago by simonpj

I wonder if you have ensured that GHC.Base is loaded? If the interface is not loaded, GHC won't see the rule.

Check what you get in the intial rule-base

comment:7 Changed 2 months ago by nomeata

I wonder if you have ensured that GHC.Base is loaded? If the interface is not loaded, GHC won't see the rule.

Yes, GHC.Base is loaded.

I was under the impression that (non-orphan) rules are attached to the ID, so should it matter? Or is that only true for locally defined rules? Or is GHC.Base special in that matter?

In the above plugin code, the result of getRuleBase is indeed empty, and so is mg_rules guts.

comment:8 Changed 2 months ago by simonpj

You'll have to look at the code, but only locally-defined Ids have their rules attached. The rest are in a global database (the rule-base). I can't explain why it's empty.

comment:9 Changed 2 months ago by nomeata

You'll have to look at the code, but only locally-defined Ids have their rules attached. The rest are in a global database (the rule-base)

Ah, thanks for clearing that up for me. I always assumed that the rule base is only for orphan rules. Now I know where to look, and will hopefully find out what’s wrong.

comment:10 Changed 2 months ago by nomeata

Resolution: worksforme
Status: newclosed

Ok, it seems that CoreM’s getRuleBase only returns the rules from the current home package. To get the rules from the external package, I have to use getHscEnv:

    hpt_rule_base <- getRuleBase
    hsc_env <- getHscEnv
    eps <- liftIO $ hscEPS hsc_env
    let rule_base1 = unionRuleBase hpt_rule_base (eps_rule_base eps)
        rule_base2 = extendRuleBaseList rule_base1 (mg_rules guts)
    vis_orphs <- getVisibleOrphanMods
    let rule_env = RuleEnv rule_base2 vis_orphs

Thanks for your help!

Note: See TracTickets for help on using tickets.