203 lines
8 KiB
Text
203 lines
8 KiB
Text
/-
|
||
Copyright (c) 2022 Microsoft Corporation. All rights reserved.
|
||
Released under Apache 2.0 license as described in the file LICENSE.
|
||
Authors: Leonardo de Moura
|
||
-/
|
||
import Lean.Compiler.Specialize
|
||
import Lean.Compiler.LCNF.FixedParams
|
||
import Lean.Compiler.LCNF.InferType
|
||
|
||
namespace Lean.Compiler.LCNF
|
||
|
||
/--
|
||
Each parameter is associated with a `SpecParamInfo`. This information is used by `LCNF/Specialize.lean`.
|
||
-/
|
||
inductive SpecParamInfo where
|
||
/--
|
||
A parameter that is an type class instance (or an arrow that produces a type class instance),
|
||
and is fixed in recursive declarations. By default, Lean always specializes this kind of argument.
|
||
-/
|
||
| fixedInst
|
||
/--
|
||
A parameter that is a function and is fixed in recursive declarations. If the user tags a declaration
|
||
with `@[specialize]` without specifying which arguments should be specialized, Lean will specialize
|
||
`.fixedHO` arguments in addition to `.fixedInst`.
|
||
-/
|
||
| fixedHO
|
||
/--
|
||
Computationally irrelevant parameters that are fixed in recursive declarations,
|
||
*and* there is a `fixedInst`, `fixedHO`, or `user` param that depends on it.
|
||
-/
|
||
| fixedNeutral
|
||
/--
|
||
An argument that has been specified in the `@[specialize]` attribute. Lean specializes it even if it is
|
||
not fixed in recursive declarations. Non-termination can happen, and Lean interrupts it with an error message
|
||
based on the stack depth.
|
||
-/
|
||
| user
|
||
/--
|
||
Parameter is not going to be specialized.
|
||
-/
|
||
| other
|
||
deriving Inhabited, Repr
|
||
|
||
instance : ToMessageData SpecParamInfo where
|
||
toMessageData
|
||
| .fixedInst => "I"
|
||
| .fixedHO => "H"
|
||
| .fixedNeutral => "N"
|
||
| .user => "U"
|
||
| .other => "O"
|
||
|
||
structure SpecState where
|
||
specInfo : SMap Name (Array SpecParamInfo) := {}
|
||
deriving Inhabited
|
||
|
||
structure SpecEntry where
|
||
declName : Name
|
||
paramsInfo : Array SpecParamInfo
|
||
deriving Inhabited
|
||
|
||
namespace SpecState
|
||
|
||
def addEntry (s : SpecState) (e : SpecEntry) : SpecState :=
|
||
match s with
|
||
| { specInfo } => { specInfo := specInfo.insert e.declName e.paramsInfo }
|
||
|
||
def switch : SpecState → SpecState
|
||
| { specInfo } => { specInfo := specInfo.switch }
|
||
|
||
end SpecState
|
||
|
||
/--
|
||
Extension for storing `SpecParamInfo` for declarations being compiled.
|
||
Remark: we only store information for declarations that will be specialized.
|
||
-/
|
||
builtin_initialize specExtension : SimplePersistentEnvExtension SpecEntry SpecState ←
|
||
registerSimplePersistentEnvExtension {
|
||
addEntryFn := SpecState.addEntry
|
||
addImportedFn := fun es => mkStateFromImportedEntries SpecState.addEntry {} es |>.switch
|
||
}
|
||
|
||
/--
|
||
Return `true` if `type` is a type tagged with `@[nospecialize]` or an arrow that produces this kind of type.
|
||
For example, this function returns true for `Inhabited Nat`, and `Nat → Inhabited Nat`.
|
||
-/
|
||
private def isNoSpecType (env : Environment) (type : Expr) : Bool :=
|
||
match type with
|
||
| .forallE _ _ b _ => isNoSpecType env b
|
||
| _ =>
|
||
if let .const declName _ := type.getAppFn then
|
||
hasNospecializeAttribute env declName
|
||
else
|
||
false
|
||
|
||
/-!
|
||
*Note*: `fixedNeutral` must have forward dependencies.
|
||
|
||
The code specializer consider a `fixedNeutral` parameter during code specialization
|
||
only if it contains forward dependecies that are tagged as `.user`, `.fixedHO`, or `.fixedInst`.
|
||
The motivation is to minimize the number of code specializations that have little or no impact on
|
||
performance. For example, let's consider the function.
|
||
```
|
||
def liftMacroM
|
||
{α : Type} {m : Type → Type}
|
||
[Monad m] [MonadMacroAdapter m] [MonadEnv m] [MonadRecDepth m] [MonadError m]
|
||
[MonadResolveName m] [MonadTrace m] [MonadOptions m] [AddMessageContext m] [MonadLiftT IO m] (x : MacroM α) : m α := do
|
||
```
|
||
The parameter `α` does not occur in any local instance, and `x` is marked as `.other` since the function
|
||
is not tagged as `[specialize]`. There is little value in considering `α` during code specialization,
|
||
but if we do many copies of this function will be generated.
|
||
Recall users may still force the code specializer to take `α` into account by using `[specialize α]` (`α` has `.user` info),
|
||
or `[specialize x]` (`α` has `.fixedNeutral` since `x` is a forward dependency tagged as `.user`),
|
||
or `[specialize]` (`α` has `.fixedNeutral` since `x` is a forward dependency tagged as `.fixedHO`).
|
||
-/
|
||
|
||
/--
|
||
Return `true` if parameter `j` of the given declaration has a forward dependency at parameter `k`,
|
||
and `k` is tagged as `.user`, `.fixedHO`, or `.fixedInst`.
|
||
|
||
See comment at `.fixedNeutral`.
|
||
-/
|
||
private def hasFwdDeps (decl : Decl) (paramsInfo : Array SpecParamInfo) (j : Nat) : Bool := Id.run do
|
||
let param := decl.params[j]!
|
||
for k in [j+1 : decl.params.size] do
|
||
if paramsInfo[k]! matches .user | .fixedHO | .fixedInst then
|
||
let param' := decl.params[k]!
|
||
if param'.type.containsFVar param.fvarId then
|
||
return true
|
||
return false
|
||
|
||
/--
|
||
Save parameter information for `decls`.
|
||
|
||
Remark: this function, similarly to `mkFixedArgMap`,
|
||
assumes that if a function `f` was declared in a mutual block, then `decls`
|
||
contains all (computationally relevant) functions in the mutual block.
|
||
-/
|
||
def saveSpecParamInfo (decls : Array Decl) : CompilerM Unit := do
|
||
let mut declsInfo := #[]
|
||
for decl in decls do
|
||
if hasNospecializeAttribute (← getEnv) decl.name then
|
||
declsInfo := declsInfo.push (mkArray decl.params.size .other)
|
||
else
|
||
let specArgs? := getSpecializationArgs? (← getEnv) decl.name
|
||
let contains (i : Nat) : Bool := specArgs?.getD #[] |>.contains i
|
||
let mut paramsInfo : Array SpecParamInfo := #[]
|
||
for i in [:decl.params.size] do
|
||
let param := decl.params[i]!
|
||
let info ←
|
||
if contains i then
|
||
pure .user
|
||
/-
|
||
If the user tagged class (e.g., `Inhabited`) with the `@[nospecialize]` attribute,
|
||
then parameters of this type should not be considered for specialization.
|
||
-/
|
||
else if isNoSpecType (← getEnv) param.type then
|
||
pure .other
|
||
else if isTypeFormerType param.type then
|
||
pure .fixedNeutral
|
||
else if (← isArrowClass? param.type).isSome then
|
||
pure .fixedInst
|
||
/-
|
||
Recall that if `specArgs? == some #[]`, then user annotated function with `@[specialize]`, but did not
|
||
specify which arguments must be specialized besides instances. In this case, we try to specialize
|
||
any "fixed higher-order argument"
|
||
-/
|
||
else if specArgs? == some #[] && param.type matches .forallE .. then
|
||
pure .fixedHO
|
||
else
|
||
pure .other
|
||
paramsInfo := paramsInfo.push info
|
||
pure ()
|
||
declsInfo := declsInfo.push paramsInfo
|
||
if declsInfo.any fun paramsInfo => paramsInfo.any (· matches .user | .fixedInst | .fixedHO) then
|
||
let m := mkFixedParamsMap decls
|
||
for i in [:decls.size] do
|
||
let decl := decls[i]!
|
||
let mut paramsInfo := declsInfo[i]!
|
||
let some mask := m.find? decl.name | unreachable!
|
||
trace[Compiler.specialize.info] "{decl.name} {mask}"
|
||
paramsInfo := paramsInfo.zipWith mask fun info fixed => if fixed || info matches .user then info else .other
|
||
for j in [:paramsInfo.size] do
|
||
let mut info := paramsInfo[j]!
|
||
if info matches .fixedNeutral && !hasFwdDeps decl paramsInfo j then
|
||
paramsInfo := paramsInfo.set! j .other
|
||
if paramsInfo.any fun info => info matches .fixedInst | .fixedHO | .user then
|
||
trace[Compiler.specialize.info] "{decl.name} {paramsInfo}"
|
||
modifyEnv fun env => specExtension.addEntry env { declName := decl.name, paramsInfo }
|
||
|
||
def getSpecParamInfoCore? (env : Environment) (declName : Name) : Option (Array SpecParamInfo) :=
|
||
(specExtension.getState env).specInfo.find? declName
|
||
|
||
def getSpecParamInfo? [Monad m] [MonadEnv m] (declName : Name) : m (Option (Array SpecParamInfo)) :=
|
||
return (specExtension.getState (← getEnv)).specInfo.find? declName
|
||
|
||
def isSpecCandidate [Monad m] [MonadEnv m] (declName : Name) : m Bool :=
|
||
return (specExtension.getState (← getEnv)).specInfo.contains declName
|
||
|
||
builtin_initialize
|
||
registerTraceClass `Compiler.specialize.info
|
||
|
||
end Lean.Compiler.LCNF
|
||
|