27 | module Compiler.Opts.CSE

 29 | import Core.CompileExpr

 30 | import Core.Context.Log

 31 | import Core.Options

 33 | import Core.Ord

 34 | import Data.String

 35 | import Data.SortedMap

 36 | import Data.Vect

 38 | import Libraries.Data.Erased

 39 | import Libraries.Data.List.SizeOf

 45 | public export

 46 | UsageMap : Type

 47 | UsageMap = SortedMap (Integer, ClosedCExp) (Name, Integer, Bool)

 59 | public export

 60 | data Count = Once | Many | C Integer

 62 | Show Count where

 63 |   show Once  = "Once"

 64 |   show Many  = "Many"

 65 |   show (C n) = "C " ++ show n

 73 | public export

 74 | ReplaceMap : Type

 75 | ReplaceMap = SortedMap Name (ClosedCExp, Count, Bool)

 77 | toReplaceMap : UsageMap -> ReplaceMap

 78 | toReplaceMap = SortedMap.fromList

 79 |              . map (\((_,exp),(n,c,d)) => (n, (exp, C c, d)))

 80 |              . SortedMap.toList

 86 | data Sts : Type where

 88 | record St where

 89 |   constructor MkSt

 90 |   map : UsageMap

 91 |   idx : Int

 92 |   inDelay : Bool

 98 | store : Ref Sts St => Integer -> ClosedCExp -> Core (Maybe Name)

 99 | store sz exp =

100 |   if sz < 5

101 |      then pure Nothing

102 |      else do

103 |        (MkSt map idx inDelay)    <- get Sts

105 |        (name,count,idx2,delayed) <-

106 |          case lookup (sz,exp) map of

107 |            Just (nm,cnt,delayed) => pure (nm, cnt+1, idx, delayed)

108 |            Nothing       => pure (MN "csegen" idx, 1, idx + 1, inDelay)

110 |        put Sts $ MkSt (insert (sz,exp) (name, count, inDelay || delayed) map) idx2 inDelay

111 |        pure (Just name)

117 | dropVar : SizeOf inner

118 |         -> {n : Nat}

119 |         -> (0 p : IsVar x n (inner ++ outer))

120 |         -> Maybe (Erased (IsVar x n inner))

121 | dropVar inn p = case locateIsVar inn p of

122 |   Left p => Just p

123 |   Right p => Nothing

129 | 0 Drop : Scoped -> Type

130 | Drop tm

131 |   = {0 inner, outer : Scope} ->

132 |     SizeOf inner ->

133 |     tm (inner ++ outer) ->

134 |     Maybe (tm inner)

138 |   dropCExp : Drop CExp

139 |   dropCExp inn (CLocal {idx} fc p) = (\ q => CLocal fc (runErased q)) <$> dropVar inn p

140 |   dropCExp inn (CRef fc x) = Just (CRef fc x)

141 |   dropCExp inn (CLam fc x y) = CLam fc x <$> dropCExp (suc inn) y

142 |   dropCExp inn (CLet fc x inlineOK y z) =

143 |     CLet fc x inlineOK <$> dropCExp inn y <*> dropCExp (suc inn) z

144 |   dropCExp inn (CApp fc x xs) = CApp fc <$> dropCExp inn x <*> traverse (dropCExp inn) xs

145 |   dropCExp inn (CCon fc x y tag xs) = CCon fc x y tag <$> traverse (dropCExp inn) xs

146 |   dropCExp inn (COp fc x xs) = COp fc x <$> traverse (dropCExp inn) xs

147 |   dropCExp inn (CExtPrim fc p xs) = CExtPrim fc p <$> traverse (dropCExp inn) xs

148 |   dropCExp inn (CForce fc x y) = CForce fc x <$> dropCExp inn y

149 |   dropCExp inn (CDelay fc x y) = CDelay fc x <$> dropCExp inn y

150 |   dropCExp inn (CConCase fc sc xs x) =

151 |     CConCase fc                  <$>

152 |     dropCExp inn sc              <*>

153 |     traverse (dropConAlt inn) xs <*>

154 |     traverse (dropCExp inn) x

156 |   dropCExp inn (CConstCase fc sc xs x) =

157 |     CConstCase fc                  <$>

158 |     dropCExp inn sc                <*>

159 |     traverse (dropConstAlt inn) xs <*>

160 |     traverse (dropCExp inn) x

162 |   dropCExp inn (CPrimVal fc x) = Just $ CPrimVal fc x

163 |   dropCExp inn (CErased fc) = Just $ CErased fc

164 |   dropCExp inn (CCrash fc x) = Just $ CCrash fc x

166 |   dropConAlt : Drop CConAlt

167 |   dropConAlt inn (MkConAlt x y tag args z) =

168 |     MkConAlt x y tag args <$>

169 |         dropCExp (mkSizeOf args + inn)

170 |         (replace {p = CExp} (appendAssociative args inner outer) z)

172 |   dropConstAlt : Drop CConstAlt

173 |   dropConstAlt inn (MkConstAlt x y) = MkConstAlt x <$> dropCExp inn y

181 |   -- Tries to convert an expression and its

182 |   -- sub-expression to closed terms (`CExp []`).

183 |   -- Every occurence of a closed term will be replaced by

184 |   -- a machine generated name and its count in the `UsageMap`

185 |   -- will be increased by one.

186 |   --

187 |   -- Because we start at the leafs and substitute all closed

188 |   -- expressions with machine generated names, we have to

189 |   -- keep track of the original expression's size, which

190 |   -- is returned in addition to the adjusted expressions.

191 |   export

192 |   analyze : Ref Sts St => CExp ns -> Core (Integer, CExp ns)

194 |   -- We ignore prim ops here, since moving them to the toplevel

195 |   -- might interfere with other optimizations, for instance

196 |   -- the one dealing with #1320.

197 |   -- Some other terms are ignored, as I (@stefan-hoeck)

198 |   -- am being conservative here,

199 |   -- not daring to inadvertently change the semantics

200 |   -- of the program.

201 |   analyze c@(COp {})      = analyzeSubExp c

202 |   analyze c@(CExtPrim {}) = analyzeSubExp c

203 |   analyze c@(CForce {})   = analyzeSubExp c

204 |   analyze c@(CDelay {})   = analyzeSubExp c

206 |   analyze exp = do

207 |     (sze, exp') <- analyzeSubExp exp

208 |     case dropCExp zero exp' of

209 |       Just e0 => do

210 |         Just nm <- store sze e0

211 |           | Nothing => pure (sze, exp')

212 |         pure (sze, CRef EmptyFC nm)

213 |       Nothing => pure (sze, exp')

215 |   analyzeList : Ref Sts St => List (CExp ns) -> Core (Integer, List (CExp ns))

216 |   analyzeList es = do

217 |     (sizes, exps) <- unzip <$> traverse analyze es

218 |     pure (sum sizes, exps)

220 |   analyzeMaybe : Ref Sts St => Maybe (CExp ns) -> Core (Integer, Maybe (CExp ns))

221 |   analyzeMaybe Nothing  = pure (0, Nothing)

222 |   analyzeMaybe (Just e) = do

223 |     (se,e') <- analyze e

224 |     pure (se, Just e')

226 |   analyzeVect : Ref Sts St => Vect n (CExp ns) -> Core (Integer, Vect n (CExp ns))

227 |   analyzeVect es = do

228 |     (sizes, exps) <- unzip <$> traverseVect analyze es

229 |     pure (sum sizes, exps)

232 |   analyzeSubExp : Ref Sts St => CExp ns -> Core (Integer, CExp ns)

233 |   analyzeSubExp e@(CLocal {}) = pure (1, e)

234 |   analyzeSubExp e@(CRef {})   = pure (1, e)

235 |   analyzeSubExp (CLam f n y)  = do

236 |     (sy, y') <- analyze y

237 |     pure (sy + 1, CLam f n y')

239 |   analyzeSubExp (CLet f n i y z) = do

240 |     (sy, y') <- analyze y

241 |     (sz, z') <- analyze z

242 |     pure (sy + sz + 1, CLet f n i y' z')

244 |   analyzeSubExp (CApp fc x xs) = do

245 |     (sx, x')   <- analyze x

246 |     (sxs, xs') <- analyzeList xs

247 |     pure (sx + sxs + 1, CApp fc x' xs')

249 |   analyzeSubExp (CCon f n c t xs) = do

250 |     (sxs, xs') <- analyzeList xs

251 |     pure (sxs + 1, CCon f n c t xs')

253 |   analyzeSubExp (COp f n xs) = do

254 |     (sxs, xs') <- analyzeVect xs

255 |     pure (sxs + 1, COp f n xs')

257 |   analyzeSubExp (CExtPrim f n xs) = do

258 |     (sxs, xs') <- analyzeList xs

259 |     pure (sxs + 1, CExtPrim f n xs')

261 |   analyzeSubExp (CForce f r y) = do

262 |     (sy, y') <- analyze y

263 |     pure (sy + 1, CForce f r y')

265 |   analyzeSubExp (CDelay f r y) = do

266 |     MkSt _ _ inDelay <- get Sts

267 |     update Sts (\(MkSt map idx _) => MkSt map idx True)

268 |     (sy, y') <- analyze y

269 |     update Sts (\(MkSt map idx _) => MkSt map idx inDelay)

270 |     pure (sy + 1, CDelay f r y')

272 |   analyzeSubExp (CConCase f sc xs x) = do

273 |     (ssc, sc') <- analyze sc

274 |     (sxs, xs') <- unzip <$> traverse analyzeConAlt xs

275 |     (sx, x')   <- analyzeMaybe x

276 |     pure (ssc + sum sxs + sx + 1, CConCase f sc' xs' x')

278 |   analyzeSubExp (CConstCase f sc xs x) = do

279 |     (ssc, sc') <- analyze sc

280 |     (sxs, xs') <- unzip <$> traverse analyzeConstAlt xs

281 |     (sx, x')   <- analyzeMaybe x

282 |     pure (ssc + sum sxs + sx + 1, CConstCase f sc' xs' x')

284 |   analyzeSubExp c@(CPrimVal {}) = pure (1, c)

285 |   analyzeSubExp c@(CErased {})  = pure (1, c)

286 |   analyzeSubExp c@(CCrash {})   = pure (1, c)

288 |   analyzeConAlt :  { auto c : Ref Sts St }

289 |                 -> CConAlt ns

290 |                 -> Core (Integer, CConAlt ns)

291 |   analyzeConAlt (MkConAlt n c t as z) = do

292 |     (sz, z') <- analyze z

293 |     pure (sz + 1, MkConAlt n c t as z')

295 |   analyzeConstAlt : Ref Sts St => CConstAlt ns -> Core (Integer, CConstAlt ns)

296 |   analyzeConstAlt (MkConstAlt c y) = do

297 |     (sy, y') <- analyze y

298 |     pure (sy + 1, MkConstAlt c y')

300 | analyzeDef : Ref Sts St => CDef -> Core CDef

301 | analyzeDef (MkFun args x)   = MkFun args . snd <$> analyze x

302 | analyzeDef d@(MkCon {})     = pure d

303 | analyzeDef d@(MkForeign {}) = pure d

304 | analyzeDef d@(MkError {})   = pure d

306 | compileName :  Ref Ctxt Defs

307 |             => Name

308 |             -> Core (Maybe (Name, FC, CDef))

309 | compileName fn = do

310 |     defs <- get Ctxt

311 |     Just def <- lookupCtxtExact fn (gamma defs)

312 |         | Nothing => do log "compile.execute" 50 $ "Couldn't find " ++ show fn

313 |                         pure Nothing

314 |     let Just cexp = compexpr def

315 |         | Nothing => do log "compile.execute" 50 $ "Couldn't compile " ++ show fn

316 |                         pure Nothing

317 |     pure $ Just (fn, location def, cexp)

324 |   -- During the analysis step, we replaced every closed

325 |   -- expression with a machine generated name. We only

326 |   -- want to keep these substitutions, if a closed term

327 |   -- appears in several distinct locations in the code

328 |   -- and does not appear inside `%delay`.

329 |   --

330 |   -- We therefore check for each machine generated name, if

331 |   -- the corresponding closed term has a count > 1. If that's

332 |   -- not the case, the machine generated name will be dropped

333 |   -- and replaces with the CSE-optimized original term.

334 |   --

335 |   -- However, during the analysis step, we might get counts > 1

336 |   -- for expressions that still are used only once.

337 |   -- Consider the following situation:

338 |   --

339 |   --    fun1 = exp1(exp2(exp3))

340 |   --    fun2 = exp4(exp2(exp3))

341 |   --

342 |   -- In the example above, `exp2` is a proper common subexpression

343 |   -- that should be lifted to the toplevel, but exp3 is not, although

344 |   -- it will also get a count of 2 in the `UsageMap`.

345 |   -- We therefore compare the count of each child expression

346 |   -- with the count of their parent expression, lifting

347 |   -- a child only if it was counted mor often than its parent.

348 |   replaceRef :  Ref ReplaceMap ReplaceMap

349 |              => Ref Ctxt Defs

350 |              => (parentCount : Integer)

351 |              -> FC

352 |              -> Name

353 |              -> Core (CExp ns)

354 |   replaceRef pc fc n = do

355 |     log "compiler.cse" 10 $ "Trying to replace " ++ show n ++ ": "

356 |     res <- lookup n <$> get ReplaceMap

357 |     case res of

358 |       -- not a name generated during CSE

359 |       Nothing          => do

360 |         log "compiler.cse" 10 $ "  not a name generated during CSE"

361 |         pure (CRef fc n)

363 |       -- Expression count has already been checked and occurs

364 |       -- several times. Replace it with the machine generated name.

365 |       Just (exp, Many, False) => do

366 |         log "compiler.cse" 10 $ "  already replaced: Occurs many times"

367 |         pure (CApp EmptyFC (CRef fc n) [])

369 |       -- Expression count has already been checked, it occurs

370 |       -- several times, but it has an occurrence inside `%delay`.

371 |       -- Substitute the machine generated name with the original

372 |       -- (CSE optimized) expression.

373 |       Just (exp, Many, True) => do

374 |         log "compiler.cse" 10 $ "  already replaced: Occurs inside %delay"

375 |         -- pure (embed exp)

376 |         pure (CForce EmptyFC LLazy (CApp EmptyFC (CRef fc n) []))

378 |       -- Expression count has already been checked and occurs

379 |       -- only once. Substitute the machine generated name with

380 |       -- the original (but CSE optimized) exp

381 |       Just (exp, Once, _) => do

382 |         log "compiler.cse" 10 $ "  already replaced: Occurs once"

383 |         pure (embed exp)

385 |       -- Expression count has not yet been compared with the

386 |       -- parent count. Do this now.

387 |       Just (exp, C c, d)  => do

388 |         log "compiler.cse" 10 $  "  expression of unknown quantity ("

389 |                               ++ show c

390 |                               ++ " occurences)"

391 |         -- We first have to replace all child expressions.

392 |         exp' <- replaceExp c exp

393 |         if c > pc

394 |            -- This is a common subexpression. We set its count to `Many`

395 |            -- and inspect its occurence in delay to check whether it

396 |            -- should be replaced or not.

397 |            then do

398 |              log "compiler.cse" 10 $ show n ++ " assigned quantity \"Many\""

399 |              update ReplaceMap (insert n (exp', Many, d))

400 |              case d of

401 |                   False => pure (CApp EmptyFC (CRef fc n) [])

402 |                   True => pure (CForce EmptyFC LLazy (CApp EmptyFC (CRef fc n) []))

404 |            -- This expression occurs only once. We set its count to `Once`

405 |            -- and keep it.

406 |            else do

407 |              log "compiler.cse" 10 $ show n ++ " assigned quantity \"Once\""

408 |              update ReplaceMap (insert n (exp', Once, d))

409 |              pure (embed exp')

412 |   replaceExp :  Ref ReplaceMap ReplaceMap

413 |              => Ref Ctxt Defs

414 |              => (parentCount : Integer)

415 |              -> CExp ns

416 |              -> Core (CExp ns)

417 |   replaceExp _ e@(CLocal {}) = pure e

418 |   replaceExp pc (CRef f n)   = replaceRef pc f n

419 |   replaceExp pc (CLam f n y) = CLam f n <$> replaceExp pc y

420 |   replaceExp pc (CLet f n i y z) =

421 |     CLet f n i <$> replaceExp pc y <*> replaceExp pc z

422 |   replaceExp pc (CApp f x xs) =

423 |     CApp f <$> replaceExp pc x <*> traverse (replaceExp pc) xs

424 |   replaceExp pc (CCon f n c t xs) =

425 |     CCon f n c t <$> traverse (replaceExp pc) xs

426 |   replaceExp pc (COp f n xs) =

427 |     COp f n <$> traverseVect (replaceExp pc) xs

428 |   replaceExp pc (CExtPrim f n xs) =

429 |     CExtPrim f n <$> traverse (replaceExp pc) xs

430 |   replaceExp pc (CForce f r y) =

431 |     CForce f r <$> replaceExp pc y

432 |   replaceExp pc (CDelay f r y) =

433 |     CDelay f r <$> replaceExp pc y

434 |   replaceExp pc (CConCase f sc xs x) =

435 |     CConCase f                     <$>

436 |     replaceExp pc sc               <*>

437 |     traverse (replaceConAlt pc) xs <*>

438 |     traverseOpt (replaceExp pc) x

440 |   replaceExp pc (CConstCase f sc xs x) = do

441 |     CConstCase f                     <$>

442 |     replaceExp pc sc                 <*>

443 |     traverse (replaceConstAlt pc) xs <*>

444 |     traverseOpt (replaceExp pc) x

446 |   replaceExp _ c@(CPrimVal {}) = pure c

447 |   replaceExp _ c@(CErased {})  = pure c

448 |   replaceExp _ c@(CCrash {})   = pure c

450 |   replaceConAlt :  Ref ReplaceMap ReplaceMap

451 |                 => Ref Ctxt Defs

452 |                 => (parentCount : Integer)

453 |                 -> CConAlt ns

454 |                 -> Core (CConAlt ns)

455 |   replaceConAlt pc (MkConAlt n c t as z) =

456 |     MkConAlt n c t as <$> replaceExp pc z

458 |   replaceConstAlt :  Ref ReplaceMap ReplaceMap

459 |                   => Ref Ctxt Defs

460 |                   => (parentCount : Integer)

461 |                   -> CConstAlt ns

462 |                   -> Core (CConstAlt ns)

463 |   replaceConstAlt pc (MkConstAlt c y) =

464 |     MkConstAlt c <$> replaceExp pc y

466 | replaceDef :  Ref ReplaceMap ReplaceMap

467 |            => Ref Ctxt Defs

468 |            => (Name, FC, CDef)

469 |            -> Core (Name, FC, CDef)

470 | replaceDef (n, fc, MkFun args x) =

471 |   (\x' => (n, fc, MkFun args x')) <$> replaceExp 1 x

472 | replaceDef (n, fc, d@(MkCon {}))     = pure (n, fc, d)

473 | replaceDef (n, fc, d@(MkForeign {})) = pure (n, fc, d)

474 | replaceDef (n, fc, d@(MkError {}))   = pure (n, fc, d)

476 | newToplevelDefs : ReplaceMap -> List (Name, FC, CDef)

477 | newToplevelDefs rm = mapMaybe toDef $ SortedMap.toList rm

478 |   where toDef : (Name,(ClosedCExp,Count,Bool)) -> Maybe (Name, FC, CDef)

479 |         toDef (nm,(exp,Many,False)) = Just (nm, EmptyFC, MkFun Scope.empty exp)

480 |         toDef (nm,(exp,Many,True)) = Just (nm, EmptyFC, MkFun Scope.empty (CDelay EmptyFC LLazy exp))

481 |         toDef _               = Nothing

483 | undefinedCount : (Name, (ClosedCExp, Count)) -> Bool

484 | undefinedCount (_, _, Once) = False

485 | undefinedCount (_, _, Many) = False

486 | undefinedCount (_, _, C x)  = True

491 | cse :  Ref Ctxt Defs

492 |     => (definitionNames : List Name)

493 |     -> (mainExpr        : CExp ns)

494 |     -> Core (List (Name, FC, CDef), CExp ns)

495 | cse defs me = do

496 |   compilerDefs <- get Ctxt

497 |   compiledDefs <- catMaybes <$> traverse compileName defs

498 |   if compilerDefs.options.session.noCSE

499 |     then pure (compiledDefs, me)

500 |     else do

501 |       log "compiler.cse" 10 $ "Analysing " ++ show (length defs) ++ " names"

502 |       s            <- newRef Sts $ MkSt empty 0 False

503 |       analyzedDefs <- traverse (traversePair (traversePair analyzeDef)) compiledDefs

504 |       MkSt um _ _  <- get Sts

505 |       srep         <- newRef ReplaceMap $ toReplaceMap um

506 |       replacedDefs <- traverse replaceDef analyzedDefs

507 |       replacedMain <- replaceExp 1 me

508 |       replaceMap   <- get ReplaceMap

509 |       let filtered = SortedMap.toList replaceMap

510 |       log "compiler.cse" 10 $ unlines $

511 |         "Found the following unadjusted subexpressions:"

512 |         ::  map (\(name,(_,cnt)) =>

513 |                       show name ++ ": count " ++ show cnt

514 |                ) filtered

515 |       let newDefs := newToplevelDefs replaceMap ++ replacedDefs

516 |       pure (newDefs, replacedMain)