<commit_before>{-# LANGUAGE GeneralizedNewtypeDeriving, NoImplicitPrelude #-}
{-# OPTIONS -Wall #-}
module Language.Paraiso.Tuning.Genetic
  (
   Genome, Species(..),
   makeSpecies,
   readGenome, overwriteGenome,
   mutate, cross, triangulate,
   generateIO
  ) where

import qualified Control.Monad.State as State
import qualified Data.Graph.Inductive                   as FGL
import qualified Data.Vector as V
import           Data.Vector ((!))
import qualified Language.Paraiso.Annotation as Anot
import qualified Language.Paraiso.Annotation.Allocation as Alloc
import qualified Language.Paraiso.Generator.Native as Native
import qualified Language.Paraiso.OM as OM
import qualified Language.Paraiso.OM.Graph as OM
import qualified Language.Paraiso.Optimization          as Opt
import           Language.Paraiso.Prelude
import qualified Language.Paraiso.Generator as Gen  (generateIO) 
import qualified Prelude as Prelude
import           System.Random 
import qualified Text.Read as Read

data Species v g = 
  Species {
    setup   :: Native.Setup v g,
    machine :: OM.OM v g Anot.Annotation
  } deriving (Show)

makeSpecies :: Native.Setup v g -> OM.OM v g Anot.Annotation -> Species v g
makeSpecies = Species

generateIO :: (Opt.Ready v g) => Species v g -> IO [(FilePath, Text)]
generateIO (Species s om) = Gen.generateIO s om

newtype Genome = Genome [Bool] deriving (Eq)

instance Show Genome where
  show (Genome xs) = show $ toDNA xs

instance Read Genome where
  readPrec = fmap (Genome . fromDNA) Read.readPrec

toDNA :: [Bool] -> String
toDNA xss@(x:xs)
  | even $ length xss = 'C' : inner xss
  | not x             = 'A' : inner xs
  | otherwise         = 'T' : inner xs
  where
    inner [] = ""
    inner (x:y:xs) = inner1 x y : inner xs
    inner _ = error "parity conservation law is broken"

    inner1 False False = 'A'
    inner1 False True  = 'C'
    inner1 True  False = 'G'
    inner1 True  True  = 'T'

fromDNA :: String -> [Bool]
fromDNA xss@(x:xs) 
  | x == 'C'  = inner xs
  | x == 'A'  = False : inner xs
  | x == 'T'  = True  : inner xs
  | otherwise = error "bad DNA"
  where
    inner = concat . map inner1
    inner1 :: Char -> [Bool]
    inner1 'A' = [False, False]
    inner1 'C' = [False, True ]
    inner1 'G' = [True , False]
    inner1 'T' = [True , True ]
    inner1 _   = error "bad DNA"

mutate :: Genome -> IO Genome
mutate (Genome xs) = do
  let oldVector = V.fromList xs
      n = length xs
      logN :: Double
      logN = log (fromIntegral n)
  logRand <- randomRIO (0, logN - 2)
  let randN :: Int
      randN = Prelude.max 1 $ ceiling $ exp logRand
      randRanges = V.replicate randN (0, n - 1)
      randUpd range = do
        idx <- randomRIO range
        return (idx, not $ oldVector ! idx)
  randUpds <- V.mapM randUpd randRanges
  return $ Genome $ V.toList $ V.update oldVector randUpds

cross :: Genome -> Genome -> IO Genome
cross (Genome xs) (Genome ys) = do
  let n = length xs
      vx = V.fromList xs
      vy = V.fromList ys
      atLeast :: Int -> IO Int
      atLeast n = do
        coin <- randomRIO (0,1)
        if coin < (0.5 :: Double) 
           then return n
           else atLeast (n+1)
  randN <- atLeast 1
  let randRanges = replicate randN (-1, n + 1)
  crossPoints <- mapM randomRIO randRanges
  let vz = V.generate n $ \i ->
           if odd $ length $ filter (<i) crossPoints then vx!i else vy!i
  return $ Genome $ V.toList $ vz

triangulate :: Genome -> Genome -> Genome -> IO Genome
triangulate (Genome base) (Genome left) (Genome right) = do
  return $ Genome $ zipWith3 f base left right
    where
      f b l r = if b/=l then l else r

readGenome :: Species v g -> Genome
readGenome spec =
  encode $ do
    let (x,y) = Native.cudaGridSize $ setup spec
    putInt 16 x
    putInt 16 y
    let om    = machine spec
        kerns = OM.kernels om
    V.mapM_ (putGraph . OM.dataflow) kerns

overwriteGenome :: (Opt.Ready v g) => Genome -> Species v g -> Species v g
overwriteGenome dna oldSpec = 
  decode dna $ do
    -- load cuda grid topology from the genome
    x <- getInt 16
    y <- getInt 16
    let oldSetup = setup oldSpec
        oldOM = machine oldSpec
        oldKernels = OM.kernels oldOM
        oldFlags = OM.globalAnnotation $ OM.setup $ oldOM
    let overwriteKernel kern = do
                 let graph = OM.dataflow kern
                 newGraph <- overwriteGraph graph
                 return $ kern{OM.dataflow = newGraph}
    -- load manifesto from the genome
    newKernels <- V.mapM overwriteKernel oldKernels
    let newGrid = (x,y)
        newSetup = oldSetup {Native.cudaGridSize = newGrid}
        -- reset the optimization flag and cause the optimization again
        newFlags = Anot.set Opt.Unoptimized oldFlags
        newOM = oldOM 
          { OM.kernels = newKernels, 
            OM.setup   = (OM.setup oldOM){ OM.globalAnnotation = newFlags }
          }
    return $ Species newSetup (Opt.optimize Opt.O3 newOM)



newtype Get a = Get { getGet :: State.State [Bool] a } deriving (Monad)
newtype Put a = Put { getPut :: State.State [Bool] a } deriving (Monad)

get :: Get Bool
get = Get $ do
  (x:xs) <- State.get
  State.put xs
  return x

put :: Bool -> Put ()
put x = Put $ do
  xs <- State.get
  State.put $ x:xs

decode :: Genome -> Get a -> a
decode (Genome dna) m = State.evalState (getGet m) dna

encode :: Put a -> Genome 
encode m = Genome $ reverse $ State.execState (getPut m) []

putInt :: Int -> Int -> Put ()
putInt bit n 
  | bit <= 0  = return ()
  | n >= val  = do
    put True
    putInt (bit-1) (n-val)
  | otherwise = do
    put False
    putInt (bit-1) n
  where 
    val :: Int
    val = 2^(fromIntegral $ bit-1)
    
getInt :: Int -> Get Int
getInt bit
  | bit <= 0  = return 0
  | otherwise = do
    x <- get
    y <- getInt (bit-1)
    return $ y + 2^(fromIntegral $ bit-1) * (if x then 1 else 0)

putGraph :: OM.Graph v g Anot.Annotation -> Put ()
putGraph graph = do
  V.mapM_ put focus
  where
    focus = 
      V.map (isManifest . snd) $
      V.filter (hasChoice . snd) anots

    anots :: V.Vector (FGL.Node, Anot.Annotation)
    anots = V.fromList $ map (\(n, lab) -> (n, OM.getA lab)) $ FGL.labNodes graph
    
    hasChoice :: Anot.Annotation -> Bool
    hasChoice anot = 
      case Anot.toMaybe anot of
        Just (Alloc.AllocationChoice _) -> True
        _                               -> False
    
    isManifest :: Anot.Annotation -> Bool
    isManifest anot = 
      case Anot.toMaybe anot of
        Just Alloc.Manifest -> True
        _                   -> False



overwriteGraph :: OM.Graph v g Anot.Annotation -> Get (OM.Graph v g Anot.Annotation)
overwriteGraph graph = do
  ovs <- V.mapM getAt focusIndices
  return $ overwritten ovs
  where
    overwritten ovs = 
      let newManifest ::  V.Vector (Maybe Alloc.Allocation)
          newManifest = 
            flip V.update ovs $
            V.map (const Nothing) anots in
      flip OM.imap graph $ \idx anot -> 
        case newManifest ! idx of
          Nothing -> anot
          Just x  -> Anot.set x anot

    getAt idx = do
      ret <- get
      return (idx, Just $ if ret then Alloc.Manifest else Alloc.Delayed)

    focusIndices = 
      V.map fst $
      V.filter (hasChoice . snd) anots

    anots :: V.Vector (FGL.Node, Anot.Annotation)
    anots = V.fromList $ map (\(n, lab) -> (n, OM.getA lab)) $ FGL.labNodes graph

    
    hasChoice :: Anot.Annotation -> Bool
    hasChoice anot = 
      case Anot.toMaybe anot of
        Just (Alloc.AllocationChoice _) -> True
        _                               -> False
    
    isManifest :: Anot.Annotation -> Bool
    isManifest anot = 
      case Anot.toMaybe anot of
        Just Alloc.Manifest -> True
        _                   -> False


<commit_msg>mutate is improved<commit_after>{-# LANGUAGE GeneralizedNewtypeDeriving, NoImplicitPrelude #-}
{-# OPTIONS -Wall #-}
module Language.Paraiso.Tuning.Genetic
  (
   Genome, Species(..),
   makeSpecies,
   readGenome, overwriteGenome,
   mutate, cross, triangulate,
   generateIO
  ) where

import qualified Control.Monad.State as State
import qualified Data.Graph.Inductive                   as FGL
import qualified Data.Vector as V
import           Data.Vector ((!))
import qualified Language.Paraiso.Annotation as Anot
import qualified Language.Paraiso.Annotation.Allocation as Alloc
import qualified Language.Paraiso.Generator.Native as Native
import qualified Language.Paraiso.OM as OM
import qualified Language.Paraiso.OM.Graph as OM
import qualified Language.Paraiso.Optimization          as Opt
import           Language.Paraiso.Prelude
import qualified Language.Paraiso.Generator as Gen  (generateIO) 
import qualified Prelude as Prelude
import           System.Random 
import qualified Text.Read as Read

data Species v g = 
  Species {
    setup   :: Native.Setup v g,
    machine :: OM.OM v g Anot.Annotation
  } deriving (Show)

makeSpecies :: Native.Setup v g -> OM.OM v g Anot.Annotation -> Species v g
makeSpecies = Species

generateIO :: (Opt.Ready v g) => Species v g -> IO [(FilePath, Text)]
generateIO (Species s om) = Gen.generateIO s om

newtype Genome = Genome [Bool] deriving (Eq)

instance Show Genome where
  show (Genome xs) = show $ toDNA xs

instance Read Genome where
  readPrec = fmap (Genome . fromDNA) Read.readPrec

toDNA :: [Bool] -> String
toDNA xss@(x:xs)
  | even $ length xss = 'C' : inner xss
  | not x             = 'A' : inner xs
  | otherwise         = 'T' : inner xs
  where
    inner [] = ""
    inner (x:y:xs) = inner1 x y : inner xs
    inner _ = error "parity conservation law is broken"

    inner1 False False = 'A'
    inner1 False True  = 'C'
    inner1 True  False = 'G'
    inner1 True  True  = 'T'

fromDNA :: String -> [Bool]
fromDNA xss@(x:xs) 
  | x == 'C'  = inner xs
  | x == 'A'  = False : inner xs
  | x == 'T'  = True  : inner xs
  | otherwise = error "bad DNA"
  where
    inner = concat . map inner1
    inner1 :: Char -> [Bool]
    inner1 'A' = [False, False]
    inner1 'C' = [False, True ]
    inner1 'G' = [True , False]
    inner1 'T' = [True , True ]
    inner1 _   = error "bad DNA"

mutate :: Genome -> IO Genome
mutate original@(Genome xs) = do
  let oldVector = V.fromList xs
      n = length xs
      logN :: Double
      logN = log (fromIntegral n)
  logRand <- randomRIO (-1, logN)
  let randN :: Int
      randN = Prelude.max 1 $ ceiling $ exp logRand
      randRanges = V.replicate randN (0, n - 1)
      randUpd range = do
        idx <- randomRIO range
        return (idx, not $ oldVector ! idx)
  randUpds <- V.mapM randUpd randRanges
  let pureMutant = Genome $ V.toList $ V.update oldVector randUpds
  if logRand > logN - 2 
     then return pureMutant >>= cross original >>= cross original
     else return pureMutant

cross :: Genome -> Genome -> IO Genome
cross (Genome xs0) (Genome ys0) = do
  swapCoin <- randomRIO (0,1)
  let 
    (xs,ys) = if swapCoin < (0.5::Double) then (xs0, ys0) else (ys0, xs0)
    n = length xs
    vx = V.fromList xs
    vy = V.fromList ys
    atLeast :: Int -> IO Int
    atLeast n = do
      coin <- randomRIO (0,1)
      if coin < (0.5 :: Double) 
         then return n
         else atLeast (n+1)
  randN <- atLeast 1
  let randRanges = replicate randN (-1, n + 1)
  crossPoints <- mapM randomRIO randRanges
  let vz = V.generate n $ \i ->
           if odd $ length $ filter (<i) crossPoints then vx!i else vy!i
  return $ Genome $ V.toList $ vz

triangulate :: Genome -> Genome -> Genome -> IO Genome
triangulate (Genome base) (Genome left) (Genome right) = do
  return $ Genome $ zipWith3 f base left right
    where
      f b l r = if b/=l then l else r

readGenome :: Species v g -> Genome
readGenome spec =
  encode $ do
    let (x,y) = Native.cudaGridSize $ setup spec
    putInt 16 x
    putInt 16 y
    let om    = machine spec
        kerns = OM.kernels om
    V.mapM_ (putGraph . OM.dataflow) kerns

overwriteGenome :: (Opt.Ready v g) => Genome -> Species v g -> Species v g
overwriteGenome dna oldSpec = 
  decode dna $ do
    -- load cuda grid topology from the genome
    x <- getInt 16
    y <- getInt 16
    let oldSetup = setup oldSpec
        oldOM = machine oldSpec
        oldKernels = OM.kernels oldOM
        oldFlags = OM.globalAnnotation $ OM.setup $ oldOM
    let overwriteKernel kern = do
                 let graph = OM.dataflow kern
                 newGraph <- overwriteGraph graph
                 return $ kern{OM.dataflow = newGraph}
    -- load manifesto from the genome
    newKernels <- V.mapM overwriteKernel oldKernels
    let newGrid = (x,y)
        newSetup = oldSetup {Native.cudaGridSize = newGrid}
        -- reset the optimization flag and cause the optimization again
        newFlags = Anot.set Opt.Unoptimized oldFlags
        newOM = oldOM 
          { OM.kernels = newKernels, 
            OM.setup   = (OM.setup oldOM){ OM.globalAnnotation = newFlags }
          }
    return $ Species newSetup (Opt.optimize Opt.O3 newOM)



newtype Get a = Get { getGet :: State.State [Bool] a } deriving (Monad)
newtype Put a = Put { getPut :: State.State [Bool] a } deriving (Monad)

get :: Get Bool
get = Get $ do
  (x:xs) <- State.get
  State.put xs
  return x

put :: Bool -> Put ()
put x = Put $ do
  xs <- State.get
  State.put $ x:xs

decode :: Genome -> Get a -> a
decode (Genome dna) m = State.evalState (getGet m) dna

encode :: Put a -> Genome 
encode m = Genome $ reverse $ State.execState (getPut m) []

putInt :: Int -> Int -> Put ()
putInt bit n 
  | bit <= 0  = return ()
  | n >= val  = do
    put True
    putInt (bit-1) (n-val)
  | otherwise = do
    put False
    putInt (bit-1) n
  where 
    val :: Int
    val = 2^(fromIntegral $ bit-1)
    
getInt :: Int -> Get Int
getInt bit
  | bit <= 0  = return 0
  | otherwise = do
    x <- get
    y <- getInt (bit-1)
    return $ y + 2^(fromIntegral $ bit-1) * (if x then 1 else 0)

putGraph :: OM.Graph v g Anot.Annotation -> Put ()
putGraph graph = do
  V.mapM_ put focus
  where
    focus = 
      V.map (isManifest . snd) $
      V.filter (hasChoice . snd) anots

    anots :: V.Vector (FGL.Node, Anot.Annotation)
    anots = V.fromList $ map (\(n, lab) -> (n, OM.getA lab)) $ FGL.labNodes graph
    
    hasChoice :: Anot.Annotation -> Bool
    hasChoice anot = 
      case Anot.toMaybe anot of
        Just (Alloc.AllocationChoice _) -> True
        _                               -> False
    
    isManifest :: Anot.Annotation -> Bool
    isManifest anot = 
      case Anot.toMaybe anot of
        Just Alloc.Manifest -> True
        _                   -> False



overwriteGraph :: OM.Graph v g Anot.Annotation -> Get (OM.Graph v g Anot.Annotation)
overwriteGraph graph = do
  ovs <- V.mapM getAt focusIndices
  return $ overwritten ovs
  where
    overwritten ovs = 
      let newManifest ::  V.Vector (Maybe Alloc.Allocation)
          newManifest = 
            flip V.update ovs $
            V.map (const Nothing) anots in
      flip OM.imap graph $ \idx anot -> 
        case newManifest ! idx of
          Nothing -> anot
          Just x  -> Anot.set x anot

    getAt idx = do
      ret <- get
      return (idx, Just $ if ret then Alloc.Manifest else Alloc.Delayed)

    focusIndices = 
      V.map fst $
      V.filter (hasChoice . snd) anots

    anots :: V.Vector (FGL.Node, Anot.Annotation)
    anots = V.fromList $ map (\(n, lab) -> (n, OM.getA lab)) $ FGL.labNodes graph

    
    hasChoice :: Anot.Annotation -> Bool
    hasChoice anot = 
      case Anot.toMaybe anot of
        Just (Alloc.AllocationChoice _) -> True
        _                               -> False
    
    isManifest :: Anot.Annotation -> Bool
    isManifest anot = 
      case Anot.toMaybe anot of
        Just Alloc.Manifest -> True
        _                   -> False



{'fraction_non_alphanumeric': 0.10119992381436099, 'fraction_numerical': 0.006602755380610755, 'mean_word_length': 2.181579478893153, 'pattern_counts': {'":': 0, '<': 45, '<?xml version=': 0, '>': 100, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7466806', 'n_tokens_mistral': 5437, 'n_tokens_neox': 5114, 'n_words': 1749}