OptimizationFuel.hs 4.91 KB
Newer Older
1
{-# LANGUAGE TypeFamilies #-}
Thomas Schilling's avatar
Thomas Schilling committed
2 3 4 5 6 7
-- | Optimisation fuel is used to control the amount of work the optimiser does.
--
-- Every optimisation step consumes a certain amount of fuel and stops when
-- it runs out of fuel.  This can be used e.g. to debug optimiser bugs: Run
-- the optimiser with varying amount of fuel to find out the exact number of
-- steps where a bug is introduced in the output.
8
module OptimizationFuel
9
    ( OptimizationFuel, amountOfFuel, tankFilledTo, unlimitedFuel, anyFuelLeft, oneLessFuel
10 11 12 13
    , OptFuelState, initOptFuelState
    , FuelConsumer, FuelUsingMonad, FuelState
    , fuelGet, fuelSet, lastFuelPass, setFuelPass
    , fuelExhausted, fuelDec1, tryWithFuel
14
    , runFuelIO, runInfiniteFuelIO, fuelConsumingPass
15
    , FuelUniqSM
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
16
    , liftUniq
17 18 19 20
    )
where

import Data.IORef
Ian Lynagh's avatar
Ian Lynagh committed
21
import Control.Monad
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
22 23
import StaticFlags (opt_Fuel)
import UniqSupply
Simon Marlow's avatar
Simon Marlow committed
24
import Panic
25
import Util
26 27 28

import Compiler.Hoopl
import Compiler.Hoopl.GHC (getFuel, setFuel)
29 30 31

#include "HsVersions.h"

dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
32 33 34 35 36 37 38 39 40 41 42 43 44 45

-- We limit the number of transactions executed using a record of flags
-- stored in an HscEnv. The flags store the name of the last optimization
-- pass and the amount of optimization fuel remaining.
data OptFuelState =
  OptFuelState { pass_ref :: IORef String
               , fuel_ref :: IORef OptimizationFuel
               }
initOptFuelState :: IO OptFuelState
initOptFuelState =
  do pass_ref' <- newIORef "unoptimized program"
     fuel_ref' <- newIORef (tankFilledTo opt_Fuel)
     return OptFuelState {pass_ref = pass_ref', fuel_ref = fuel_ref'}

46 47 48
type FuelConsumer a = OptimizationFuel -> (a, OptimizationFuel)

tankFilledTo :: Int -> OptimizationFuel
49 50 51 52
amountOfFuel :: OptimizationFuel -> Int

anyFuelLeft :: OptimizationFuel -> Bool
oneLessFuel :: OptimizationFuel -> OptimizationFuel
53
unlimitedFuel :: OptimizationFuel
54 55 56 57 58

newtype OptimizationFuel = OptimizationFuel Int
  deriving Show

tankFilledTo = OptimizationFuel
59 60 61
amountOfFuel (OptimizationFuel f) = f

anyFuelLeft (OptimizationFuel f) = f > 0
62
oneLessFuel (OptimizationFuel f) = ASSERT (f > 0) (OptimizationFuel (f - 1))
63
unlimitedFuel = OptimizationFuel infiniteFuel
64

65 66
data FuelState = FuelState { fs_fuel :: OptimizationFuel, fs_lastpass :: String }
newtype FuelUniqSM a = FUSM { unFUSM :: FuelState -> UniqSM (a, FuelState) }
67

68 69 70
fuelConsumingPass :: String -> FuelConsumer a -> FuelUniqSM a
fuelConsumingPass name f = do setFuelPass name
                              fuel <- fuelGet
71
                              let (a, fuel') = f fuel
72
                              fuelSet fuel'
73 74
                              return a

75 76
runFuelIO :: OptFuelState -> FuelUniqSM a -> IO a
runFuelIO fs (FUSM f) =
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
77 78 79 80 81 82 83
    do pass <- readIORef (pass_ref fs)
       fuel <- readIORef (fuel_ref fs)
       u    <- mkSplitUniqSupply 'u'
       let (a, FuelState fuel' pass') = initUs_ u $ f (FuelState fuel pass)
       writeIORef (pass_ref fs) pass'
       writeIORef (fuel_ref fs) fuel'
       return a
84

85 86 87 88 89 90 91 92 93 94
-- ToDo: Do we need the pass_ref when we are doing infinite fueld
-- transformations?
runInfiniteFuelIO :: OptFuelState -> FuelUniqSM a -> IO a
runInfiniteFuelIO fs (FUSM f) =
    do pass <- readIORef (pass_ref fs)
       u <- mkSplitUniqSupply 'u'
       let (a, FuelState _ pass') = initUs_ u $ f (FuelState unlimitedFuel pass)
       writeIORef (pass_ref fs) pass'
       return a

95 96 97
instance Monad FuelUniqSM where
  FUSM f >>= k = FUSM (\s -> f s >>= \(a, s') -> unFUSM (k a) s')
  return a     = FUSM (\s -> return (a, s))
98

99
instance MonadUnique FuelUniqSM where
dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
100 101 102
    getUniqueSupplyM = liftUniq getUniqueSupplyM
    getUniqueM       = liftUniq getUniqueM
    getUniquesM      = liftUniq getUniquesM
103 104 105

liftUniq :: UniqSM x -> FuelUniqSM x
liftUniq x = FUSM (\s -> x >>= (\u -> return (u, s)))
106

dias@eecs.harvard.edu's avatar
dias@eecs.harvard.edu committed
107
class Monad m => FuelUsingMonad m where
108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142
  fuelGet      :: m OptimizationFuel
  fuelSet      :: OptimizationFuel -> m ()
  lastFuelPass :: m String
  setFuelPass  :: String -> m ()

fuelExhausted :: FuelUsingMonad m => m Bool
fuelExhausted = fuelGet >>= return . anyFuelLeft

fuelDec1 :: FuelUsingMonad m => m ()
fuelDec1 = fuelGet >>= fuelSet . oneLessFuel

tryWithFuel :: FuelUsingMonad m => a -> m (Maybe a)
tryWithFuel r = do f <- fuelGet
                   if anyFuelLeft f then fuelSet (oneLessFuel f) >> return (Just r)
                                    else return Nothing

instance FuelUsingMonad FuelUniqSM where
  fuelGet          = extract fs_fuel
  lastFuelPass     = extract fs_lastpass
  fuelSet fuel     = FUSM (\s -> return ((), s { fs_fuel     = fuel }))
  setFuelPass pass = FUSM (\s -> return ((), s { fs_lastpass = pass }))

extract :: (FuelState -> a) -> FuelUniqSM a
extract f = FUSM (\s -> return (f s, s))

instance FuelMonad FuelUniqSM where
  getFuel = liftM amountOfFuel fuelGet
  setFuel = fuelSet . tankFilledTo

-- Don't bother to checkpoint the unique supply; it doesn't matter
instance CheckpointMonad FuelUniqSM where
    type Checkpoint FuelUniqSM = FuelState
    checkpoint = FUSM $ \fuel -> return (fuel, fuel) 
    restart fuel = FUSM $ \_ -> return ((), fuel)