Transform indirect calls to direct ones in the importer

src/coreclr/jit/compiler.h

@@ -3830,6 +3830,8 @@ class Compiler
     bool impCanPInvokeInlineCallSite(BasicBlock* block);
     void impCheckForPInvokeCall(
         GenTreeCall* call, CORINFO_METHOD_HANDLE methHnd, CORINFO_SIG_INFO* sig, unsigned mflags, BasicBlock* block);
+
+    bool impCanSubstituteSig(CORINFO_SIG_INFO* sourceSig, CORINFO_SIG_INFO* targetSig, var_types sourceThis, var_types targetThis);
     GenTreeCall* impImportIndirectCall(CORINFO_SIG_INFO* sig, const DebugInfo& di = DebugInfo());
     void impPopArgsForUnmanagedCall(GenTreeCall* call, CORINFO_SIG_INFO* sig);
 
@@ -3858,6 +3860,7 @@ class Compiler
 
     GenTree* impInitClass(CORINFO_RESOLVED_TOKEN* pResolvedToken);
 
+    GenTree* impGetNodeFromLocal(GenTree* node);
     GenTree* impImportStaticReadOnlyField(CORINFO_FIELD_HANDLE field, CORINFO_CLASS_HANDLE ownerCls);
     GenTree* impImportCnsTreeFromBuffer(uint8_t* buffer, var_types valueType);
 

src/coreclr/jit/importer.cpp

@@ -3764,6 +3764,65 @@ GenTree* Compiler::impInitClass(CORINFO_RESOLVED_TOKEN* pResolvedToken)
     return node;
 }
 
+//------------------------------------------------------------------------
+// impGetNodeFromLocal: Tries to return the node that's assigned
+// to the provided local.
+//
+// Arguments:
+//    node - GT_LCL_VAR whose value is searched for
+//
+// Return Value:
+//    The tree representing the node assigned to the variable when possible,
+//    nullptr otherwise.
+//
+GenTree* Compiler::impGetNodeFromLocal(GenTree* node)
+{
+    assert(node != nullptr);
+    assert(node->OperIs(GT_LCL_VAR));
+
+    unsigned lclNum = node->AsLclVarCommon()->GetLclNum();
+
+    if (lvaTable[lclNum].lvSingleDef == 0)
+    {
+        return nullptr;
+    }
+
+    auto findValue = [&](Statement* stmtList) -> GenTree* {
+        for (Statement* stmt : StatementList(stmtList))
+        {
+            GenTree* root = stmt->GetRootNode();
+            if (root->OperIs(GT_ASG) && root->AsOp()->gtOp1->OperIs(GT_LCL_VAR) &&
+                root->AsOp()->gtOp1->AsLclVarCommon()->GetLclNum() == lclNum)
+            {
+                return root->AsOp()->gtOp2;
+            }
+        }
+        return nullptr;
+    };
+
+    GenTree*    valueNode = findValue(impStmtList);
+    BasicBlock* bb        = fgFirstBB;
+    while (valueNode == nullptr && bb != nullptr)
+    {
+        valueNode = findValue(bb->bbStmtList);
+        if (valueNode == nullptr && bb->NumSucc(this) == 1)
+        {
+            bb = bb->GetSucc(0, this);
+        }
+        else
+        {
+            bb = nullptr;
+        }
+    }
+
+    if (valueNode != nullptr && valueNode->OperIs(GT_LCL_VAR))
+    {
+        return impGetNodeFromLocal(valueNode);
+    }
+
+    return valueNode;
+}
+
 //------------------------------------------------------------------------
 // impImportStaticReadOnlyField: Tries to import 'static readonly' field
 //    as a constant if the host type is statically initialized.

src/coreclr/jit/importercalls.cpp

@@ -104,13 +104,17 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
     bool checkForSmallType  = false;
     bool bIntrinsicImported = false;
 
-    CORINFO_SIG_INFO calliSig;
+    CORINFO_SIG_INFO originalSig = {};
     NewCallArg       extraArg;
 
-    /*-------------------------------------------------------------------------
-     * First create the call node
-     */
+    // run transformations when instrumenting to not pollute PGO data
+    bool                  optimizedOrInstrumented = opts.OptimizationEnabled() || opts.IsInstrumented();
+    CORINFO_METHOD_HANDLE replacementMethod       = nullptr;
+    GenTree*              newThis                 = nullptr;
+    var_types             oldThis                 = TYP_UNDEF;
+    var_types             targetThis              = TYP_UNDEF;
 
+    // handle special import cases
     if (opcode == CEE_CALLI)
     {
         if (IsTargetAbi(CORINFO_NATIVEAOT_ABI))
@@ -125,25 +129,107 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
         }
 
         /* Get the call site sig */
-        eeGetSig(pResolvedToken->token, pResolvedToken->tokenScope, pResolvedToken->tokenContext, &calliSig);
+        eeGetSig(pResolvedToken->token, pResolvedToken->tokenScope, pResolvedToken->tokenContext, &originalSig);
+
+        if (!optimizedOrInstrumented)
+        {
+            // ignore
+        }
+        else if (originalSig.getCallConv() == CORINFO_CALLCONV_DEFAULT)
+        {
+            JITDUMP("\n\nimpImportCall trying to import calli as call\n");
+            GenTree* fptr = impStackTop().val;
+            if (fptr->OperIs(GT_LCL_VAR))
+            {
+                JITDUMP("impImportCall trying to import calli as call - trying to substitute LCL_VAR\n");
+                GenTree* lclValue = impGetNodeFromLocal(fptr);
+                if (lclValue != nullptr)
+                {
+                    fptr = lclValue;
+                }
+            }
+            if (fptr->OperIs(GT_FTN_ADDR))
+            {
+                replacementMethod = fptr->AsFptrVal()->gtFptrMethod;
+            }
+            else
+            {
+                JITDUMP("impImportCall failed to import calli as call - address node not found\n");
+            }
+        }
+        else
+        {
+            JITDUMP("\n\nimpImportCall failed to import calli as call - call conv %u is not managed\n",
+                    originalSig.getCallConv());
+        }
+    }
+
+    if (replacementMethod != nullptr)
+    {
+        JITDUMP("impImportCall trying to transform call - found target method %s\n",
+                eeGetMethodName(replacementMethod));
+        CORINFO_SIG_INFO     methodSig;
+        CORINFO_CLASS_HANDLE targetClass = info.compCompHnd->getMethodClass(replacementMethod);
+        info.compCompHnd->getMethodSig(replacementMethod, &methodSig, targetClass);
+
+        if (methodSig.hasImplicitThis() && targetThis == TYP_UNDEF)
+        {
+            targetThis = eeIsValueClass(targetClass) ? TYP_BYREF : TYP_REF;
+        }
+
+        unsigned replacementFlags = info.compCompHnd->getMethodAttribs(replacementMethod);
 
-        callRetTyp = JITtype2varType(calliSig.retType);
+        if ((replacementFlags & CORINFO_FLG_PINVOKE) != 0)
+        {
+            JITDUMP("impImportCall aborting transformation - found PInvoke\n");
+        }
+        else if (impCanSubstituteSig(&originalSig, &methodSig, oldThis, targetThis))
+        {
+            impPopStack();
+            if (newThis != nullptr)
+            {
+                assert(oldThis == TYP_REF);
+                assert(targetThis == TYP_REF);
+                CORINFO_CLASS_HANDLE thisCls = NO_CLASS_HANDLE;
+                info.compCompHnd->getArgType(&methodSig, methodSig.args, &thisCls);
+                impPushOnStack(newThis, typeInfo(TI_REF, thisCls));
+            }
+            JITDUMP("impImportCall transforming call\n");
+            pResolvedToken->hMethod = replacementMethod;
+            pResolvedToken->hClass  = targetClass;
+
+            callInfo->sig         = methodSig;
+            callInfo->hMethod     = replacementMethod;
+            callInfo->methodFlags = replacementFlags;
+            callInfo->classFlags  = info.compCompHnd->getClassAttribs(targetClass);
 
-        call = impImportIndirectCall(&calliSig, di);
+            return impImportCall(CEE_CALL, pResolvedToken, nullptr, nullptr, prefixFlags, callInfo, rawILOffset);
+        }
+    }
+
+    /*-------------------------------------------------------------------------
+     * First create the call node
+     */
+
+    if (opcode == CEE_CALLI)
+    {
+        callRetTyp = JITtype2varType(originalSig.retType);
+
+        call = impImportIndirectCall(&originalSig, di);
 
         // We don't know the target method, so we have to infer the flags, or
         // assume the worst-case.
-        mflags = (calliSig.callConv & CORINFO_CALLCONV_HASTHIS) ? 0 : CORINFO_FLG_STATIC;
+        mflags = (originalSig.callConv & CORINFO_CALLCONV_HASTHIS) ? 0 : CORINFO_FLG_STATIC;
 
 #ifdef DEBUG
         if (verbose)
         {
-            unsigned structSize = (callRetTyp == TYP_STRUCT) ? eeTryGetClassSize(calliSig.retTypeSigClass) : 0;
+            unsigned structSize = (callRetTyp == TYP_STRUCT) ? eeTryGetClassSize(originalSig.retTypeSigClass) : 0;
             printf("\nIn Compiler::impImportCall: opcode is %s, kind=%d, callRetType is %s, structSize is %u\n",
                    opcodeNames[opcode], callInfo->kind, varTypeName(callRetTyp), structSize);
         }
 #endif
-        sig = &calliSig;
+        sig = &originalSig;
     }
     else // (opcode != CEE_CALLI)
     {
@@ -1622,6 +1708,138 @@ GenTree* Compiler::impFixupCallStructReturn(GenTreeCall* call, CORINFO_CLASS_HAN
 #endif // FEATURE_MULTIREG_RET
 }
 
+//-----------------------------------------------------------------------------------
+//  impCanSubstituteSig: Checks whether it's safe to replace a call with another one.
+//  This DOES NOT check if the calls are actually compatible, it only checks if their trees are.
+//  Use ONLY when code will call the method with target sig anyway.
+//
+//  Arguments:
+//    sourceSig      - original call signature
+//    targetSig      - new call signature
+//    transformation - transformations performed on the original signature
+//
+//  Return Value:
+//    Whether it's safe to change the IR to call the target method
+//
+bool Compiler::impCanSubstituteSig(CORINFO_SIG_INFO* sourceSig,
+                                   CORINFO_SIG_INFO* targetSig,
+                                   var_types         sourceThis,
+                                   var_types         targetThis)
+{
+    assert(sourceSig->hasImplicitThis() || sourceThis == TYP_UNDEF);
+    assert(targetSig->hasImplicitThis() || targetThis == TYP_UNDEF || targetThis == TYP_VOID);
+    assert(!targetSig->hasImplicitThis() || targetThis != TYP_VOID);
+    assert(sourceThis != TYP_VOID);
+
+    if (sourceSig->getCallConv() != targetSig->getCallConv())
+    {
+        JITDUMP("impCanSubstituteSig returning false - call conv %u != %u\n", sourceSig->callConv, targetSig->callConv);
+        return false;
+    }
+
+    if ((sourceSig->hasImplicitThis() && sourceThis == TYP_UNDEF) ||
+        (targetSig->hasImplicitThis() && targetThis == TYP_UNDEF))
+    {
+        JITDUMP("impCanSubstituteSig returning false - unknown this type\n");
+        return false;
+    }
+
+    unsigned sourceArgCount = sourceSig->totalILArgs();
+    if (targetThis == TYP_VOID)
+    {
+        assert(sourceSig->hasImplicitThis() && sourceThis == TYP_REF);
+        sourceArgCount--;
+    }
+
+    if (sourceArgCount != targetSig->totalILArgs())
+    {
+        JITDUMP("impCanSubstituteSig returning false - args count %u != %u\n", sourceArgCount,
+                targetSig->totalILArgs());
+        return false;
+    }
+
+    if (sourceSig->retType != targetSig->retType)
+    {
+        JITDUMP("impCanSubstituteSig returning false - return type %u != %u\n", (unsigned)sourceSig->retType,
+                (unsigned)targetSig->retType);
+        return false;
+    }
+
+    if (sourceSig->retType == CORINFO_TYPE_VALUECLASS || sourceSig->retType == CORINFO_TYPE_REFANY)
+    {
+        ClassLayout* layoutRetA = typGetObjLayout(sourceSig->retTypeClass);
+        ClassLayout* layoutRetB = typGetObjLayout(targetSig->retTypeClass);
+
+        if (!ClassLayout::AreCompatible(layoutRetA, layoutRetB))
+        {
+            JITDUMP("impCanSubstituteSig returning false - return type %u is incompatible with %u\n",
+                    (unsigned)sourceSig->retType, (unsigned)targetSig->retType);
+            return false;
+        }
+    }
+
+    CORINFO_ARG_LIST_HANDLE sourceArg = sourceSig->args;
+    CORINFO_ARG_LIST_HANDLE targetArg = targetSig->args;
+
+    unsigned numArgs = targetSig->totalILArgs();
+
+    if ((sourceSig->hasImplicitThis() && targetThis != TYP_VOID) || targetSig->hasImplicitThis())
+    {
+        if (sourceThis == TYP_UNDEF)
+        {
+            sourceThis = eeGetArgType(sourceArg, sourceSig);
+            sourceArg  = info.compCompHnd->getArgNext(sourceArg);
+            numArgs--;
+        }
+        else
+        {
+            targetThis = eeGetArgType(targetArg, targetSig);
+            targetArg  = info.compCompHnd->getArgNext(targetArg);
+        }
+        assert(sourceThis == TYP_REF || sourceThis == TYP_BYREF || sourceThis == TYP_I_IMPL);
+        assert(targetThis == TYP_REF || targetThis == TYP_BYREF || targetThis == TYP_I_IMPL);
+        if (sourceThis != targetThis)
+        {
+            JITDUMP("impCanSubstituteSig returning false - this type %s != %s\n", varTypeName(sourceThis),
+                    varTypeName(targetThis));
+            return false;
+        }
+    }
+
+    for (unsigned i = 0; i < numArgs;
+         i++, sourceArg = info.compCompHnd->getArgNext(sourceArg), targetArg = info.compCompHnd->getArgNext(targetArg))
+    {
+        var_types sourceType = eeGetArgType(sourceArg, sourceSig);
+        var_types targetType = eeGetArgType(targetArg, targetSig);
+        if (sourceType != targetType)
+        {
+            JITDUMP("impCanSubstituteSig returning false - parameter %u type %s != %s\n", i, varTypeName(sourceType),
+                    varTypeName(targetType));
+            return false;
+        }
+
+        if (varTypeIsStruct(sourceType) && varTypeIsStruct(targetType))
+        {
+            CORINFO_CLASS_HANDLE sourceClassHnd = NO_CLASS_HANDLE;
+            CORINFO_CLASS_HANDLE targetClassHnd = NO_CLASS_HANDLE;
+            info.compCompHnd->getArgType(sourceSig, sourceArg, &sourceClassHnd);
+            info.compCompHnd->getArgType(targetSig, targetArg, &targetClassHnd);
+
+            ClassLayout* sourceLayout = typGetObjLayout(sourceClassHnd);
+            ClassLayout* targetLayout = typGetObjLayout(targetClassHnd);
+
+            if (!ClassLayout::AreCompatible(sourceLayout, targetLayout))
+            {
+                JITDUMP("impCanSubstituteSig returning false - parameter %u type %s is inconmpatible with %s\n", i,
+                        varTypeName(sourceType), varTypeName(targetType));
+                return false;
+            }
+        }
+    }
+
+    return true;
+}
+
 GenTreeCall* Compiler::impImportIndirectCall(CORINFO_SIG_INFO* sig, const DebugInfo& di)
 {
     var_types callRetTyp = JITtype2varType(sig->retType);

src/tests/JIT/opt/Devirtualization/CMakeLists.txt

@@ -0,0 +1,5 @@
+project (IndirectNative)
+set(SOURCES IndirectNative.cpp)
+add_library (IndirectNative SHARED ${SOURCES})
+# add the install targets
+install (TARGETS IndirectNative DESTINATION bin)