From f4da87d36cf3c965ad6692288a08374a64126356 Mon Sep 17 00:00:00 2001
From: Andy Ayers <andya@microsoft.com>
Date: Tue, 10 Dec 2024 15:15:21 -0800
Subject: [PATCH] JIT: boost inlining when callee unboxes an arg

Especially so if the caller is passing an exact type. This may lead
to the JIT being able to stack allocate the box and promote the underlying
payload.

Fixes #104479
---
 src/coreclr/jit/fgbasic.cpp      | 18 ++++++++++++++++++
 src/coreclr/jit/inline.def       |  2 ++
 src/coreclr/jit/inlinepolicy.cpp | 32 ++++++++++++++++++++++++++++++++
 src/coreclr/jit/inlinepolicy.h   |  4 ++++
 4 files changed, 56 insertions(+)

diff --git a/src/coreclr/jit/fgbasic.cpp b/src/coreclr/jit/fgbasic.cpp
index 30a3a83581a3b..c3ac9e7bf9d11 100644
--- a/src/coreclr/jit/fgbasic.cpp
+++ b/src/coreclr/jit/fgbasic.cpp
@@ -1106,6 +1106,24 @@ void Compiler::fgFindJumpTargets(const BYTE* codeAddr, IL_OFFSET codeSize, Fixed
                 break;
             }
 
+            case CEE_UNBOX:
+            case CEE_UNBOX_ANY:
+            {
+                if (makeInlineObservations)
+                {
+                    FgStack::FgSlot slot = pushedStack.Top();
+                    if (FgStack::IsExactArgument(slot, impInlineInfo))
+                    {
+                        compInlineResult->Note(InlineObservation::CALLSITE_UNBOX_EXACT_ARG);
+                    }
+                    else if (FgStack::IsArgument(slot))
+                    {
+                        compInlineResult->Note(InlineObservation::CALLEE_UNBOX_ARG);
+                    }
+                }
+                break;
+            }
+
             case CEE_CASTCLASS:
             case CEE_ISINST:
             {
diff --git a/src/coreclr/jit/inline.def b/src/coreclr/jit/inline.def
index b1152d6f9ecc0..efacbd4deb27b 100644
--- a/src/coreclr/jit/inline.def
+++ b/src/coreclr/jit/inline.def
@@ -106,6 +106,7 @@ INLINE_OBSERVATION(NUMBER_OF_ARGUMENTS,       int,    "number of arguments",
 INLINE_OBSERVATION(NUMBER_OF_BASIC_BLOCKS,    int,    "number of basic blocks",               INFORMATION, CALLEE)
 INLINE_OBSERVATION(NUMBER_OF_LOCALS,          int,    "number of locals",                     INFORMATION, CALLEE)
 INLINE_OBSERVATION(RANDOM_ACCEPT,             bool,   "random accept",                        INFORMATION, CALLEE)
+INLINE_OBSERVATION(UNBOX_ARG,                 int,    "callee unboxes arg",                   INFORMATION, CALLEE)
 INLINE_OBSERVATION(UNSUPPORTED_OPCODE,        bool,   "unsupported opcode",                   INFORMATION, CALLEE)
 
 // ------ Caller Correctness -------
@@ -194,6 +195,7 @@ INLINE_OBSERVATION(LOG_REPLAY_ACCEPT,         bool,   "accepted by log replay",
 INLINE_OBSERVATION(PROFILE_FREQUENCY,         double, "frequency from profile data",          INFORMATION, CALLSITE)
 INLINE_OBSERVATION(RANDOM_ACCEPT,             bool,   "random accept",                        INFORMATION, CALLSITE)
 INLINE_OBSERVATION(WEIGHT,                    int,    "frequency from block weight",          INFORMATION, CALLSITE)
+INLINE_OBSERVATION(UNBOX_EXACT_ARG,           int,    "unbox of arg with exact class",        INFORMATION, CALLSITE)
 
 // ------ Final Sentinel -------
 
diff --git a/src/coreclr/jit/inlinepolicy.cpp b/src/coreclr/jit/inlinepolicy.cpp
index d7b7bf6e1a024..d22676a62a3ea 100644
--- a/src/coreclr/jit/inlinepolicy.cpp
+++ b/src/coreclr/jit/inlinepolicy.cpp
@@ -1362,6 +1362,14 @@ void ExtendedDefaultPolicy::NoteBool(InlineObservation obs, bool value)
             m_IsCallsiteInNoReturnRegion = value;
             break;
 
+        case InlineObservation::CALLEE_UNBOX_ARG:
+            m_ArgUnbox++;
+            break;
+
+        case InlineObservation::CALLSITE_UNBOX_EXACT_ARG:
+            m_ArgUnboxExact++;
+            break;
+
         default:
             DefaultPolicy::NoteBool(obs, value);
             break;
@@ -1716,6 +1724,30 @@ double ExtendedDefaultPolicy::DetermineMultiplier()
         JITDUMP("\nPrejit root candidate has arg that feeds a conditional.  Multiplier increased to %g.", multiplier);
     }
 
+    if (m_ArgUnboxExact > 0)
+    {
+        // Callee has unbox(arg), caller supplies exact type (a box)
+        // We can likely optimize
+        multiplier += 4.0;
+        JITDUMP("\nInline candidate has %d exact arg unboxes.  Multiplier increased to %g.", m_ArgUnboxExact,
+                multiplier);
+    }
+
+    if (m_ArgUnbox > 0)
+    {
+        // Callee has unbox(arg), caller arg not known type
+        if (m_IsPrejitRoot)
+        {
+            // Assume these might be met with exact type args
+            multiplier += 4.0;
+        }
+        else
+        {
+            multiplier += 1.0;
+        }
+        JITDUMP("\nInline candidate has %d arg unboxes.  Multiplier increased to %g.", m_ArgUnboxExact, multiplier);
+    }
+
     switch (m_CallsiteFrequency)
     {
         case InlineCallsiteFrequency::RARE:
diff --git a/src/coreclr/jit/inlinepolicy.h b/src/coreclr/jit/inlinepolicy.h
index a8d8e67f1db3c..d08fbf7b32309 100644
--- a/src/coreclr/jit/inlinepolicy.h
+++ b/src/coreclr/jit/inlinepolicy.h
@@ -219,6 +219,8 @@ class ExtendedDefaultPolicy : public DefaultPolicy
         , m_UnrollableMemop(0)
         , m_Switch(0)
         , m_DivByCns(0)
+        , m_ArgUnbox(0)
+        , m_ArgUnboxExact(0)
         , m_ReturnsStructByValue(false)
         , m_IsFromValueClass(false)
         , m_NonGenericCallsGeneric(false)
@@ -272,6 +274,8 @@ class ExtendedDefaultPolicy : public DefaultPolicy
     unsigned m_UnrollableMemop;
     unsigned m_Switch;
     unsigned m_DivByCns;
+    unsigned m_ArgUnbox;
+    unsigned m_ArgUnboxExact;
     bool     m_ReturnsStructByValue       : 1;
     bool     m_IsFromValueClass           : 1;
     bool     m_NonGenericCallsGeneric     : 1;