LexFilter's hwTokenFetch is too large for the JIT to optimize and zeroes a large stack frame for every token

[nojaf]

LexFilterImpl.hwTokenFetch in LexFilter.fs compiles to a single method of 110,356 bytes of IL with 565 locals. That is over the JIT's size limit for optimizing a method (60,000 bytes of IL, DEFAULT_MIN_OPTS_CODE_SIZE), so the JIT compiles it with MinOpts and never recompiles it. The method runs once per token, and each call sets up and zeroes a stack frame of 11 KB to 20 KB depending on how it was compiled.

Splitting the method in three, without changing any rule, made parsing 28% faster in Fantomas, which compiles the same source as FCS.

Repro steps

1. Parse a set of files with the compiler from the SDK, with ReadyToRun disabled so that the method is JIT-compiled, and ask the JIT for its disassembly:

   { echo --parseonly; echo --nologo; find src/Compiler -name '*.fs' | head -300; } > files.rsp
   DOTNET_ReadyToRun=0 DOTNET_JitDisasm=hwTokenFetch DOTNET_JitStdOutFile=disasm.txt \
       dotnet <sdk>/FSharp/fsc.dll @files.rsp

2. The top of disasm.txt (SDK 10.0.401, Linux x64):

   ; Assembly listing for method FSharp.Compiler.LexFilter+LexFilterImpl:hwTokenFetch(bool):FSharp.Compiler.Parser+token:this (Tier0-FullOpts-MinOpts)
   ; Tier-0 switched to FullOpts, then to MinOpts code
   ...
          lea      rbp, [rsp+0x4930]
   ...
   ; Total bytes of code 173570
   

   An 18 KB frame, zeroed in the prologue on every call, around 173 KB of unoptimized code.

Expected behavior

The offside rule is optimized by the JIT like the rest of the lexer and parser, with a frame small enough that zeroing it per token does not matter.

Actual behavior

Where the code is JIT-compiled, hwTokenFetch stays at MinOpts. That covers the FSharp.Compiler.Service NuGet package, which is not ReadyToRun (checked for 43.9.300), and so every editor and tool built on it. Every local gets its own stack slot, and the prologue zeroes the whole frame.

I measured this in Fantomas, which compiles the FCS syntax tree sources from commit 74ec4f7. That version of hwTokenFetch is identical to the one on main. With perf on Linux, parsing the 688 source files of this repository:

• hwTokenFetch alone was 28% of parse time.
• The loop that zeroes its 20 KB frame was the hottest code in the parse: 14% to 20% of all samples, over two runs.

The SDK's FSharp.Compiler.Service.dll is ReadyToRun. Its precompiled hwTokenFetch is optimized, but still has an 11.7 KB frame, of which it zeroes 0x2a90 bytes (10.9 KB) per call:

lea    -0x2db8(%rsp),%r11
call   *0x12f1970(%rip)          ; stack probe
...
movabs $0xffffffffffffd570,%rax
movaps %xmm8,-0x70(%rbp,%rax,1)
movaps %xmm8,-0x60(%rbp,%rax,1)
movaps %xmm8,-0x50(%rbp,%rax,1)
add    $0x30,%rax
jne    ...

In fsc --parseonly over the files above, that loop was 6.9% of the samples taken while parsing. A dotnet build pays this as well, just less than the JIT-compiled package does.

Proposed fix

Split the rule chain of the big match token, offsideStack with into consecutive parts, each a local function that the previous part falls through to:

let rec hwTokenFetch useBlockRule =
    let tokenTup = popNextTokenTup()
    // ... the same locals and helpers as today ...

    let hwTokenFetchPart3 () =
        match token, offsideStack with
        | LARROW, _ when isControlFlowOrNotSameLine() -> ...   // the rules from here to the end, unchanged
        ...
        | _ -> returnToken tokenLexbufState token

    let hwTokenFetchPart2 () =
        match token, offsideStack with
        | _, CtxtInterfaceHead offsidePos :: _ when ... -> ... // the middle rules, unchanged
        ...
        | _ -> hwTokenFetchPart3 ()

    match token, offsideStack with
    | _ when tokensThatNeedNoProcessingCount > 0 -> ...         // the first rules, unchanged
    ...
    | _ -> hwTokenFetchPart2 ()

A match takes the first rule that matches, and a token that matches nothing in one part reaches the next part with the same token and offsideStack, as no rule has run yet. So every token is handled by the same rule as before. The split points are between whole rules, never inside an or-pattern: after the rule for a token offside from CtxtDo, and before the LARROW rule. Apart from the two new let lines and the two | _ -> lines, the diff only indents two blocks of rules by four spaces, so it is small with whitespace ignored.

The F# compiler lifts both local functions to static methods that take the captured values as arguments, so there is no closure allocated per token. In Fantomas's copy:

	IL	Frame	JIT
hwTokenFetch before	110 KB	20 KB	MinOpts
hwTokenFetch after	20.6 KB	904 bytes	optimized
hwTokenFetchPart2	37.4 KB	392 bytes	optimized
hwTokenFetchPart3	9.6 KB	504 bytes	optimized

Parsing the 688 source files of this repository that Fantomas formats, without defines, averaged over warm runs:

	Time	Allocated
Before	about 1.59 s	732 MB
After	about 1.14 s	731 MB

Fantomas's output for 902 files in 4 configurations (3608 outputs) is byte for byte the same with and without the split, and its test suites pass. I have not yet built the compiler itself with the change or run its test suite.

The parts can grow past the limit again as rules are added, so a comment at the split saying why it is there would help keep it.

Known workarounds

None from outside the compiler. Setting the JIT's MinOpts thresholds through DOTNET_JITMinOpts* has no effect on a release runtime.

Related information

• Linux x64, .NET runtime 10.0.11, SDK 10.0.401
• FSharp.Compiler.Service 43.9.300 from NuGet (not ReadyToRun), and the SDK's copy (ReadyToRun)
• Found while profiling Fantomas (Allocate less while formatting fsprojects/fantomas#3499)

[auduchinok]

@nojaf Nice analysis! Do you think it's possible to split it in more logical parts instead of just splitting it like this?

[T-Gro]

Could we not "anti inline" the individual case handlers ?
Either as a new form of compiler optimization, or just directly in the source code?

So that the method would become mere router - just following the pattern match and for each case immediately calling a different .NET IL method.

And if the current IL does already contain parts that were inlined by the optimized but are separate functions in the source, [<NoCompilerInlining>] could be used as well?

[T-Gro]

And we might also turn the if debug into a [<Conditional>] ?

[nojaf]

Thanks both. Before choosing an approach I measured the IL of hwTokenFetch in a Release netstandard2.0 build of current main (110,341 bytes, 509 locals). The size comes from the decision tree duplicating guards, not from large rule bodies.

Where the 110 KB comes from

A when guard can fail, and when it does the match compiler compiles the remaining rules again on that path. Each guard that fails early therefore carries a copy of the tree below it. Here is how often some calls appear in the IL compared with the source:

Call	Source	IL
thereIsACtxtMemberBodyOnTheStackAndWeShouldPopStackForUpcomingMember (the one MEMBER guard)	1 use	168
isSameLine	4	27
isTryBlockContinuator	1	14
The CtxtMatchClauses guard, errorR included	1	14
HeadOrDefault (tests on offsideStack)		736
dprintf (these sit in rule bodies)	92	84

Rule bodies are shared targets and are emitted once. The copies are guards and stack tests. That also explains why the three-way split cut the total from 110 KB to 67.6 KB: each | _ -> next () stops the copying at the boundary.

On the suggestions

@T-Gro, moving the case handlers into separate methods would shrink the stack frame, because their locals go with them. It would not remove much IL, because the bodies are not duplicated. What is left would be a router of about 60 KB or more of copied tests. Your idea does pay off for the large guards: hoisted into named local functions, each copy becomes a single call. The CtxtMatchClauses guard is the clearest case, an inline lambda with an errorR in it, emitted 14 times.

[<NoCompilerInlining>] would not change much here. The helpers already show up as calls (isSameLine, isInfix, the is*Continuator functions), so they are not being inlined into the method.

On [<Conditional>] for if debug, I don't think it fits, for two reasons:

• debug is not a compile-time constant. Its value comes from the hidden --tokenize-debug option and from fsi, and it prints to stdout in release builds today. Conditional would make that option do nothing in shipped compilers. That may be acceptable, but it should be a deliberate decision.
• dprintf returns a curried printf function, and Conditional needs a method that returns unit, so all 92 call sites would need rewriting.

The debug code is in the cold bodies, so it is a small win at best. I would keep it out of this change.

@auduchinok, agreed. The only constraint is that the parts have to stay contiguous and in the original order, since the rules overlap. The rules already fall into three sections:

1. Preprocessing, tokens that close a context (IN, DONE, closing parentheses, END), and the namespace and module heads.
2. The offside rules, which pop a context (_, CtxtX offsidePos :: _ when ... <= offsidePos.Column), starting at the CtxtSeqBlock offside rule.
3. Keywords that push a context, from NAMESPACE to the end.

The split points in the issue (the CtxtDo offside rule and LARROW) cut through sections 2 and 3. Splitting at the section boundaries instead gives parts that can be named for what they do, for example hwTokenFetchClosing, hwTokenFetchOffside and hwTokenFetchPush. The risk is section 2: it holds nearly all the guards and so most of the copying. It needs measuring, and it may need the guard hoisting to stay well under 60 KB.

Proposed approach

1. Split at the section boundaries above, with names and a comment explaining why the split exists. Measure each part's IL.
2. If the offside part is too large, hoist its biggest guards into local functions, in the same PR or the next one.
3. Add a test that fails if any hwTokenFetch* method goes over a limit, such as 50 KB, by reading the method body size through System.Reflection.Metadata. That keeps rules added later from pushing a part past the JIT limit again, which a comment alone would not.
4. Leave Conditional for a separate discussion.

A general fix in the match compiler, jumping to a shared continuation when a guard fails instead of re-expanding the remaining rules, would help all F# code with large guarded matches. That would be the "new form of compiler optimization", but it is a separate and much larger project.

Script used for the measurements

Run as dotnet fsi probe.fsx <path to FSharp.Compiler.Service.dll> "LexFilterImpl::hwTokenFetch". It prints the IL size, local count and max stack of each matching method, followed by its most frequent call targets.

open System
open System.IO
open System.Reflection
open System.Reflection.Emit
open System.Reflection.Metadata
open System.Reflection.PortableExecutable
open System.Reflection.Metadata.Ecma335
open System.Collections.Generic

let path = fsi.CommandLineArgs[1]
let filter = fsi.CommandLineArgs[2]
let pe = new PEReader(File.OpenRead path)
let md = pe.GetMetadataReader()

let one = Dictionary<int16, OpCode>()
let two = Dictionary<int16, OpCode>()
for f in typeof<OpCodes>.GetFields(BindingFlags.Public ||| BindingFlags.Static) do
    let op = f.GetValue(null) :?> OpCode
    if op.Size = 1 then one[op.Value] <- op else two[int16 (op.Value &&& 0xffs)] <- op

let operandSize (op: OpCode) (il: byte[]) i =
    match op.OperandType with
    | OperandType.InlineNone -> 0
    | OperandType.ShortInlineBrTarget | OperandType.ShortInlineI | OperandType.ShortInlineVar -> 1
    | OperandType.InlineVar -> 2
    | OperandType.InlineI8 | OperandType.InlineR -> 8
    | OperandType.InlineSwitch -> 4 + 4 * BitConverter.ToInt32(il, i)
    | _ -> 4

let memberName (h: EntityHandle) =
    match h.Kind with
    | HandleKind.MethodDefinition ->
        let m = md.GetMethodDefinition(MethodDefinitionHandle.op_Explicit h)
        let t = md.GetTypeDefinition(m.GetDeclaringType())
        md.GetString t.Name + "::" + md.GetString m.Name
    | HandleKind.MemberReference ->
        let m = md.GetMemberReference(MemberReferenceHandle.op_Explicit h)
        let parent =
            match m.Parent.Kind with
            | HandleKind.TypeReference -> md.GetString (md.GetTypeReference(TypeReferenceHandle.op_Explicit m.Parent).Name)
            | HandleKind.TypeDefinition -> md.GetString (md.GetTypeDefinition(TypeDefinitionHandle.op_Explicit m.Parent).Name)
            | _ -> "?"
        parent + "::" + md.GetString m.Name
    | HandleKind.MethodSpecification ->
        let s = md.GetMethodSpecification(MethodSpecificationHandle.op_Explicit h)
        match s.Method.Kind with
        | HandleKind.MethodDefinition ->
            let m = md.GetMethodDefinition(MethodDefinitionHandle.op_Explicit s.Method) in md.GetString m.Name + "<>"
        | _ ->
            let m = md.GetMemberReference(MemberReferenceHandle.op_Explicit s.Method) in md.GetString m.Name + "<>"
    | _ -> "?"

for th in md.TypeDefinitions do
    let t = md.GetTypeDefinition th
    for mh in t.GetMethods() do
        let m = md.GetMethodDefinition mh
        let name = md.GetString t.Name + "::" + md.GetString m.Name
        if name.Contains(filter, StringComparison.Ordinal) && m.RelativeVirtualAddress <> 0 then
            let body = pe.GetMethodBody m.RelativeVirtualAddress
            let il = body.GetILBytes()
            let locals =
                if body.LocalSignature.IsNil then 0
                else
                    let sa = md.GetStandaloneSignature body.LocalSignature
                    let mutable r = md.GetBlobReader sa.Signature
                    r.ReadSignatureHeader() |> ignore
                    r.ReadCompressedInteger()
            printfn "%s  IL=%d locals=%d maxstack=%d" name il.Length locals body.MaxStack
            let calls = Dictionary<string, int>()
            let mutable i = 0
            while i < il.Length do
                let op =
                    if il[i] = 0xFEuy then (i <- i + 1; two[int16 il[i]]) else one[int16 il[i]]
                i <- i + 1
                if op.OperandType = OperandType.InlineMethod then
                    let tok = BitConverter.ToInt32(il, i)
                    let h = MetadataTokens.EntityHandle tok
                    let n = op.Name + " " + memberName h
                    calls[n] <- (match calls.TryGetValue n with | true, c -> c + 1 | _ -> 1)
                i <- i + operandSize op il i
            for KeyValue(k, v) in calls |> Seq.sortByDescending (fun kv -> kv.Value) |> Seq.truncate 40 do
                printfn "   %5d  %s" v k