perf: do not import non-meta IR

.github/workflows/build-template.yml

@@ -107,6 +107,7 @@ jobs:
             ${{ matrix.name == 'Linux Lake' && false && 'build/stage1/**/*.trace
             build/stage1/**/*.olean*
             build/stage1/**/*.ilean
+            build/stage1/**/*.ir
             build/stage1/**/*.c
             build/stage1/**/*.c.o*' || '' }}
           key: ${{ matrix.name }}-build-v3-${{ github.sha }}
@@ -248,6 +249,7 @@ jobs:
             ${{ matrix.name == 'Linux Lake' && false && 'build/stage1/**/*.trace
             build/stage1/**/*.olean*
             build/stage1/**/*.ilean
+            build/stage1/**/*.ir
             build/stage1/**/*.c
             build/stage1/**/*.c.o*' || '' }}
           key: ${{ steps.restore-cache.outputs.cache-primary-key }}

src/Lean/Compiler/ExternAttr.lean

@@ -17,7 +17,8 @@ inductive ExternEntry where
   | adhoc    (backend : Name)
   | inline   (backend : Name) (pattern : String)
   | standard (backend : Name) (fn : String)
-  | foreign  (backend : Name) (fn : String)
+  /-- Call to a Lean function without exported IR. -/
+  | opaque   (fn : Name)
   deriving BEq, Hashable
 
 /--
@@ -112,7 +113,7 @@ def ExternEntry.backend : ExternEntry → Name
   | ExternEntry.adhoc n      => n
   | ExternEntry.inline n _   => n
   | ExternEntry.standard n _ => n
-  | ExternEntry.foreign n _  => n
+  | ExternEntry.opaque ..    => `all
 
 def getExternEntryForAux (backend : Name) : List ExternEntry → Option ExternEntry
   | []    => none
@@ -139,7 +140,6 @@ def getExternNameFor (env : Environment) (backend : Name) (fn : Name) : Option S
   let entry ← getExternEntryFor data? backend
   match entry with
   | ExternEntry.standard _ n => pure n
-  | ExternEntry.foreign _ n  => pure n
   | _ => failure
 
 private def getExternConstArity (declName : Name) : CoreM Nat := do

src/Lean/Compiler/IR/CompilerM.lean

@@ -7,6 +7,8 @@ prelude
 import Lean.Environment
 import Lean.Compiler.IR.Basic
 import Lean.Compiler.IR.Format
+import Lean.Compiler.MetaAttr
+import Lean.Compiler.ExportAttr
 
 namespace Lean.IR
 
@@ -82,13 +84,73 @@ private abbrev findAtSorted? (decls : Array Decl) (declName : Name) : Option Dec
   let tmpDecl := Decl.extern declName #[] default default
   decls.binSearch tmpDecl declLt
 
+namespace CollectUsedFDecls
+
+abbrev M := StateM NameSet
+
+@[inline] def collect (f : FunId) : M Unit :=
+  modify fun s => s.insert f
+
+partial def collectFnBody : FnBody → M Unit
+  | .vdecl _ _ v b   =>
+    match v with
+    | .fap f _ => collect f *> collectFnBody b
+    | .pap f _ => collect f *> collectFnBody b
+    | _        => collectFnBody b
+  | .jdecl _ _ v b   => collectFnBody v *> collectFnBody b
+  | .case _ _ _ alts => alts.forM fun alt => collectFnBody alt.body
+  | e => unless e.isTerminal do collectFnBody e.body
+
+def collectDecl : Decl → M NameSet
+  | .fdecl (body := b) .. => collectFnBody b *> get
+  | .extern .. => get
+
+end CollectUsedFDecls
+
+/-- Adds to `used` all `Decl.fdecl`s referenced directly by `decl`. -/
+def collectUsedFDecls (decl : Decl) (used : NameSet := {}) : NameSet :=
+  (CollectUsedFDecls.collectDecl decl).run' used
+
+/-- Computes the closure of `Decl.fdecl`s referenced by `decl`. -/
+def getFDeclClosure (m : DeclMap) (decls : Array Decl) : NameSet := Id.run do
+  let mut toVisit := decls.map (·.name) |>.toList
+  let mut res : NameSet := .ofList toVisit
+  while !toVisit.isEmpty do
+    let n :: toVisit' := toVisit | continue
+    toVisit := toVisit'
+    let some d := m.find? n | continue
+    for d' in collectUsedFDecls d do
+      if !res.contains d' then
+        res := res.insert d'
+        toVisit := d' :: toVisit
+  return res
+
 builtin_initialize declMapExt : SimplePersistentEnvExtension Decl DeclMap ←
   registerSimplePersistentEnvExtension {
     addImportedFn := fun _ => {}
     addEntryFn    := fun s d => s.insert d.name d
-    toArrayFn     := fun s =>
-      let decls := s.foldl (init := #[]) fun decls decl => decls.push decl
-      sortDecls decls
+    -- Store `meta` closure only in `.olean`, turn all other decls into opaque externs.
+    -- Leave storing the remainder for `meta import` and server `#eval` to `exportIREntries` below.
+    exportEntriesFnEx? := some fun env s entries level =>
+      let decls := entries.foldl (init := #[]) fun decls decl => decls.push decl
+      let entries := sortDecls decls
+      let metaClosure := getFDeclClosure s (decls.filter (isMeta env ·.name))
+      if env.header.isModule then
+        entries.map fun
+          | d@(.fdecl f xs ty b info) =>
+            -- The interpreter may call the boxed variant even if the IR does not directly reference
+            -- it.
+            let n := match f with
+              | .str n "_boxed" => n
+              | n => n
+            if metaClosure.contains n then
+              d
+            else if let some (.str _ s) := getExportNameFor? env n then
+              .extern f xs ty { arity? := xs.size, entries := [.standard `all s] }
+            else
+              .extern f xs ty { arity? := xs.size, entries := [.opaque f] }
+          | d => d
+      else entries
     -- Written to on codegen environment branch but accessed from other elaboration branches when
     -- calling into the interpreter. We cannot use `async` as the IR declarations added may not
     -- share a name prefix with the top-level Lean declaration being compiled, e.g. from
@@ -97,12 +159,32 @@ builtin_initialize declMapExt : SimplePersistentEnvExtension Decl DeclMap ←
     replay?       := some <| SimplePersistentEnvExtension.replayOfFilter (!·.contains ·.name) (fun s d => s.insert d.name d)
   }
 
-@[export lean_ir_find_env_decl]
+@[export lean_ir_export_entries]
+private def exportIREntries (env : Environment) : Array (Name × Array EnvExtensionEntry) :=
+  let decls := declMapExt.getEntries env |>.foldl (init := #[]) fun decls decl => decls.push decl
+  -- safety: cast to erased type
+  let entries : Array EnvExtensionEntry := unsafe unsafeCast <| sortDecls decls
+  #[(``declMapExt, entries)]
+
+/-- Retrieves IR for codegen purposes, i.e. independent of `meta import`. -/
 def findEnvDecl (env : Environment) (declName : Name) : Option Decl :=
   match env.getModuleIdxFor? declName with
   | some modIdx => findAtSorted? (declMapExt.getModuleEntries env modIdx) declName
   | none        => declMapExt.getState env |>.find? declName
 
+@[export lean_ir_find_env_decl]
+private def findInterpreterDecl (env : Environment) (declName : Name) : Option Decl :=
+  match env.getModuleIdxFor? declName with
+  | some modIdx => do
+    let decl ←
+      -- `meta import` and server `#eval`
+      findAtSorted? (declMapExt.getModuleIREntries env modIdx) declName <|>
+      -- (closure of) `meta def`; will report `.extern`s for other `def`s so needs to come second
+      findAtSorted? (declMapExt.getModuleEntries env modIdx) declName
+    guard !decl matches .extern _ _ _ { entries := [.opaque _], .. }
+    return decl
+  | none => declMapExt.getState env |>.find? declName
+
 def findDecl (n : Name) : CompilerM (Option Decl) :=
   return findEnvDecl (← get).env n
 

src/Lean/Compiler/IR/EmitC.lean

@@ -422,18 +422,17 @@ def emitExternCall (f : FunId) (ps : Array Param) (extData : ExternAttrData) (ys
   match getExternEntryFor extData `c with
   | some (ExternEntry.standard _ extFn) => emitSimpleExternalCall extFn ps ys
   | some (ExternEntry.inline _ pat)     => do emit (expandExternPattern pat (toStringArgs ys)); emitLn ";"
-  | some (ExternEntry.foreign _ extFn)  => emitSimpleExternalCall extFn ps ys
   | _ => throw s!"failed to emit extern application '{f}'"
 
 def emitFullApp (z : VarId) (f : FunId) (ys : Array Arg) : M Unit := do
   emitLhs z
   let decl ← getDecl f
   match decl with
-  | Decl.extern _ ps _ extData => emitExternCall f ps extData ys
-  | _ =>
+  | .fdecl .. | .extern _ _ _ { entries := [.opaque _], .. } =>
     emitCName f
     if ys.size > 0 then emit "("; emitArgs ys; emit ")"
     emitLn ";"
+  | Decl.extern _ ps _ extData => emitExternCall f ps extData ys
 
 def emitPartialApp (z : VarId) (f : FunId) (ys : Array Arg) : M Unit := do
   let decl ← getDecl f

src/Lean/Compiler/IR/EmitLLVM.lean

@@ -665,7 +665,7 @@ def emitExternCall (builder : LLVM.Builder llvmctx)
   | some (ExternEntry.standard _ extFn) => emitSimpleExternalCall builder extFn ps ys retty name
   | some (ExternEntry.inline `llvm _pat) => throw "Unimplemented codegen of inline LLVM"
   | some (ExternEntry.inline _ pat) => throw s!"Cannot codegen non-LLVM inline code '{pat}'."
-  | some (ExternEntry.foreign _ extFn)  => emitSimpleExternalCall builder extFn ps ys retty name
+  | some (ExternEntry.opaque _)  => unreachable!
   | _ => throw s!"Failed to emit extern application '{f}'."
 
 def getFunIdTy (f : FunId) : M llvmctx (LLVM.LLVMType llvmctx) := do
@@ -740,10 +740,7 @@ def emitFullApp (builder : LLVM.Builder llvmctx)
   let (__zty, zslot) ← emitLhsSlot_ z
   let decl ← getDecl f
   match decl with
-  | Decl.extern _ ps retty extData =>
-     let zv ← emitExternCall builder f ps extData ys retty
-     LLVM.buildStore builder zv zslot
-  | Decl.fdecl .. =>
+  | .fdecl .. | .extern _ _ _ { entries := [.opaque _] } =>
     if ys.size > 0 then
         let fv ← getOrAddFunIdValue builder f
         let ys ←  ys.mapM (fun y => do
@@ -755,6 +752,9 @@ def emitFullApp (builder : LLVM.Builder llvmctx)
     else
        let zv ← getOrAddFunIdValue builder f
        LLVM.buildStore builder zv zslot
+  | Decl.extern _ ps retty extData =>
+     let zv ← emitExternCall builder f ps extData ys retty
+     LLVM.buildStore builder zv zslot
 
 -- Note that this returns a *slot*, just like `emitLhsSlot_`.
 def emitLit (builder : LLVM.Builder llvmctx)

src/Lean/DocString/Extension.lean

@@ -58,9 +58,11 @@ structure ModuleDoc where
 private builtin_initialize moduleDocExt : SimplePersistentEnvExtension ModuleDoc (PersistentArray ModuleDoc) ← registerSimplePersistentEnvExtension {
   addImportedFn := fun _ => {}
   addEntryFn    := fun s e => s.push e
-  exportEntriesFnEx? := some fun _ _ es => fun
-    | .exported => #[]
-    | _         => es.toArray
+  exportEntriesFnEx? := some fun _ _ es level =>
+    if level < .server then
+      #[]
+    else
+      es.toArray
 }
 
 def addMainModuleDoc (env : Environment) (doc : ModuleDoc) : Environment :=

src/Lean/Environment.lean

@@ -254,6 +254,11 @@ structure Environment where
   -/
   private extensions      : Array EnvExtensionState
   /--
+  Additional imported environment extension state for the interpreter. Access via
+  `getModuleIREntries`.
+  -/
+  private irBaseExts      : Array EnvExtensionState
+  /--
   Constant names to be saved in the field `extraConstNames` at `ModuleData`.
   It contains auxiliary declaration names created by the code generator which are not in `constants`.
   When importing modules, we want to insert them at `const2ModIdx`.
@@ -1470,6 +1475,7 @@ def mkEmptyEnvironment (trustLevel : UInt32 := 0) : IO Environment := do
       header          := { trustLevel }
       extraConstNames := {}
       extensions      := exts
+      irBaseExts      := exts
     }
     realizedImportedConsts? := none
   }
@@ -1583,10 +1589,18 @@ language server. Higher levels will return the data of the maximum imported leve
 -/
 def getModuleEntries {α β σ : Type} [Inhabited σ] (ext : PersistentEnvExtension α β σ)
     (env : Environment) (m : ModuleIdx) (level := OLeanLevel.exported) : Array α :=
-  let exts := if level = .exported then env.base.private.extensions else env.serverBaseExts
+  let exts := match level with
+    | .exported => env.base.private.extensions
+    | _         => env.serverBaseExts
   -- safety: as in `getStateUnsafe`
   unsafe (ext.toEnvExtension.getStateImpl exts).importedEntries[m]!
 
+/-- Retrieves additional IR extension state for the interpreter. -/
+def getModuleIREntries {α β σ : Type} [Inhabited σ] (ext : PersistentEnvExtension α β σ)
+    (env : Environment) (m : ModuleIdx) : Array α :=
+  -- safety: as in `getStateUnsafe`
+  unsafe (ext.toEnvExtension.getStateImpl env.base.private.irBaseExts).importedEntries[m]!
+
 def addEntry {α β σ : Type} (ext : PersistentEnvExtension α β σ) (env : Environment) (b : β) : Environment :=
   ext.toEnvExtension.modifyState env fun s =>
     let state   := ext.addEntryFn s.state b;
@@ -1755,6 +1769,18 @@ def mkModuleData (env : Environment) (level : OLeanLevel := .private) : IO Modul
     constNames, constants, entries
   }
 
+@[extern "lean_ir_export_entries"]
+private opaque exportIREntries (env : Environment) : Array (Name × Array EnvExtensionEntry)
+
+private def mkIRData (env : Environment) : IO ModuleData := do
+  -- TODO: should we use a more specific/efficient data format for IR?
+  return { env.header with
+    entries := exportIREntries env
+    constants := default
+    constNames := default
+    extraConstNames := default
+  }
+
 def writeModule (env : Environment) (fname : System.FilePath) : IO Unit := do
   if env.header.isModule then
     let mkPart (level : OLeanLevel) :=
@@ -1763,6 +1789,7 @@ def writeModule (env : Environment) (fname : System.FilePath) : IO Unit := do
       (← mkPart .exported),
       (← mkPart .server),
       (← mkPart .private)]
+    saveModuleData (fname.withExtension "ir") env.mainModule (← mkIRData env)
   else
     saveModuleData fname env.mainModule (← mkModuleData env)
 
@@ -1846,21 +1873,37 @@ private structure ImportedModule extends EffectiveImport where
   /-- All loaded incremental compacted regions. -/
   parts     : Array (ModuleData × CompactedRegion)
 
-/-- The main module data that will eventually be used to construct the kernel environment. -/
-private def ImportedModule.mainModule? (self : ImportedModule) : Option ModuleData := do
-  let (baseMod, _) ← self.parts[0]?
-  self.parts[if baseMod.isModule && self.importAll then 2 else 0]?.map (·.1)
-
 /-- The main module data that will eventually be used to construct the publicly accessible constants. -/
 private def ImportedModule.publicModule? (self : ImportedModule) : Option ModuleData := do
   let (baseMod, _) ← self.parts[0]?
   return baseMod
 
+private def ImportedModule.getData? (self : ImportedModule) (level : OLeanLevel) : Option ModuleData := do
+  -- Without the module system, we only have the exported level.
+  let level := if (← self.publicModule?).isModule then level else .exported
+  self.parts[level.toCtorIdx]?.map (·.1)
+
+/-- The main module data that will eventually be used to construct the kernel environment. -/
+private def ImportedModule.mainModule? (self : ImportedModule) : Option ModuleData :=
+  self.getData? (if self.importAll then .private else .exported)
+
 /-- The module data that should be used for server purposes. -/
 private def ImportedModule.serverData? (self : ImportedModule) (level : OLeanLevel) :
-    Option ModuleData := do
-  let (baseMod, _) ← self.parts[0]?
-  self.parts[if baseMod.isModule && level != .exported then 1 else 0]?.map (·.1)
+    Option ModuleData :=
+  -- fall back to `exported` outside the server
+  self.getData? (if level ≥ .server then level else .exported)
+
+/-- The module data that should be used for codegen purposes. -/
+private def ImportedModule.loadIRData? (self : ImportedModule) (arts : NameMap ModuleArtifacts)
+    (level : OLeanLevel):
+    IO (Option (ModuleData × CompactedRegion)) := do
+  if (level < .server && self.irPhases == .runtime) || !self.mainModule?.any (·.isModule) then
+    return none
+  let fname ←
+    arts.find? self.module |>.bind (·.ir?) |>.getDM do
+      let mFile ← findOLean self.module
+      return mFile.withExtension "ir"
+  readModuleData fname
 
 structure ImportState where
   private moduleNameMap : Std.HashMap Name ImportedModule := {}
@@ -1948,7 +1991,10 @@ where go (imports : Array Import) (importAll isExported isMeta : Bool) := do
     let importAll := !isModule || (importAll && i.importAll)
     -- `B ≥ public`?
     let isExported := isExported && i.isExported
-    let irPhases := if isMeta || i.isMeta then .comptime else .runtime
+    let irPhases :=
+      if !isModule then .all
+      else if isMeta || i.isMeta then .comptime
+      else .runtime
     let goRec imports := do
       go (importAll := importAll) (isExported := isExported) (isMeta := isMeta || i.isMeta) imports
     if let some mod := (← get).moduleNameMap[i.module]? then
@@ -2015,7 +2061,8 @@ Constructs environment from `importModulesCore` results.
 See also `importModules` for parameter documentation.
 -/
 def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (trustLevel : UInt32 := 0)
-    (leakEnv loadExts : Bool) (level := OLeanLevel.private) : IO Environment := do
+    (leakEnv loadExts : Bool) (level := OLeanLevel.private) (arts : NameMap ModuleArtifacts := {}) :
+    IO Environment := do
   let isModule := level != .private
   let modules := s.moduleNames.filterMap (s.moduleNameMap[·]?)
   let moduleData ← modules.mapM fun mod => do
@@ -2066,6 +2113,7 @@ def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (
     quotInit        := !imports.isEmpty -- We assume `Init.Prelude` initializes quotient module
     extraConstNames := {}
     extensions      := exts
+    irBaseExts      := exts
     header     := {
       trustLevel, imports, moduleData, isModule
       modules      := modules.map (·.toEffectiveImport)
@@ -2074,14 +2122,23 @@ def finalizeImport (s : ImportState) (imports : Array Import) (opts : Options) (
   }
   let publicConstants : ConstMap := SMap.fromHashMap publicConstantMap false
   let publicBase := { privateBase with constants := publicConstants, header.regions := #[] }
+  let extensions ← setImportedEntries privateBase.extensions moduleData
+  let irData? ← modules.mapM (·.loadIRData? arts level)
+  -- fall back to basic data when not in server
+  let serverData := modules.mapIdx (fun idx mod => mod.serverData? level |>.getD moduleData[idx]!)
+  let irData := irData?.mapIdx (fun idx ir => ir.map (·.1) |>.getD moduleData[idx]!)
+  let irRegions := irData?.filterMap (·.map (·.2))
+  let privateBase := { privateBase with
+    extensions
+    irBaseExts := (← setImportedEntries privateBase.extensions irData)
+    header.regions := privateBase.header.regions ++ irRegions
+  }
   let mut env : Environment := {
     base.private := privateBase
     base.public  := publicBase
     realizedImportedConsts? := none
+    serverBaseExts := (← setImportedEntries privateBase.extensions serverData)
   }
-  env := env.setCheckedSync { env.base.private with extensions := (← setImportedEntries env.base.private.extensions moduleData) }
-  let serverData := modules.filterMap (·.serverData? level)
-  env := { env with serverBaseExts := (← setImportedEntries env.base.private.extensions serverData) }
   if leakEnv then
     /- Mark persistent a first time before `finalizePersistenExtensions`, which
        avoids costly MT markings when e.g. an interpreter closure (which

src/Lean/Setup.lean

@@ -65,6 +65,7 @@ structure ModuleArtifacts where
   oleanServer? : Option System.FilePath := none
   oleanPrivate? : Option System.FilePath := none
   ilean? : Option System.FilePath := none
+  ir? : Option System.FilePath := none
   c? : Option System.FilePath := none
   bc? : Option System.FilePath := none
   deriving Repr, Inhabited, ToJson, FromJson

src/library/compiler/ir_interpreter.cpp

@@ -885,6 +885,7 @@ private:
             throw exception(sstream() << "cannot evaluate `[init]` declaration '" << fn << "' in the same module");
         }
         push_frame(e.m_decl, m_arg_stack.size());
+        lean_always_assert(decl_tag(e.m_decl) == decl_kind::Fun);
         value r = eval_body(decl_fun_body(e.m_decl));
         pop_frame(r, decl_type(e.m_decl));
         if (!type_is_scalar(t)) {

stage0/src/stdlib_flags.h

@@ -11,7 +11,7 @@ options get_default_options() {
     opts = opts.update({"debug", "terminalTacticsAsSorry"}, false);
     // switch to `true` for ABI-breaking changes affecting meta code;
     // see also next option!
-    opts = opts.update({"interpreter", "prefer_native"}, false);
+    opts = opts.update({"interpreter", "prefer_native"}, true);
     // switch to `false` when enabling `prefer_native` should also affect use
     // of built-in parsers in quotations; this is usually the case, but setting
     // both to `true` may be necessary for handling non-builtin parsers with

tests/bench/speedcenter.exec.velcom.yaml

@@ -40,6 +40,8 @@
       find ${BUILD:-build/release}/stage2/lib/lean -name '*.olean.server' -print0 | wc -c --files0-from=- | tail -1 | cut -d' ' -f 1
       echo -n 'bytes .olean.private: '
       find ${BUILD:-build/release}/stage2/lib/lean -name '*.olean.private' -print0 | wc -c --files0-from=- | tail -1 | cut -d' ' -f 1
+      echo -n 'bytes .ir: '
+      find ${BUILD:-build/release}/stage2/lib/lean -name '*.ir' -print0 | wc -c --files0-from=- | tail -1 | cut -d' ' -f 1
       echo -n 'lines C: '
       find ${BUILD:-build/release}/stage2/lib/temp -name '*.c' -print0 | wc -l --files0-from=- | tail -1 | cut -d' ' -f 1
       echo -n 'lines C++: '

tests/lean/run/structure.lean

@@ -24,7 +24,7 @@ inductive D
   | mk (x y z : Nat) : D
 
 /--
-info: #[constants, quotInit, diagnostics, const2ModIdx, extensions, extraConstNames, header]
+info: #[constants, quotInit, diagnostics, const2ModIdx, extensions, irBaseExts, extraConstNames, header]
 #[toS2, toS1, x, y, z, toS3, w, s]
 (some [S4.toS2, S2.toS1])
 #[S2, S3]