Closure and closure context field loads are not folded in AOT

[osa1]

In https://github.com/osa1/protobuf.dart/tree/inlining_issue I add a bunch of inline pragmas to two higher-order functions and a few small functions.

Inlining these higher-order functions should eliminate indirect calls when calling the function arguments as in each call site the closure being passed is a function expression.

However currently AOT does not seem to fold field loads from closures and contexts.

If I manually inline _readPacked (by copy-pasting it into each of the call sites), that improves runtime of from_binary benchmark 7%. The diff for this change is at the end.

To repro:

• Clone the repo and switch to the inlining_issue branch.
• In benchmarks directory:
  • ./tool/compile_protos.sh (requires protoc with Dart plugin)
  • dart pub get
  • dart compile exe bin/from_binary.dart

I'm not fluent in VM's flow graph IR, but with --extra-gen-snapshot-options=--print-flow-graph-filter=_mergeFromCodedBufferReader,--print-flow-graph-optimized,--compiler-passes=\*Inlining and inline pragmas, I see code like

v525 <- Constant(#Function '<anonymous closure>': static.) T{Function}
...
v519 <- AllocateContext:10(num_variables=2) T{Context}
StoreField(v519 . input = v84 T{CodedBufferReader})
StoreField(v519 . readFunc = v82 T{_Closure})
...
v524 <- AllocateClosure:14(v525 T{Function}, v519) T{_Closure}
...
v4632 <- LoadField(v524 T{_Closure} . Closure.context {final}) T{!null}
...
v4636 <- LoadField(v4632 . readFunc {final}) T{(dynamic) => void}
v4637 <- ClosureCall:18( closure=v4636<0>, v4636)

(this is in an "after inlining" block)

Here LoadField(v524, ...) could be folded, which would then allow folding LoadField(v4632, ...), which I suspect in the end would give us the optimizations.

Inlining _readPacked manually is not ideal as that would potentially increase binary sizes on the browsers without improving runtime perf. Ideally I want to do this with backend-specific pragmas.

Diff for manually inlining `_readPacked`

diff --git a/protobuf/lib/src/protobuf/coded_buffer.dart b/protobuf/lib/src/protobuf/coded_buffer.dart
index a102cbb..ca5d26d 100644
--- a/protobuf/lib/src/protobuf/coded_buffer.dart
+++ b/protobuf/lib/src/protobuf/coded_buffer.dart
@@ -130,7 +130,11 @@ void _mergeFromCodedBufferReader(BuilderInfo meta, _FieldSet fs,
       case PbFieldType._REPEATED_BOOL:
         final list = fs._ensureRepeatedField(meta, fi);
         if (wireType == WIRETYPE_LENGTH_DELIMITED) {
-          _readPacked(input, () => list.add(input.readBool()));
+          input._withLimit(input.readInt32(), () {
+            while (!input.isAtEnd()) {
+              list.add(input.readBool());
+            }
+          });
         } else {
           list.add(input.readBool());
         }
@@ -146,7 +150,11 @@ void _mergeFromCodedBufferReader(BuilderInfo meta, _FieldSet fs,
       case PbFieldType._REPEATED_FLOAT:
         final list = fs._ensureRepeatedField(meta, fi);
         if (wireType == WIRETYPE_LENGTH_DELIMITED) {
-          _readPacked(input, () => list.add(input.readFloat()));
+          input._withLimit(input.readInt32(), () {
+            while (!input.isAtEnd()) {
+              list.add(input.readFloat());
+            }
+          });
         } else {
           list.add(input.readFloat());
         }
@@ -154,7 +162,11 @@ void _mergeFromCodedBufferReader(BuilderInfo meta, _FieldSet fs,
       case PbFieldType._REPEATED_DOUBLE:
         final list = fs._ensureRepeatedField(meta, fi);
         if (wireType == WIRETYPE_LENGTH_DELIMITED) {
-          _readPacked(input, () => list.add(input.readDouble()));
+          input._withLimit(input.readInt32(), () {
+            while (!input.isAtEnd()) {
+              list.add(input.readDouble());
+            }
+          });
         } else {
           list.add(input.readDouble());
         }
@@ -173,7 +185,11 @@ void _mergeFromCodedBufferReader(BuilderInfo meta, _FieldSet fs,
       case PbFieldType._REPEATED_INT32:
         final list = fs._ensureRepeatedField(meta, fi);
         if (wireType == WIRETYPE_LENGTH_DELIMITED) {
-          _readPacked(input, () => list.add(input.readInt32()));
+          input._withLimit(input.readInt32(), () {
+            while (!input.isAtEnd()) {
+              list.add(input.readInt32());
+            }
+          });
         } else {
           list.add(input.readInt32());
         }
@@ -181,7 +197,11 @@ void _mergeFromCodedBufferReader(BuilderInfo meta, _FieldSet fs,
       case PbFieldType._REPEATED_INT64:
         final list = fs._ensureRepeatedField(meta, fi);
         if (wireType == WIRETYPE_LENGTH_DELIMITED) {
-          _readPacked(input, () => list.add(input.readInt64()));
+          input._withLimit(input.readInt32(), () {
+            while (!input.isAtEnd()) {
+              list.add(input.readInt64());
+            }
+          });
         } else {
           list.add(input.readInt64());
         }
@@ -189,7 +209,11 @@ void _mergeFromCodedBufferReader(BuilderInfo meta, _FieldSet fs,
       case PbFieldType._REPEATED_SINT32:
         final list = fs._ensureRepeatedField(meta, fi);
         if (wireType == WIRETYPE_LENGTH_DELIMITED) {
-          _readPacked(input, () => list.add(input.readSint32()));
+          input._withLimit(input.readInt32(), () {
+            while (!input.isAtEnd()) {
+              list.add(input.readSint32());
+            }
+          });
         } else {
           list.add(input.readSint32());
         }
@@ -197,7 +221,11 @@ void _mergeFromCodedBufferReader(BuilderInfo meta, _FieldSet fs,
       case PbFieldType._REPEATED_SINT64:
         final list = fs._ensureRepeatedField(meta, fi);
         if (wireType == WIRETYPE_LENGTH_DELIMITED) {
-          _readPacked(input, () => list.add(input.readSint64()));
+          input._withLimit(input.readInt32(), () {
+            while (!input.isAtEnd()) {
+              list.add(input.readSint64());
+            }
+          });
         } else {
           list.add(input.readSint64());
         }
@@ -205,7 +233,11 @@ void _mergeFromCodedBufferReader(BuilderInfo meta, _FieldSet fs,
       case PbFieldType._REPEATED_UINT32:
         final list = fs._ensureRepeatedField(meta, fi);
         if (wireType == WIRETYPE_LENGTH_DELIMITED) {
-          _readPacked(input, () => list.add(input.readUint32()));
+          input._withLimit(input.readInt32(), () {
+            while (!input.isAtEnd()) {
+              list.add(input.readUint32());
+            }
+          });
         } else {
           list.add(input.readUint32());
         }
@@ -213,7 +245,11 @@ void _mergeFromCodedBufferReader(BuilderInfo meta, _FieldSet fs,
       case PbFieldType._REPEATED_UINT64:
         final list = fs._ensureRepeatedField(meta, fi);
         if (wireType == WIRETYPE_LENGTH_DELIMITED) {
-          _readPacked(input, () => list.add(input.readUint64()));
+          input._withLimit(input.readInt32(), () {
+            while (!input.isAtEnd()) {
+              list.add(input.readUint64());
+            }
+          });
         } else {
           list.add(input.readUint64());
         }
@@ -221,7 +257,11 @@ void _mergeFromCodedBufferReader(BuilderInfo meta, _FieldSet fs,
       case PbFieldType._REPEATED_FIXED32:
         final list = fs._ensureRepeatedField(meta, fi);
         if (wireType == WIRETYPE_LENGTH_DELIMITED) {
-          _readPacked(input, () => list.add(input.readFixed32()));
+          input._withLimit(input.readInt32(), () {
+            while (!input.isAtEnd()) {
+              list.add(input.readFixed32());
+            }
+          });
         } else {
           list.add(input.readFixed32());
         }
@@ -229,7 +269,11 @@ void _mergeFromCodedBufferReader(BuilderInfo meta, _FieldSet fs,
       case PbFieldType._REPEATED_FIXED64:
         final list = fs._ensureRepeatedField(meta, fi);
         if (wireType == WIRETYPE_LENGTH_DELIMITED) {
-          _readPacked(input, () => list.add(input.readFixed64()));
+          input._withLimit(input.readInt32(), () {
+            while (!input.isAtEnd()) {
+              list.add(input.readFixed64());
+            }
+          });
         } else {
           list.add(input.readFixed64());
         }
@@ -237,7 +281,11 @@ void _mergeFromCodedBufferReader(BuilderInfo meta, _FieldSet fs,
       case PbFieldType._REPEATED_SFIXED32:
         final list = fs._ensureRepeatedField(meta, fi);
         if (wireType == WIRETYPE_LENGTH_DELIMITED) {
-          _readPacked(input, () => list.add(input.readSfixed32()));
+          input._withLimit(input.readInt32(), () {
+            while (!input.isAtEnd()) {
+              list.add(input.readSfixed32());
+            }
+          });
         } else {
           list.add(input.readSfixed32());
         }
@@ -245,7 +293,11 @@ void _mergeFromCodedBufferReader(BuilderInfo meta, _FieldSet fs,
       case PbFieldType._REPEATED_SFIXED64:
         final list = fs._ensureRepeatedField(meta, fi);
         if (wireType == WIRETYPE_LENGTH_DELIMITED) {
-          _readPacked(input, () => list.add(input.readSfixed64()));
+          input._withLimit(input.readInt32(), () {
+            while (!input.isAtEnd()) {
+              list.add(input.readSfixed64());
+            }
+          });
         } else {
           list.add(input.readSfixed64());
         }
@@ -269,14 +321,6 @@ void _mergeFromCodedBufferReader(BuilderInfo meta, _FieldSet fs,
   }
 }

-void _readPacked(CodedBufferReader input, void Function() readFunc) {
-  input._withLimit(input.readInt32(), () {
-    while (!input.isAtEnd()) {
-      readFunc();
-    }
-  });
-}
-
 void _readPackableToListEnum(
     List list,
     BuilderInfo meta,

cc @mraleph

[mkustermann]

/cc @alexmarkov

[Manishtechbee]

@osa1 @mkustermann
Possible Approaches:
1 Manually Inlining (Not Ideal)

As you noted, manually inlining _readPacked at call sites gives a 7% performance boost but at the cost of increased binary size, which isn't ideal for browser deployment.
2 Using @pragma('vm:prefer-inline') More Aggressively

You've already used inline pragmas, but try ensuring that both _readPacked and its calling functions are aggressively inlined.
Also, check if the closure itself (i.e., the function expression passed) can be marked with @pragma('vm:prefer-inline').
3 Eliminating Closure Allocation

The allocation of AllocateClosure and AllocateContext suggests that Dart's AOT compiler isn't optimizing the closure away.
Instead of passing a closure, can you refactor _readPacked to take a function pointer directly?
Example:

void _readPacked(FunctionType readFunc) { ... }

This might avoid closure allocations altogether.
4 Checking Closure Escape Analysis

If the closure escapes the function, the compiler may not optimize it as aggressively.
Try using an explicit inline function pointer instead of a closure.
5 Investigating the Flow Graph Further

Since you've already enabled flow graph printing, analyze whether ClosureCall gets devirtualized.
If not, that might explain the performance issue.
6 Raising an Issue with the Dart VM Team

Since you're working with AOT compilation and flow graph analysis, this might be an optimization limitation in Dart.
You already tagged @mraleph (Vyacheslav Egorov), who is deeply familiar with Dart's compiler internals. If this turns out to be an actual missing optimization, it could be reported and possibly improved.

[mraleph]

@Manishtechbee please no useless LLM generated garbage on this issue tracker or you will be banned for spam.

[osa1]

We tried this making these changes with @mkustermann but it didn't quite give us the same performance as manually inlining _readPacked even though it eliminated all of the closures that we wanted to eliminate:

diff --git a/runtime/vm/compiler/backend/il.cc b/runtime/vm/compiler/backend/il.cc
index 42d89ce95f1..35258314ff3 100644
--- a/runtime/vm/compiler/backend/il.cc
+++ b/runtime/vm/compiler/backend/il.cc
@@ -2969,6 +2969,11 @@ Definition* LoadFieldInstr::Canonicalize(FlowGraph* flow_graph) {
         }
       }
       break;
+    case Slot::Kind::kClosure_context:
+      if (auto closure_allocation = orig_instance->AsAllocateClosure()) {
+        return closure_allocation->context()->definition();
+      }
+      break;
     default:
       break;
   }
@@ -3008,6 +3013,22 @@ Definition* LoadFieldInstr::Canonicalize(FlowGraph* flow_graph) {
     }
   }
 
+  if (auto instr = orig_instance->AsAllocateContext()) {
+    if (IsImmutableLoad()) {
+      for (auto use : instr->input_uses()) {
+        if (auto store = use->instruction()->AsStoreField()) {
+          if (IsDominatedBy(store)) {
+            if ((use->use_index() == StoreFieldInstr::kInstancePos) &&
+                store->slot().IsIdentical(slot())) {
+              // RELEASE_ASSERT(IsDominatedBy(store));
+              return store->value()->definition();
+            }
+          }
+        }
+      }
+    }
+  }
+
   return this;
 }
 
diff --git a/runtime/vm/compiler/backend/inliner.cc b/runtime/vm/compiler/backend/inliner.cc
index 2909c6915fc..e24e17e583f 100644
--- a/runtime/vm/compiler/backend/inliner.cc
+++ b/runtime/vm/compiler/backend/inliner.cc
@@ -464,7 +464,14 @@ class CallSites : public ValueObject {
         worklist.Add(call_info.call);
       }
     }
-    RELEASE_ASSERT(worklist.length() == calls_->length());
+    for (auto& call_info : closure_calls_) {
+      if (call_info.call->HasSSATemp()) {
+        add_to_worklist(call_info.call);
+      } else {
+        worklist.Add(call_info.call);
+      }
+    }
+    // RELEASE_ASSERT(worklist.length() == calls_->length());
     add_transitive_dependencies_to_worklist(0);
 
     // Step 2: canonicalize each definition from the worklist. We process
@@ -588,13 +602,13 @@ class CallSites : public ValueObject {
           HandleStaticCall(static_call, inline_only_profitable_methods, graph,
                            depth, nesting_depth, inlined_info);
         } else if (auto closure_call = current->AsClosureCall()) {
-          if (!inline_only_profitable_methods) {
-            // Consider closure for further inlining. Note that it will
-            // still be subject to all the inlining heuristics.
-            closure_calls_.Add({graph, closure_call, depth, nesting_depth});
-          } else {
-            // No longer consider the closure because inlining is too deep.
-          }
+          // if (!inline_only_profitable_methods) {
+          // Consider closure for further inlining. Note that it will
+          // still be subject to all the inlining heuristics.
+          closure_calls_.Add({graph, closure_call, depth, nesting_depth});
+          // } else {
+          //   // No longer consider the closure because inlining is too deep.
+          // }
         }
       }
     }

We improved the benchmark from 400us to 395us, but manually inlining _readPacked gives us 370us.

There's also something strange in this change which is the IsDominatedBy check before returning the store value. We thought the check shouldn't be necessary (and it doesn't exist for other cases in the same function) as you shouldn't be even allowed to read a field unless it's stored in a dominator block, but removing that check causes causes assertion errors.

It's also strange that eliminating a bunch of closure allocations in a hot loop gave us only a tiny improvement.

Update: the benchmarks I used didn't use the optimized code path at all, so the numbers are just noise. However flow graphs are correct: VM's inliner does not eliminate the closures, but manually inlining _readPacked (along with some other prefer-inlines) does.

[alexmarkov]

This is likely a pass ordering problem: full-blown store-load forwarding happens after inlining is already finished. Simple store-loading forwarding during canonicalization (which happens during inlining) is one way to solve it. However, such store-load forwarding does not respect control flow and usually applicable only to final fields.

With captured final local variables you need to be a bit more careful as there could be multiple assignments (stores) to a final local variable. For example:

  final int foo;
  if (...) {
    foo = 2;
  } else {
    foo = 3;
  }
  () { print(foo); } ();

That's why you probably need a IsDominatedBy check.

[a-siva]

@osa1 are you going to work on landing a change with the changes you have in the comment above ?