118 lines
3.7 KiB
Text
118 lines
3.7 KiB
Text
/-
|
||
Copyright (c) 2022 Henrik Böving. All rights reserved.
|
||
Released under Apache 2.0 license as described in the file LICENSE.
|
||
Authors: Henrik Böving
|
||
-/
|
||
import Lean.Compiler.LCNF.PassManager
|
||
import Lean.Compiler.LCNF.PullLetDecls
|
||
import Lean.Compiler.LCNF.CSE
|
||
import Lean.Compiler.LCNF.Simp
|
||
import Lean.Compiler.LCNF.PullFunDecls
|
||
import Lean.Compiler.LCNF.ReduceJpArity
|
||
import Lean.Compiler.LCNF.JoinPoints
|
||
import Lean.Compiler.LCNF.Specialize
|
||
import Lean.Compiler.LCNF.PhaseExt
|
||
import Lean.Compiler.LCNF.ToMono
|
||
import Lean.Compiler.LCNF.LambdaLifting
|
||
|
||
namespace Lean.Compiler.LCNF
|
||
|
||
open PassInstaller
|
||
|
||
def init : Pass where
|
||
name := `init
|
||
run := fun decls => do
|
||
decls.forM (·.saveBase)
|
||
return decls
|
||
phase := .base
|
||
|
||
-- Helper pass used for debugging purposes
|
||
def trace (phase := Phase.base) : Pass where
|
||
name := `trace
|
||
run := pure
|
||
phase := phase
|
||
|
||
def normalizeFVarIds (decl : Decl) : CoreM Decl := do
|
||
let ngenSaved ← getNGen
|
||
setNGen {}
|
||
try
|
||
CompilerM.run <| decl.internalize
|
||
finally
|
||
setNGen ngenSaved
|
||
|
||
def saveBase : Pass :=
|
||
.mkPerDeclaration `saveBase (fun decl => do (← normalizeFVarIds decl).saveBase; return decl) .base
|
||
|
||
def saveMono : Pass :=
|
||
.mkPerDeclaration `saveMono (fun decl => do (← normalizeFVarIds decl).saveMono; return decl) .mono
|
||
|
||
def builtinPassManager : PassManager := {
|
||
passes := #[
|
||
init,
|
||
pullInstances,
|
||
cse,
|
||
simp,
|
||
findJoinPoints,
|
||
pullFunDecls,
|
||
reduceJpArity,
|
||
simp { etaPoly := true, inlinePartial := true, implementedBy := true } (occurrence := 1),
|
||
-- eagerLambdaLifting,
|
||
specialize,
|
||
simp (occurrence := 2),
|
||
cse,
|
||
saveBase, -- End of base phase
|
||
toMono,
|
||
simp (occurrence := 3) (phase := .mono),
|
||
reduceJpArity (phase := .mono),
|
||
extendJoinPointContext,
|
||
simp (occurrence := 4) (phase := .mono),
|
||
lambdaLifting,
|
||
simp (occurrence := 5) (phase := .mono),
|
||
-- TODO: reduce function arity
|
||
saveMono -- End of mono phase
|
||
]
|
||
}
|
||
|
||
def runImportedDecls (importedDeclNames : Array (Array Name)) : CoreM PassManager := do
|
||
let mut m := builtinPassManager
|
||
for declNames in importedDeclNames do
|
||
for declName in declNames do
|
||
m ← runFromDecl m declName
|
||
return m
|
||
|
||
builtin_initialize passManagerExt : PersistentEnvExtension Name (Name × PassManager) (List Name × PassManager) ←
|
||
registerPersistentEnvExtension {
|
||
mkInitial := return ([], builtinPassManager)
|
||
addImportedFn := fun ns => return ([], ← ImportM.runCoreM <| runImportedDecls ns)
|
||
addEntryFn := fun (installerDeclNames, _) (installerDeclName, managerNew) => (installerDeclName :: installerDeclNames, managerNew)
|
||
exportEntriesFn := fun s => s.1.reverse.toArray
|
||
}
|
||
|
||
def getPassManager : CoreM PassManager :=
|
||
return passManagerExt.getState (← getEnv) |>.2
|
||
|
||
def addPass (declName : Name) : CoreM Unit := do
|
||
let info ← getConstInfo declName
|
||
match info.type with
|
||
| .const `Lean.Compiler.LCNF.PassInstaller .. =>
|
||
let managerNew ← runFromDecl (← getPassManager) declName
|
||
modifyEnv fun env => passManagerExt.addEntry env (declName, managerNew)
|
||
| _ =>
|
||
throwError "invalid 'cpass' only 'PassInstaller's can be added via the 'cpass' attribute: {info.type}"
|
||
|
||
builtin_initialize
|
||
registerBuiltinAttribute {
|
||
name := `cpass
|
||
descr := "compiler passes for the code generator"
|
||
add := fun declName stx kind => do
|
||
Attribute.Builtin.ensureNoArgs stx
|
||
unless kind == AttributeKind.global do throwError "invalid attribute 'cpass', must be global"
|
||
discard <| addPass declName
|
||
applicationTime := .afterCompilation
|
||
}
|
||
|
||
builtin_initialize
|
||
registerTraceClass `Compiler.saveBase (inherited := true)
|
||
registerTraceClass `Compiler.trace (inherited := true)
|
||
|
||
end Lean.Compiler.LCNF
|