Bug 1646378: Add WarpInlinedCall snapshot r=jandem

I based the InliningTree / CompileInfo code on IonBuilder::inlineScriptedCall.

In the future, if we want to inline cases where we didn't allocate a new ICScript, it should just be a matter of changing a few lines in maybeInlineIC. We can use the default ICScript off the target script.

(Note: we now check in maybeInlineCall that we don't already have a CallInlinedFunction. This should not normally happen, but could happen if we reset the warmup counter for a script.)

Differential Revision: https://phabricator.services.mozilla.com/D83184

js/src/jit/TrialInlining.cpp

@@ -72,7 +72,7 @@ ICStub* TrialInliner::maybeSingleStub(const ICEntry& entry) {
 }
 
 Maybe<InlinableCallData> FindInlinableCallData(ICStub* stub) {
-  InlinableCallData data;
+  Maybe<InlinableCallData> data;
 
   const CacheIRStubInfo* stubInfo = stub->cacheIRStubInfo();
   const uint8_t* stubData = stub->cacheIRStubData();
@@ -80,7 +80,6 @@ Maybe<InlinableCallData> FindInlinableCallData(ICStub* stub) {
   ObjOperandId calleeGuardOperand;
   CallFlags flags;
   uint32_t targetOffset = 0;
-  bool foundCall = false;
 
   CacheIRReader reader(stubInfo);
   while (reader.more()) {
@@ -92,7 +91,7 @@ Maybe<InlinableCallData> FindInlinableCallData(ICStub* stub) {
     switch (op) {
       case CacheOp::GuardSpecificFunction:
         // If we see a guard, remember which operand we are guarding.
-        MOZ_ASSERT(!foundCall);
+        MOZ_ASSERT(data.isNothing());
         calleeGuardOperand = reader.objOperandId();
         targetOffset = reader.stubOffset();
         mozilla::Unused << reader.stubOffset();  // nargsAndFlags
@@ -102,23 +101,35 @@ Maybe<InlinableCallData> FindInlinableCallData(ICStub* stub) {
         // operand. If it is, we know the target and can inline.
         ObjOperandId calleeOperand = reader.objOperandId();
         mozilla::DebugOnly<Int32OperandId> argcId = reader.int32OperandId();
-        CallFlags flags = reader.callFlags();
+        flags = reader.callFlags();
 
         if (calleeOperand == calleeGuardOperand) {
-          MOZ_ASSERT(!foundCall);
           MOZ_ASSERT(static_cast<OperandId&>(argcId).id() == 0);
-          foundCall = true;
-          data.calleeOperand = calleeOperand;
-          data.callFlags = flags;
-          data.endOfSharedPrefix = opStart;
-          uintptr_t rawTarget =
-              stubInfo->getStubRawWord(stubData, targetOffset);
-          data.target = reinterpret_cast<JSFunction*>(rawTarget);
+          MOZ_ASSERT(data.isNothing());
+          data.emplace();
+          data->endOfSharedPrefix = opStart;
+        }
+        break;
+      }
+      case CacheOp::CallInlinedFunction: {
+        ObjOperandId calleeOperand = reader.objOperandId();
+        mozilla::DebugOnly<Int32OperandId> argcId = reader.int32OperandId();
+        uint32_t icScriptOffset = reader.stubOffset();
+        flags = reader.callFlags();
+
+        if (calleeOperand == calleeGuardOperand) {
+          MOZ_ASSERT(static_cast<OperandId&>(argcId).id() == 0);
+          MOZ_ASSERT(data.isNothing());
+          data.emplace();
+          data->endOfSharedPrefix = opStart;
+          uintptr_t rawICScript =
+              stubInfo->getStubRawWord(stubData, icScriptOffset);
+          data->icScript = reinterpret_cast<ICScript*>(rawICScript);
         }
         break;
       }
       default:
-        if (foundCall) {
+        if (data.isSome()) {
           MOZ_ASSERT(op == CacheOp::ReturnFromIC ||
                      op == CacheOp::TypeMonitorResult);
         }
@@ -128,10 +139,13 @@ Maybe<InlinableCallData> FindInlinableCallData(ICStub* stub) {
     MOZ_ASSERT(opStart + 1 + argLength == reader.currentPosition());
   }
 
-  if (!foundCall) {
-    return mozilla::Nothing();
+  if (data.isSome()) {
+    data->calleeOperand = calleeGuardOperand;
+    data->callFlags = flags;
+    uintptr_t rawTarget = stubInfo->getStubRawWord(stubData, targetOffset);
+    data->target = reinterpret_cast<JSFunction*>(rawTarget);
   }
-  return mozilla::Some(data);
+  return data;
 }
 
 bool TrialInliner::shouldInline(JSFunction* target) {
@@ -199,6 +213,10 @@ bool TrialInliner::maybeInlineCall(const ICEntry& entry, BytecodeLocation loc) {
   if (data.isNothing()) {
     return true;
   }
+  // Ensure that we haven't already trial-inlined this callsite.
+  if (data->icScript) {
+    return true;
+  }
 
   JitSpew(JitSpew_WarpTrialInlining,
           "Inlining candidate JSOp::%s: callee function %p",

js/src/jit/TrialInlining.h

@@ -64,6 +64,7 @@ class InlinableCallData {
   CallFlags callFlags;
   const uint8_t* endOfSharedPrefix = nullptr;
   JSFunction* target = nullptr;
+  ICScript* icScript = nullptr;
 };
 
 mozilla::Maybe<InlinableCallData> FindInlinableCallData(ICStub* stub);

js/src/jit/WarpOracle.cpp

@@ -14,6 +14,7 @@
 #include "jit/CacheIR.h"
 #include "jit/CacheIRCompiler.h"
 #include "jit/CacheIROpsGenerated.h"
+#include "jit/CompileInfo.h"
 #include "jit/JitScript.h"
 #include "jit/JitSpewer.h"
 #include "jit/MIRGenerator.h"
@@ -32,6 +33,8 @@
 using namespace js;
 using namespace js::jit;
 
+using mozilla::Maybe;
+
 // WarpScriptOracle creates a WarpScriptSnapshot for a single JSScript. Note
 // that a single WarpOracle can use multiple WarpScriptOracles when scripts are
 // inlined.
@@ -739,6 +742,9 @@ AbortReasonOr<Ok> WarpScriptOracle::maybeInlineIC(WarpOpSnapshotList& snapshots,
   // * If the Baseline IC has a single ICStub we can inline, add a WarpCacheIR
   //   snapshot to transpile it to MIR.
   //
+  // * If that single ICStub is a call IC with a known target, instead add a
+  //   WarpInline snapshot to transpile the guards to MIR and inline the target.
+  //
   // * If the Baseline IC is cold (never executed), add a WarpBailout snapshot
   //   so that we can collect information in Baseline.
   //
@@ -874,6 +880,56 @@ AbortReasonOr<Ok> WarpScriptOracle::maybeInlineIC(WarpOpSnapshotList& snapshots,
 
   JitCode* jitCode = stub->jitCode();
 
+  Maybe<InlinableCallData> callData;
+  if (loc.isInvokeOp()) {
+    callData = FindInlinableCallData(stub);
+  }
+  if (callData.isSome() && callData->icScript) {
+    RootedFunction targetFunction(cx_, callData->target);
+    RootedScript targetScript(cx_, targetFunction->nonLazyScript());
+    ICScript* icScript = callData->icScript;
+    MOZ_ASSERT(targetScript->jitScript() == icScript->jitScript());
+
+    // Add the inlined script to the inline script tree.
+    LifoAlloc* lifoAlloc = alloc_.lifoAlloc();
+    InlineScriptTree* inlineScriptTree = info_->inlineScriptTree()->addCallee(
+        &alloc_, loc.toRawBytecode(), targetScript);
+    if (!inlineScriptTree) {
+      return abort(AbortReason::Alloc);
+    }
+
+    // Create a CompileInfo for the inlined script.
+    jsbytecode* osrPc = nullptr;
+    bool needsArgsObj = false;
+    CompileInfo* info = lifoAlloc->new_<CompileInfo>(
+        mirGen_.runtime, targetScript, targetFunction, osrPc,
+        info_->analysisMode(), needsArgsObj, inlineScriptTree);
+    if (!info) {
+      return abort(AbortReason::Alloc);
+    }
+
+    // Take a snapshot of the CacheIR.
+    WarpCacheIR* cacheIRSnapshot = new (alloc_.fallible())
+        WarpCacheIR(offset, jitCode, stubInfo, stubDataCopy);
+    if (!cacheIRSnapshot) {
+      return abort(AbortReason::Alloc);
+    }
+
+    // Take a snapshot of the inlined script (which may do more
+    // inlining recursively).
+    WarpScriptOracle scriptOracle(cx_, oracle_, targetScript, info, icScript);
+
+    WarpScriptSnapshot* scriptSnapshot;
+    MOZ_TRY_VAR(scriptSnapshot, scriptOracle.createScriptSnapshot());
+    oracle_->addScriptSnapshot(scriptSnapshot);
+
+    if (!AddOpSnapshot<WarpInlinedCall>(
+            alloc_, snapshots, offset, cacheIRSnapshot, scriptSnapshot, info)) {
+      return abort(AbortReason::Alloc);
+    }
+    return Ok();
+  }
+
   if (!AddOpSnapshot<WarpCacheIR>(alloc_, snapshots, offset, jitCode, stubInfo,
                                   stubDataCopy)) {
     return abort(AbortReason::Alloc);

js/src/jit/WarpSnapshot.cpp

@@ -182,6 +182,12 @@ void WarpCacheIR::dumpData(GenericPrinter& out) const {
   out.printf("(CacheIR spew unavailable)\n");
 #  endif
 }
+
+void WarpInlinedCall::dumpData(GenericPrinter& out) const {
+  out.printf("    scriptSnapshot: 0x%p\n", scriptSnapshot_);
+  cacheIRSnapshot_->dumpData(out);
+  // TODO: dump callInfo
+}
 #endif  // JS_JITSPEW
 
 template <typename T>
@@ -295,3 +301,8 @@ void WarpCacheIR::traceData(JSTracer* trc) {
   // TODO: trace pointers in stub data. Beware of nursery indexes in the stub
   // data. See WarpObjectField.
 }
+
+void WarpInlinedCall::traceData(JSTracer* trc) {
+  // Note: scriptSnapshot_ is traced through WarpSnapshot.
+  cacheIRSnapshot_->trace(trc);
+}

js/src/jit/WarpSnapshot.h

@@ -23,6 +23,8 @@ class ModuleEnvironmentObject;
 namespace jit {
 
 class CacheIRStubInfo;
+class CompileInfo;
+class WarpScriptSnapshot;
 
 #define WARP_OP_SNAPSHOT_LIST(_) \
   _(WarpArguments)               \
@@ -35,7 +37,8 @@ class CacheIRStubInfo;
   _(WarpNewArray)                \
   _(WarpNewObject)               \
   _(WarpBailout)                 \
-  _(WarpCacheIR)
+  _(WarpCacheIR)                 \
+  _(WarpInlinedCall)
 
 // Wrapper for GC things stored in WarpSnapshot. Asserts the GC pointer is not
 // nursery-allocated. These pointers must be traced using TraceWarpGCPtr.
@@ -337,6 +340,36 @@ class WarpObjectField {
   }
 };
 
+// Information for inlining a scripted call IC.
+class WarpInlinedCall : public WarpOpSnapshot {
+  // Used for generating the correct guards.
+  WarpCacheIR* cacheIRSnapshot_;
+
+  // Used for generating the inlined code.
+  WarpScriptSnapshot* scriptSnapshot_;
+  CompileInfo* info_;
+
+ public:
+  static constexpr Kind ThisKind = Kind::WarpInlinedCall;
+
+  WarpInlinedCall(uint32_t offset, WarpCacheIR* cacheIRSnapshot,
+                  WarpScriptSnapshot* scriptSnapshot, CompileInfo* info)
+      : WarpOpSnapshot(ThisKind, offset),
+        cacheIRSnapshot_(cacheIRSnapshot),
+        scriptSnapshot_(scriptSnapshot),
+        info_(info) {}
+
+  WarpCacheIR* cacheIRSnapshot() const { return cacheIRSnapshot_; }
+  WarpScriptSnapshot* scriptSnapshot() const { return scriptSnapshot_; }
+  CompileInfo* info() const { return info_; }
+
+  void traceData(JSTracer* trc);
+
+#ifdef JS_JITSPEW
+  void dumpData(GenericPrinter& out) const;
+#endif
+};
+
 // Template object for JSOp::Rest.
 class WarpRest : public WarpOpSnapshot {
   WarpGCPtr<ArrayObject*> templateObject_;

js/src/vm/BytecodeLocation.h

@@ -198,6 +198,8 @@ class BytecodeLocation {
 
   bool isSpreadOp() const { return IsSpreadOp(getOp()); }
 
+  bool isInvokeOp() const { return IsInvokeOp(getOp()); }
+
   bool resultIsPopped() const {
     MOZ_ASSERT(StackDefs(rawBytecode_) == 1);
     return BytecodeIsPopped(rawBytecode_);