Use loop depth to decide whether to JIT functions

Summary:
Avoid JIT for functions that don't have loops if they're not called
much.

This is configured using `-Xjit-threshold`, which will be `1 << 5` by
default.

Reviewed By: tmikov

Differential Revision: D65497432

fbshipit-source-id: f5763871a9a19f1b5dfd92e0ce0f43b4b180f07f

include/hermes/ConsoleHost/ConsoleHost.h

@@ -142,6 +142,9 @@ struct ExecuteOptions {
   /// Force JIT on all functions.
   bool forceJIT{false};
 
+  /// JIT compilation threshold.
+  uint32_t jitThreshold{1 << 5};
+
   /// Dump JIT'ed code.
   unsigned dumpJITCode{0};
 

include/hermes/VM/JIT/JIT.h

@@ -94,6 +94,12 @@ class JITContext {
   /// Set the flag to force jitting of all functions.
   void setForceJIT(bool force) {}
 
+  /// Set the default threshold for function execution count before a function
+  /// is compiled. On a per-function basis, the count may be altered based on
+  /// internal heuristics.
+  /// Can be overridden by setForceJIT(true).
+  void setDefaultExecThreshold(uint32_t threshold) {}
+
   /// Set the flag to emit asserts in the JIT'ed code.
   void setEmitAsserts(bool emitAsserts) {}
 

include/hermes/VM/JIT/arm64/JIT.h

@@ -53,6 +53,14 @@ class JITContext {
     enabled_ = enabled;
   }
 
+  /// Set the default threshold for function execution count before a function
+  /// is compiled. On a per-function basis, the count may be altered based on
+  /// internal heuristics.
+  /// Can be overridden by setForceJIT(true).
+  void setDefaultExecThreshold(uint32_t threshold) {
+    defaultExecThreshold_ = threshold;
+  }
+
   /// Enable or disable dumping JIT'ed Code.
   void setDumpJITCode(unsigned dump) {
     dumpJITCode_ = dump;
@@ -75,7 +83,7 @@ class JITContext {
 
   /// Set the flag to force jitting of all functions.
   void setForceJIT(bool force) {
-    execThreshold_ = force ? 0 : DEFAULT_EXEC_THRESHOLD;
+    forceJIT_ = force;
   }
 
   /// Set the flag to emit asserts in the JIT'ed code.
@@ -105,12 +113,13 @@ class JITContext {
   bool crashOnError_{false};
   /// Whether to emit asserts in the JIT'ed code.
   bool emitAsserts_{false};
+  /// Whether to force jitting of all functions.
+  /// If true, ignores the default exec threshold completely.
+  bool forceJIT_{false};
 
-  /// Execution threshold before a function is compiled.
-  unsigned execThreshold_ = 0;
-
-  /// The JIT default threshold for function execution count
-  static constexpr uint32_t DEFAULT_EXEC_THRESHOLD = 0;
+  /// The JIT threshold for function execution count.
+  /// Lowered based on the loop depth before deciding whether to JIT.
+  uint32_t defaultExecThreshold_ = 1 << 5;
 };
 
 LLVM_ATTRIBUTE_ALWAYS_INLINE
@@ -124,8 +133,18 @@ inline JITCompiledFunctionPtr JITContext::compile(
     return nullptr;
   if (LLVM_LIKELY(codeBlock->getDontJIT()))
     return nullptr;
-  if (LLVM_LIKELY(codeBlock->getExecutionCount() < DEFAULT_EXEC_THRESHOLD))
+
+  uint32_t loopDepth = codeBlock->getFunctionHeader().loopDepth();
+  // It's possible that if the loop depth is too high, we will set the
+  // execThreshold to 0 for this function, but that's OK because we want to JIT
+  // it immediately.
+  assert(loopDepth <= 3 && "loopDepth is larger than expected");
+  uint32_t execThreshold =
+      forceJIT_ ? 0 : (defaultExecThreshold_ >> (loopDepth * 2));
+
+  if (LLVM_LIKELY(codeBlock->getExecutionCount() < execThreshold))
     return nullptr;
+
   return compileImpl(runtime, codeBlock);
 }
 

include/hermes/VM/RuntimeFlags.h

@@ -219,6 +219,13 @@ struct VMOnlyRuntimeFlags {
       llvh::cl::desc("force JIT compilation of every function"),
       llvh::cl::init(false)};
 
+  llvh::cl::opt<uint32_t> JITThreshold{
+      "Xjit-threshold",
+      llvh::cl::Hidden,
+      llvh::cl::cat(RuntimeCategory),
+      llvh::cl::desc("default minimum number of invocations to JIT compile"),
+      llvh::cl::init(1 << 5)};
+
   /// To get the value of this CLI option, use the method below.
   llvh::cl::opt<unsigned> DumpJITCode{
       "Xdump-jitcode",

lib/ConsoleHost/ConsoleHost.cpp

@@ -398,6 +398,7 @@ bool executeHBCBytecodeImpl(
 
   // TODO: surely this should use RuntimeConfig?
   runtime->getJITContext().setForceJIT(options.forceJIT);
+  runtime->getJITContext().setDefaultExecThreshold(options.jitThreshold);
   runtime->getJITContext().setDumpJITCode(options.dumpJITCode);
   runtime->getJITContext().setCrashOnError(options.jitCrashOnError);
   runtime->getJITContext().setEmitAsserts(options.jitEmitAsserts);

test/jit/count.js

@@ -0,0 +1,106 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// RUN: %hermes -Xforce-jit=0 -fno-inline -Xjit -Xjit-crash-on-error -Xdump-jitcode=2 %s | %FileCheckOrRegen --match-full-lines %s
+// REQUIRES: jit
+
+// Demonstrate the various JIT levels.
+// Update if JIT heuristics are updated to illustrate that they still work.
+
+function loop0(arr) {
+  var sum = 0;
+  return sum;
+}
+
+function loop1(arr) {
+  var sum = 0;
+  for (var i = 0; i < arr.length; ++i)
+    sum += arr[i];
+  return sum;
+}
+
+function loop2(arr) {
+  var sum = 0;
+  for (var i = 0; i < arr.length; ++i)
+    for (var j = 0; j < arr.length; ++j)
+      sum += arr[i] + arr[j];
+  return sum;
+}
+
+var arr = Array(10);
+arr.fill(0, 0, 10)
+for (var i = 0; i < 50; ++i) {
+  print(i);
+  loop0(arr);
+  loop1(arr);
+  loop2(arr);
+}
+
+// Auto-generated content below. Please do not modify manually.
+
+// CHECK:0
+// CHECK-NEXT:1
+// CHECK-NEXT:2
+
+// CHECK:JIT compilation of FunctionID 3, 'loop2'
+
+// CHECK:JIT successfully compiled FunctionID 3, 'loop2'
+// CHECK-NEXT:3
+// CHECK-NEXT:4
+// CHECK-NEXT:5
+// CHECK-NEXT:6
+// CHECK-NEXT:7
+// CHECK-NEXT:8
+
+// CHECK:JIT compilation of FunctionID 2, 'loop1'
+
+// CHECK:JIT successfully compiled FunctionID 2, 'loop1'
+// CHECK-NEXT:9
+// CHECK-NEXT:10
+// CHECK-NEXT:11
+// CHECK-NEXT:12
+// CHECK-NEXT:13
+// CHECK-NEXT:14
+// CHECK-NEXT:15
+// CHECK-NEXT:16
+// CHECK-NEXT:17
+// CHECK-NEXT:18
+// CHECK-NEXT:19
+// CHECK-NEXT:20
+// CHECK-NEXT:21
+// CHECK-NEXT:22
+// CHECK-NEXT:23
+// CHECK-NEXT:24
+// CHECK-NEXT:25
+// CHECK-NEXT:26
+// CHECK-NEXT:27
+// CHECK-NEXT:28
+// CHECK-NEXT:29
+// CHECK-NEXT:30
+// CHECK-NEXT:31
+// CHECK-NEXT:32
+
+// CHECK:JIT compilation of FunctionID 1, 'loop0'
+
+// CHECK:JIT successfully compiled FunctionID 1, 'loop0'
+// CHECK-NEXT:33
+// CHECK-NEXT:34
+// CHECK-NEXT:35
+// CHECK-NEXT:36
+// CHECK-NEXT:37
+// CHECK-NEXT:38
+// CHECK-NEXT:39
+// CHECK-NEXT:40
+// CHECK-NEXT:41
+// CHECK-NEXT:42
+// CHECK-NEXT:43
+// CHECK-NEXT:44
+// CHECK-NEXT:45
+// CHECK-NEXT:46
+// CHECK-NEXT:47
+// CHECK-NEXT:48
+// CHECK-NEXT:49

tools/hermes/hermes.cpp

@@ -117,6 +117,7 @@ static int executeHBCBytecodeFromCL(
   options.stopAfterInit = false;
   options.timeLimit = flags.ExecutionTimeLimit;
   options.forceJIT = flags.ForceJIT;
+  options.jitThreshold = flags.JITThreshold;
   options.dumpJITCode = flags.DumpJITCode;
   options.jitCrashOnError = flags.JITCrashOnError;
   options.jitEmitAsserts = flags.JITEmitAsserts;