Consider adding a cache of well known tokens to the JIT

[tannergooding]

Today, the JIT uses a set of heuristics to determine whether or not a method is profitable to inline. However, due to tokens being unique per module and due to the cost associated with resolving tokens the JIT (when inlining) can't easily see when certain code patterns will result in constants or otherwise dead code elimination. This can hinder the ability for the JIT to correctly inline methods with things like Isa.IsSupported, if (typeof(T) == typeof(...)) checks, or even string.Length.

As such, I propose we create a cache (effectively a Dictionary<TokenId, MethodInfo> or similar, per module) of well known methods (i.e. certain NamedIntrinsics) so the inliner can correctly determine if the inlining is profitable.

This table could be built up during Tier 0 and would be available during Tier 1 so that the inliner can correctly see when a call is constant and do the appropriate dead code elimination and size adjustments.

[tannergooding]

CC. @AndyAyersMS, @EgorBo

I had brought this idea up a couple times on Discord and decided to log a formal issue. I think it will largely handle the issue around the most critical scenarios in a mostly pay for play manner

[tannergooding]

On a meta point, there has also been previous discussion around concepts like "jit hints" and this would allow such work to be prototyped since the helper method would be well-known during importation.

[EgorBo]

Example of such an heuristic: EgorBo@3810c21
The idea is to increase the benefit multiplier if we feed get_Length with a string literal:

Inline candidate has const arg that feeds string.Length (+ const test).  Multiplier increased to XX.

It was originally rejected due to "resolvetoken" call but can be re-considered once we implement that cache of tokens ^.

[AndyAyersMS]

It would be better, I suspect, to implement a cache like this on the runtime side of the jit interface, as the runtime is more aware of various complications (collectible assemblies and such), has means for managing caches that might need periodic trimming or cleaning, and is able to use locks and other synchronization mechanisms that wouldn't be appropriate in the jit.

If that's the case, then what we're talking about seems like a "lightweight" version of resolveToken that has limited resolution abilities, and keeps a cache with

• positive mappings (perhaps returning a limited set of kinds/properties instead of handles) for the tokens that map to the "well known" objects of interest
• negative mappings for other tokens (things that are not of particular interest)

Given the above, it would be interesting to sketch what the jit interface API might look like.

Note if the aim of this is to improve inlining, we'll often be unable to make strong immediate inferences with the information we get back. Suppose during the IL scan this new API lets the jit know a particular call returns true. Connecting that with potential "savings" from being able to partially import the method is not easy as we don't know how the return value influences future computation and don't understand the control flow even if we can tell that value influences a branch.

So instead what we'd need to do is simply note that the IL contains "known return value" calls that we suspect are conveying valuable bits of information, and then use those notes to boost the viability of the method as an inline candidate. Then we'd revisit the method at the next stage of the inlinee evaluation pipeline, where we'd import the method for real and build IR and optimize based on that and verify that good things indeed happened (or more likely, just assume based on the hints that good things would happen).

[tannergooding]

@AndyAyersMS, that makes sense to me.

I believe the most important information for the purpose of inlining is NamedIntrinsicId. There is a lot we could theoretically do with CORINFO_RESOLVED_TOKEN or CORINFO_CALL_INFO, but I think most of the immediately interesting optimizations are limited to the named intrinsics we have in the JIT. This also keeps the maximum size of such caches restricted and fairly small and gives us a good place to "register" a token, as we can simply do it as part of lookupNamedIntrinsic for those we are interested in.

I think, therefore, that the surface area for the JIT/EE interface might be relatively simplistic such as:

void RegisterWellKnownToken(const CORINFO_RESOLVED_TOKEN* pResolvedToken, NamedIntrinsicId intrinsicId);
bool TryResolveWellKnownToken(mdToken token, CORINFO_RESOLVED_TOKEN* pResolvedToken, NamedIntrinsicId* intrinsicId);

We would then flow the resolved token down into impImportCall and impIntrinsic. This would be passed into lookupNamedIntrinsic where interesting named intrinsics can use RegisterWellKnownToken to place themselves into the cache. To start with, this would be cases such as Isa.IsSupported for hardware intrinsics, Type.op_Equality for generic checks, and string.Length for constant strings. The importer, when it encounters a call would attempt to use TryResolveWellKnownToken. The call can be skipped if optimizations are not enabled and should quickly return if the token is not in the cache. If there was a named intrinsic in the cache, the inliner can then use the returned intrinsicId to improve its own heuristics.

Say, for example, that there is a if (typeof(...) == typeof(...)) check. If we are in a generic context and T is constrained to struct then there is a very high likelihood that this represents a constant check and the inline multiplier can be increased. Ideally the inliner can also adjust itself if it sees many of these Type.op_Equality calls.

Similar heuristics can happen for Isa.IsSupported calls as the lookupNamedIntrinsic will have converted them to NI_IsSupported_True or NI_IsSupported_False letting us know if they are "constant true" or "constant false" and therefore allowing the inline multipler to be increased.

Having the cache store and return the CORINFO_RESOLVED_TOKEN allows future expansion if appropriate.

[jkotas]

void RegisterWellKnownToken(const CORINFO_RESOLVED_TOKEN* pResolvedToken, NamedIntrinsicId intrinsicId);

Should the first argument be just CORINFO_METHOD_HANDLE ?

[jkotas]

We may want to look at this as bag of cached key/value pairs attached to CORINFO_METHOD_HANDLE:

bool GetCachedData(CORINFO_METHOD_HANDLE method, intptr_t key, intptr_t* pValue);
void SetCachedData(CORINFO_METHOD_HANDLE method, intptr_t key, intptr_t value);

The JIT can define as many keys as it needs and the JIT owns interpretation of the value for given key, e.g. it can be the intrinsic IDs or cached observations about inlinees. The VM side would be free to flush the cache at will, e.g. when it becomes too big or when unloading occurs.

[AndyAyersMS]

Should the first argument be just CORINFO_METHOD_HANDLE ?

There are merits to both token-based lookup and handle based cache lookup.

Tanner's proposal is aimed at providing a cheap form of token resolution for use by the very early stages of inline candidate evaluation, where we're just scanning the IL stream, so that we can recognize some calls as being of special interest (like the Isa.IsSupported calls that will always resolve to a constant) and encourage the inliner to take a deeper look.

A more general handle keyed cache could be use to amortize costs of more detailed later stages of analysis, eg how much of an inlinee actually ends up getting imported, given that some of the calls it makes resolve to known constants at jit time.

Longer term, it also could be used by an AOT jit querying into interprocedural analysis data, where the shape of the data will likely be dictated by an earlier analysis performed by the jit.

[tannergooding]

Should the first argument be just CORINFO_METHOD_HANDLE ?

It certainly could be and that would be more extensible overall. That being said, I was hoping to attack a more specific issue first and something which I think may be better suited to be an independent API from a "you can cache any data you want" feature.

Namely, metadata tokens are unique per module and are somewhat "expensive" to resolve so we don't want to do so when doing an "is it profitable to inline this" check. This is often fine because most calls aren't "interesting" and won't have significant impact on inlining.
However, we also have a number of well-known methods which are more often than not a NamedIntrinsic and so we could maintain a small limited cache of these well-known tokens per module to assist in known problematic inlining scenarios used in perf critical/sensitive code.

If tracked on the VM side, this would effectively require a lazily allocated dictionary/map of metadata tokens to resolved token and NamedIntrinsic identifier. For a module that does not use any of these "well-known" intrinsics, there would be no additional allocations/overhead and for a module that does use them, it would be limited by the number of intrinsics we want to recognize. The most interesting ones being those that play into constant folding such as Isa.IsSupported, Type.op_Equality, string.Length, IsReferenceOrContainsReferences, and a couple others. You could imagine a similar "trivial" cache for well-known types such as string and the other primitive types (int32, char, potentially Vector4, etc).

Unlike a general-purpose "cache any information you want API", this API is known to be small and bounded and so there wouldn't need to be a huge concern around flushing it, since the overhead will never be more than a handful of entries per module and most modules will likely be empty or only contain a couple entries.

[jkotas]

There are merits to both token-based lookup and handle based cache lookup.

The expensive part of CORINFO_RESOLVED_TOKEN is dealing with the context, creating method instantiations, etc. The context-free token->HANDLE lookup is cheap if the token has been resolved already. It is basically just an lookup in an array using RID part of the token. We can streamline and aggresively inline this path to make it very fast if necessary.

If you would like to have a raw API that just can attach information for token, the API should explicitly take just CORINFO_MODULE_HANDLE and mdToken instead of the whole CORINFO_RESOLVED_TOKEN to make it clear that it does not take the context into account.

NamedIntrinsic and so we could maintain a small limited cache of these well-known tokens per module

If you do not find the token in cache, how are you going to tell whether you have not seen this token yet vs. you have seen the token and it is not a well-known method?

[tannergooding]

If you do not find the token in cache, how are you going to tell whether you have not seen this token yet vs. you have seen the token and it is not a well-known method?

The intent was that lookupNamedIntrinsic would pass in the CORINFO_RESOLVED_TOKEN and the intrinsic ID. The CORINFO_RESOLVED_TOKEN provides the Module, mdToken, and token kind and so the token/intrinsic ID would be added to the cache in the Module type. This would therefore not involve caching every token, but only a limited subset of intrinsic IDs that are explicitly "registered" to the cache when resolved..

In tier 0, it is likely there would likely be no entries in this cache and such entries wouldn't be used due to lack of inlining or other optimizations happening. However, methods that are rejitted to tier 1 would have had these self-same tokens already cached (from the tier 0 pass) if they were intrinsic IDs.

The inliner has to do a basic pass over the opcodes in a method it is determining inline profitability for and so it has the current Module, mdToken, and token kind. It can therefore construct the "input" CORINFO_RESOLVED_TOKEN and pass that into the cache. If an entry in the cache exists, it would return the cached RESOLVED_TOKEN and NamedIntrinsic and the inliner can do further heuristics as appropriate. If an entry does not exist in the cache, it would return NI_Illegal and the inliner would assume this is not a well known token and would continue as it does today. In the common "success" case, this will be NI_IsSupported_True, NI_IsSupported_False, or potentially the intrinsic ID for string.Length or Type.op_Equality.

The input does not strictly need to be a CORINFO_RESOLVED_TOKEN, this was simply a suggestion given that this contains roughly the information that is used to get eeCallInfo and ultimately the NamedIntrinsic in the normal importer scenario (where CEE_CALL is decoded).

[jkotas]

The intent was that lookupNamedIntrinsic would pass in the CORINFO_RESOLVED_TOKEN and the intrinsic ID. The CORINFO_RESOLVED_TOKEN provides the Module, mdToken, and token kind and so the token/intrinsic ID would be added to the cache in the Module type. This would therefore not involve caching every token, but only a limited subset of intrinsic IDs that are explicitly "registered" to the cache when resolved..

I am sorry but I do not see how this can work. I think you are saying that you can get from Module, mdToken to NamedIntrinsic id using a small cache and without doing expensive operation like resolving the token to HANDLE. Could you please walk through step-by-stepn on how it is going to work?

BTW: We have done experiments on catching resolveToken results in the past, but we were not able to come up with something that is a clear improvement. The problem is that resolveToken has a very long tail of tokens that are resolved just once or a few times that ends up killing the cache benefits.

[tannergooding]

Could you please walk through step-by-stepn on how it is going to work?

My understanding is that metadata tokens are effectively unique 32-bit unsigned integers per "module" (it is scoped to a specific "metadata binary"). This means that in the context of a single Module, the 4-byte token that follows CALL is unique to the module and any two methods in that module that have call 0x######## will be calling the same method. During normal importation this token is resolved by using the it as a lookup in the relevant metadata tables. For call this token will represent a methodref, methoddef, or methodspec entry.

This means that once something in the module has already resolved a call token to a well known NamedIntrinsic, you should effectively be able to use the same token ID in the same module to quickly get the same NamedIntrinsic without having to do the more in depth resolution. In pseudo-code, you would effectively have:

class Module
{
    private Dictionary<uint, NamedIntrinsic> _cachedIntrinsics;

    public static void RegisterWellKnownIntrinsic(uint token, NamedIntrinsic intrinsic) {
        if (_cachedIntrinsics is null) {
            _cachedIntrinsics = new Dictionary<uint, NamedIntrinsic>();
        }
        _cachedIntrinsics[token] = intrinsic;
    }

    public static bool TryGetWellKnownIntrinsic(uint token, out NamedIntrinsic intrinsic)
    {
        intrinsic = NI_Illegal;
        return (_cachedIntrinsic is not null) && _cachedIntrinsic.TryGetValue(token, out intrinsic);
    }
}

So, during tier 0 when an "interesting" named intrinsic is found, lookupNamedIntrinsic would call a JIT/EE interface method and say "for this module, register this token ID to be this named intrinsic", this is simple a map<mdToken, NamedIntrinsic> in the VM Module type (or similar).

During tier 1 when inlining is happening, the inliner would use a sibling JIT/EE interface method to say "is this token ID in this module a well known named intrinsic", that method effectively returns true with the corresponding NamedIntrinsic or false with NI_Illegal.