chore: fixes from #13103 "enable separate codegen" (#13241)

.vscode/tasks.json

@@ -11,6 +11,15 @@
 				"isDefault": true
 			}
 		},
+		{
+			"label": "build stage2",
+			"type": "shell",
+			"command": "make -C build/release stage2 -j$(nproc 2>/dev/null || sysctl -n hw.logicalcpu 2>/dev/null || echo 4)",
+			"problemMatcher": [],
+			"group": {
+				"kind": "build"
+			}
+		},
 		{
 			"label": "build-old",
 			"type": "shell",

src/Lean/Compiler/LCNF/Basic.lean

@@ -1230,7 +1230,14 @@ def instantiateRevRangeArgs (e : Expr) (beginIdx endIdx : Nat) (args : Array (Ar
   else
     e.instantiateRevRange beginIdx endIdx (args.map (·.toExpr))
 
-/-- Lookup function for compiler extensions with sorted persisted state that works in both `lean` and `leanir`. -/
+/--
+Lookup function for compiler extensions with sorted persisted state that works in both `lean` and
+`leanir`.
+
+`preferImported` defaults to false because in `leanir`, we do not want to mix information from
+`meta` compilation in `lean` with our own state. But in `lean`, setting `preferImported` can help
+with avoiding unnecessary task blocks.
+-/
 @[inline] def findExtEntry? [Inhabited σ] (env : Environment) (ext : PersistentEnvExtension α β σ) (declName : Name)
     (findAtSorted? : Array α → Name → Option α')
     (findInState? : σ → Name → Option α') : Option α' :=

src/Lean/Compiler/LCNF/Main.lean

@@ -78,9 +78,13 @@ def isValidMainType (type : Expr) : Bool :=
     isValidResultName resultName
   | _ => false
 
+/-- A postponed call of `compileDecls`. -/
 structure PostponedCompileDecls where
+  /-- Declaration names of this mutual group. -/
   declNames : Array Name
-deriving BEq, Hashable
+  /-- Options at time of original call, to be restored for tracing etc. -/
+  options : Options
+deriving BEq
 
 /--
 Saves postponed `compileDecls` calls.
@@ -101,16 +105,20 @@ builtin_initialize postponedCompileDeclsExt : SimplePersistentEnvExtension Postp
       { exported := #[], server := #[], «private» := es.toArray }
   }
 
-def resumeCompilation (declName : Name) : CoreM Unit := do
+def resumeCompilation (declName : Name) (baseOpts : Options) : CoreM Unit := do
   let some decls := postponedCompileDeclsExt.getState (← getEnv) |>.find? declName | return
+  let opts := baseOpts.mergeBy (fun _ base _ => base) decls.options
+  let opts := compiler.postponeCompile.set opts false
   modifyEnv (postponedCompileDeclsExt.modifyState · fun s => decls.declNames.foldl (·.erase) s)
-  withOptions (compiler.postponeCompile.set · false) do
+  -- NOTE: we *must* throw away the current options as they could depend on the specific recursion
+  -- we did to get here.
+  withOptions (fun _ => opts) do
   Core.prependError m!"Failed to compile `{declName}`" do
-    (← compileDeclsRef.get) decls.declNames
+    (← compileDeclsRef.get) decls.declNames baseOpts
 
 namespace PassManager
 
-partial def run (declNames : Array Name) : CompilerM Unit := withAtLeastMaxRecDepth 8192 do
+partial def run (declNames : Array Name) (baseOpts : Options) : CompilerM Unit := withAtLeastMaxRecDepth 8192 do
   /-
   Note: we need to increase the recursion depth because we currently do to save phase1
   declarations in .olean files. Then, we have to recursively compile all dependencies,
@@ -141,11 +149,14 @@ partial def run (declNames : Array Name) : CompilerM Unit := withAtLeastMaxRecDe
 
   -- Now that we have done all input checks, check for postponement
   if (← getEnv).header.isModule && (← compiler.postponeCompile.getM) then
-    modifyEnv (postponedCompileDeclsExt.addEntry · { declNames := decls.map (·.name) })
+    modifyEnv (postponedCompileDeclsExt.addEntry · { declNames := decls.map (·.name), options := ← getOptions })
     -- meta defs are compiled locally so they are available for execution/compilation without
     -- importing `.ir` but still marked for `leanir` compilation so that we do not have to persist
     -- module-local compilation information between the two processes
-    if !decls.any (isMarkedMeta (← getEnv) ·.name) then
+    if decls.any (isMarkedMeta (← getEnv) ·.name) then
+      -- avoid re-compiling the meta defs in this process; the entry for `leanir` is not affected
+      modifyEnv (postponedCompileDeclsExt.modifyState · fun s => decls.foldl (·.erase ·.name) s)
+    else
       trace[Compiler] "postponing compilation of {decls.map (·.name)}"
       return
 
@@ -157,7 +168,7 @@ partial def run (declNames : Array Name) : CompilerM Unit := withAtLeastMaxRecDe
       let .let { value := .const c .., .. } .. := c | return
       -- Need to do some lookups to get the actual name passed to `compileDecls`
       let c := Compiler.getImplementedBy? (← getEnv) c |>.getD c
-      resumeCompilation c
+      resumeCompilation c baseOpts
 
   let decls := markRecDecls decls
   let manager ← getPassManager
@@ -200,9 +211,9 @@ where
 
 end PassManager
 
-def main (declNames : Array Name) : CoreM Unit := do
+def main (declNames : Array Name) (baseOpts : Options) : CoreM Unit := do
   withTraceNode `Compiler (fun _ => return m!"compiling: {declNames}") do
-    CompilerM.run <| PassManager.run declNames
+    CompilerM.run <| PassManager.run declNames baseOpts
 
 builtin_initialize
   compileDeclsRef.set main

src/Lean/Compiler/Main.lean

@@ -19,7 +19,7 @@ that fulfill the requirements of `shouldGenerateCode`.
 def compile (declNames : Array Name) : CoreM Unit := do profileitM Exception "compiler new" (← getOptions) do
   withOptions (compiler.postponeCompile.set · false) do
   withTraceNode `Compiler (fun _ => return m!"compiling: {declNames}") do
-    LCNF.main declNames
+    LCNF.main declNames {}
 
 builtin_initialize
   registerTraceClass `Compiler

src/Lean/CoreM.lean

@@ -711,11 +711,11 @@ breaks the cycle by making `compileDeclsImpl` a "dynamic" call through the ref t
 to the linker. In the compiler there is a matching `builtin_initialize` to set this ref to the
 actual implementation of compileDeclsRef.
 -/
-builtin_initialize compileDeclsRef : IO.Ref (Array Name → CoreM Unit) ←
+builtin_initialize compileDeclsRef : IO.Ref (Array Name → Options → CoreM Unit) ←
-  IO.mkRef (fun _ => throwError m!"call to compileDecls with uninitialized compileDeclsRef")
+  IO.mkRef (fun _ _ => throwError m!"call to compileDecls with uninitialized compileDeclsRef")
 
 private def compileDeclsImpl (declNames : Array Name) : CoreM Unit := do
-  (← compileDeclsRef.get) declNames
+  (← compileDeclsRef.get) declNames {}
 
 -- `ref?` is used for error reporting if available
 def compileDecls (decls : Array Name) (logErrors := true) : CoreM Unit := do

src/Lean/Shell.lean

@@ -12,6 +12,7 @@ import Lean.Server.Watchdog
 import Lean.Server.FileWorker
 import Lean.Compiler.LCNF.EmitC
 import Init.System.Platform
+import Lean.Compiler.Options
 
 /-  Lean companion to  `shell.cpp` -/
 
@@ -340,7 +341,10 @@ def ShellOptions.process (opts : ShellOptions)
   | 'I' => -- `-I, --stdin`
     return {opts with useStdin := true}
   | 'r' => -- `--run`
-    return {opts with run := true}
+    return {opts with
+      run := true
+      -- can't get IR if it's postponed
+      leanOpts := Compiler.compiler.postponeCompile.set opts.leanOpts false }
   | 'o' => -- `--o, olean=fname`
     return {opts with oleanFileName? := ← checkOptArg "o" optArg?}
   | 'i' => -- `--i, ilean=fname`

src/LeanIR.lean

@@ -57,14 +57,18 @@ def setConfigOption (opts : Options) (arg : String) : IO Options := do
 
 public def main (args : List String) : IO UInt32 := do
   let setupFile::irFile::c::optArgs := args | do
-    IO.println s!"usage: leanir <setup.json> <module> <output.ir> <output.c> <-Dopt=val>..."
+    IO.println s!"usage: leanir <setup.json> <output.ir> <output.c> [--stat] <-Dopt=val>..."
     return 1
 
   let setup ← ModuleSetup.load setupFile
   let modName := setup.name
 
+  let mut printStats := false
   let mut opts := setup.options.toOptions
   for optArg in optArgs do
+    if optArg == "--stat" then
+      printStats := true
+    else
       opts ← setConfigOption opts optArg
   opts := Compiler.compiler.inLeanIR.set opts true
   opts := maxHeartbeats.set opts 0
@@ -127,12 +131,15 @@ public def main (args : List String) : IO UInt32 := do
     modifyEnv (postponedCompileDeclsExt.setState · (decls.foldl (fun s e => e.declNames.foldl (·.insert · e) s) {}))
     for decl in decls do
       for decl in decl.declNames do
-        resumeCompilation decl
+        try
+          resumeCompilation decl (← getOptions)
+        finally
+          addTraceAsMessages
+      for msg in (← Core.getAndEmptyMessageLog).unreported do
+        IO.eprintln (← msg.toString)
   catch e =>
     unless e.isInterrupt do
       logError e.toMessageData
-  finally
-    addTraceAsMessages
 
   let .ok (_, s) := res? | unreachable!
   let env := s.env
@@ -155,4 +162,6 @@ public def main (args : List String) : IO UInt32 := do
     out.write data.toUTF8
 
   displayCumulativeProfilingTimes
+  if printStats then
+    env.displayStats
   return 0

tests/bench/build/fake_root/bin/leanir

@@ -62,7 +62,7 @@ def run_leanir(module: str) -> None:
         cmd=[
             "leanir",
             *sys.argv[1:],
-            *("-Dprofiler=true", "-Dprofiler.threshold=9999999"),
+            "--stat", "-Dprofiler=true", "-Dprofiler.threshold=9999999",
         ],
         output=WRAPPER_OUT,
         topics=[f"{BENCHMARK}/module/ir/{module}"],
@@ -81,6 +81,19 @@ def run_leanir(module: str) -> None:
                 seconds = seconds / 1000
             save_measurement(f"{BENCHMARK}/profile/ir/{name}//wall-clock", seconds, "s")
 
+    # Output of `leanir --stat`
+    stat = Counter[str]()
+    for line in stdout.splitlines():
+        if match := re.fullmatch(r"number of (imported .*):\s+(\d+)$", line):
+            name = match.group(1)
+            count = int(match.group(2))
+            stat[name] += count
+    for name, count in stat.items():
+        if count > 0:
+            if name.endswith("bytes"):
+                save_measurement(f"{BENCHMARK}/stat/ir/{name}//bytes", count, "B")
+            else:
+                save_measurement(f"{BENCHMARK}/stat/ir/{name}//amount", count)
 
 def main() -> None:
     parser = argparse.ArgumentParser()

tests/compile/run_bench.sh

@@ -12,7 +12,7 @@ rm -f "$1.measurements.jsonl"
 if [[ -n $DO_COMPILE ]]; then
   run_before "$1"
 
-  lean --c="$1.c" "${TEST_LEAN_ARGS[@]}" "$1" || fail "Failed to compile $1 into $1.c"
+  lean --c="$1.c" -Dcompiler.postponeCompile=false "${TEST_LEAN_ARGS[@]}" "$1" || fail "Failed to compile $1 into $1.c"
   leanc ${LEANC_OPTS-} -O3 -DNDEBUG -o "$1.out" "${TEST_LEANC_ARGS[@]}" "$1.c" || fail "Failed to compile $1.c"
 
   # Measure .out binary size

tests/compile/run_test.sh

@@ -18,7 +18,7 @@ if [[ -n $DO_COMPILE ]]; then
   echo "Compiling and executing lean file"
   run_before "$1"
 
-  lean --c="$1.c" "${TEST_LEAN_ARGS[@]}" "$1" || fail "Failed to compile $1 into $1.c"
+  lean --c="$1.c" -Dcompiler.postponeCompile=false "${TEST_LEAN_ARGS[@]}" "$1" || fail "Failed to compile $1 into $1.c"
   leanc ${LEANC_OPTS-} -O3 -DNDEBUG -o "$1.out" "${TEST_LEANC_ARGS[@]}" "$1.c" || fail "Failed to compile $1.c"
 
   capture_only "$1" \

tests/compile_bench/run_bench.sh

@@ -22,7 +22,7 @@ if [[ -n $DO_COMPILE ]]; then
 
   TOPIC="compiled/$(basename "$1" .lean)"
 
-  lean --c="$1.c" "${TEST_LEAN_ARGS[@]}" "$1" || fail "Failed to compile $1 into $1.c"
+  lean --c="$1.c" -Dcompiler.postponeCompile=false "${TEST_LEAN_ARGS[@]}" "$1" || fail "Failed to compile $1 into $1.c"
   leanc ${LEANC_OPTS-} -O3 -DNDEBUG -o "$1.out" "${TEST_LEANC_ARGS[@]}" "$1.c" || fail "Failed to compile $1.c"
 
   capture_only "$1" \

tests/elab/run_test.sh

@@ -3,8 +3,9 @@ run_before "$1"
 
 # `--root` to infer same private names as in the server
 # Elab.inServer to allow for arbitrary `#eval`
+# compiler.postponeCompile for immediate trace output
 capture_only "$1" \
-  lean --root=.. -DprintMessageEndPos=true -Dlinter.all=false -DElab.inServer=true "${TEST_LEAN_ARGS[@]}" "$1"
+  lean --root=.. -DprintMessageEndPos=true -Dlinter.all=false -DElab.inServer=true -Dcompiler.postponeCompile=false "${TEST_LEAN_ARGS[@]}" "$1"
 normalize_mvar_suffixes
 normalize_reference_urls
 normalize_measurements

tests/misc_dir/plugin/run_test.sh

@@ -1,5 +1,5 @@
 # LEAN_EXPORTING needs to be defined for .c files included in shared libraries
-lean --c=SnakeLinter.c SnakeLinter.lean
+lean --c=SnakeLinter.c -Dcompiler.postponeCompile=false SnakeLinter.lean
 leanc ${LEANC_OPTS-} -O3 -DNDEBUG -DLEAN_EXPORTING -shared -o SnakeLinter.so SnakeLinter.c
 
 capture_only SnakeLinter.lean \