JIT: optimizations for multi-use boxes

[AndyAyersMS]

A fairly common pattern (especially after inlining) is to see a box that feeds an isinst and if that succeeds, an unbox.any. For example:

using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;

internal class ObjectEqualityComparer<T> : EqualityComparer<T>
{
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public override bool Equals(T x, T y)
    {
        if (x != null)
        {
            if (y != null) return x.Equals(y);
            return false;
        }
        if (y != null) return false;
        return true;
    }
    
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public override int GetHashCode(T obj) => obj?.GetHashCode() ?? 0;
    
    // Equals method for the comparer itself.
    public override bool Equals(Object obj) =>
    obj != null && GetType() == obj.GetType();
    
    public override int GetHashCode() =>
    GetType().GetHashCode();
}

class C
{
    public static int Main()
    {
        var comp = new ObjectEqualityComparer<int>();
        bool result = comp.Equals(3, 4);
        return result ? 0 : 100;
    }
}

We get pretty far when optimizing Main here -- we can devirtualize the call to Equals, inline it and remove the null checks since we have a value type, then inline the inner call to Equals. But along the way we have to box y and the inner Equals has the following IL:

IL_0000  03                ldarg.1     
IL_0001  75 f1 00 00 02    isinst       0x20000F1
IL_0006  2d 02             brtrue.s     2 (IL_000a)
IL_0008  16                ldc.i4.0    
IL_0009  2a                ret         
IL_000a  02                ldarg.0     
IL_000b  4a                ldind.i4    
IL_000c  03                ldarg.1     
IL_000d  a5 f1 00 00 02    unbox.any    0x20000F1
IL_0012  fe 01             ceq         
IL_0014  2a                ret         

With the advent of dotnet/coreclr#14420 the jit will now optimize away the isinst, but the box cleanup opts for unbox.any don't fire because there is usually a temp in the way, and so we generate the following code for Main:

G_M4930_IG01:
       57                   push     rdi
       56                   push     rsi
       4883EC28             sub      rsp, 40

G_M4930_IG02:
; ** BOX (y) **
       48B9086014E2FA7F0000 mov      rcx, 0x7FFAE2146008
       E8BB3C835F           call     CORINFO_HELP_NEWSFAST   
       C7400804000000       mov      dword ptr [rax+8], 4
       488BF0               mov      rsi, rax
       4885F6               test     rsi, rsi                 ; gratuitous null check ?
       7504                 jne      SHORT G_M4930_IG03
       33FF                 xor      edi, edi
       EB2D                 jmp      SHORT G_M4930_IG05

G_M4930_IG03:
; * UNBOX.ANY type check
       48BA086014E2FA7F0000 mov      rdx, 0x7FFAE2146008
       483916               cmp      qword ptr [rsi], rdx
       7412                 je       SHORT G_M4930_IG04
; * call helper if type check fails (which it won't)
       488BD6               mov      rdx, rsi
       48B9086014E2FA7F0000 mov      rcx, 0x7FFAE2146008
       E8470B395F           call     CORINFO_HELP_UNBOX

G_M4930_IG04:
       837E0803             cmp      dword ptr [rsi+8], 3
       400F94C7             sete     dil
       400FB6FF             movzx    rdi, dil

G_M4930_IG05:
       85FF                 test     edi, edi
       750C                 jne      SHORT G_M4930_IG07
       B864000000           mov      eax, 100

G_M4930_IG06:
       4883C428             add      rsp, 40
       5E                   pop      rsi
       5F                   pop      rdi
       C3                   ret

G_M4930_IG07:
       33C0                 xor      eax, eax

G_M4930_IG08:
       4883C428             add      rsp, 40
       5E                   pop      rsi
       5F                   pop      rdi
       C3                   ret

If when optimizing a successful cast we copy the result to a new more strongly typed temp (see #9117) we might be able to optimize away the type equality check in the downstream unbox.any (see dotnet/coreclr#14473). And perhaps if we are lucky and the box is simple we might be able to propagate the value to be boxed through the box/unbox to the ultimate use, and so not need the unbox. But the box would remain as it is difficult to remove unless it is known to be dead and whatever transformation makes it dead explicitly cleans it up.

A couple of ways we could approach this:

• The optimizer should be able to reason about and propagate boxes and perhaps trigger the box/unbox.any peephole, and turn the result into a simple copy.
• BOX is just an expression "wrapper" in the spirit of JIT: some ideas on high-level representation of runtime operations in IR #9056. So we could allow the inliner to give BOX(y) the same treatment as y and duplicate it within the inlinee body (essentially, generalize the logic in impInlineFetchArg that begins with else if (argInfo.argIsLclVar && !argCanBeModified) to also apply to BOX(y)). If we added suitable "reference counting" to boxes to track the duplicates then optimizing away the last use of the box could trigger the box cleanup. We have this today but the reference count is implicit and always = 1 since we don't duplicate the boxed values.

If all this kicked in, the code for Main above would collapse to simply returning a constant.

category:cq
theme:importer
skill-level:expert
cost:medium

[AndyAyersMS]

Considering for 2.1.

[mikedn]

4885F6 test rsi, rsi ; gratuitous null check ?

Yes. Can be removed by proper value numbering of GT_BOX. See dotnet/coreclr#15666.

I think that the subsequent unbox can be removed too - we need a proper VN for IND(JIT_New) and we should be able to remove CORINFO_HELP_UNBOX in optAssertionProp_Call.

But that still leaves us with the box:

G_M33790_IG01:
       4883EC28             sub      rsp, 40
G_M33790_IG02:
       48B968D0F2DEFC7F0000 mov      rcx, 0x7FFCDEF2D068
       E83DAF805F           call     CORINFO_HELP_NEWSFAST
       C7400804000000       mov      dword ptr [rax+8], 4
       83780803             cmp      dword ptr [rax+8], 3
       0F94C0               sete     al
       0FB6C0               movzx    rax, al
       85C0                 test     eax, eax
       750A                 jne      SHORT G_M33790_IG04
       B864000000           mov      eax, 100
G_M33790_IG03:
       4883C428             add      rsp, 40
       C3                   ret
G_M33790_IG04:
       33C0                 xor      eax, eax
G_M33790_IG05:
       4883C428             add      rsp, 40
       C3                   ret

It would be nice if the optimizer could handle this but I don't think there's any mechanism in the JIT that would allow the elimination of useless stores. So even with the fixes I mentioned above it would still be preferable to get rid of the box sooner.

Another related example:

struct RGB { public byte r, g, b; }
static int Main() => Test<RGB>() ? 42 : 0;
static bool Test<T>() => return default(T) == null;

Test generates:

G_M1515_IG01:
       57                   push     rdi
       56                   push     rsi
       4883EC28             sub      rsp, 40
G_M1515_IG02:
       33F6                 xor      esi, esi
       33FF                 xor      edi, edi
       48B9805AFC7FFC7F0000 mov      rcx, 0x7FFC7FFC5A80
       E827AF845F           call     CORINFO_HELP_NEWSFAST
       488D5008             lea      rdx, bword ptr [rax+8]
       408832               mov      byte  ptr [rdx], sil
       40887A01             mov      byte  ptr [rdx+1], dil
       40887A02             mov      byte  ptr [rdx+2], dil
       4885C0               test     rax, rax
       0F94C0               sete     al
       0FB6C0               movzx    rax, al
G_M1515_IG03:
       4883C428             add      rsp, 40
       5E                   pop      rsi
       5F                   pop      rdi
       C3                   ret

The null check can be eliminated but we're again stuck with the box. And this time it's completely useless, the code doesn't even attempt to read from the allocated memory.

It's not clear to me why the importer is so eager to expand box. It seems that it would be easier if the importer generates a single GT_BOX node that's then expanded during global morph.

[AndyAyersMS]

Expansion of some runtime constructs like box can cause inlining to fail. So this creates tension in the jit: expanding too early reveals ugly details that inhibit or complicates high-level optimizations; expanding too late means the jit can't easily undo inlines that are now invalid.

Hence the wrapper idea from #9056 where we try and have our cake and eat it too. We at least need to do enough sanity checking in importation that we can spot the cases where inlining will fail.

[AndyAyersMS]

Looked a bit at the reference counting idea. For instance genTreeBox nodes that share the same assign and copy could be linked in a circular list that is updated by gtClone and by deletion of box nodes.

Having IR cross links is obviously questionable but we already have these in genTreeBox and are potentially skating on thin ice because we don't have any ref counting.

At any rate, assuming we can somehow track how many uses there are, gtTryRemoveBoxUpstreamEffects should probably fail if there is more than one use. However, callers cannot tell why this method fails. Failing because the box is multi-use may still allow the optimization the caller has in mind to go forward, then allow them to remove their use of the box, and then potentially allow full cleanup to be triggered elsewhere. Hmm.

[AndyAyersMS]

In the RGB example the jit doesn't try optimizing the box early as it is in a return tree. The importer currently only calls gtFoldExpr on trees feeding conditional branches.

We could presumably invoke the expression folder more often during importation and catch more cases. I suppose the question then is "how much more often" -- likely whether the size reduction wins counteract the extra cost of searching for foldable patterns. I'll play around with this and see how likely more eager folding is to lead to changes.

We also try to fold the box tree again during morph but the struct assignment to copy the box payload has been expanded into an exception-flagged comma laden tree that the box cleanup utility does not expect, so it bails out.

Attempting to optimize BOX(valueType) EQ null [000021]
gtTryRemoveBoxUpstreamEffects: attempting to remove side effects of BOX (valuetype) [000019] (assign/newobj [000011] copy [000017])

In this case the exception flags on the morphed copy tree are bogus as we're copying from a local struct to the box payload. The copy tree doesn't have an exception flag initially but fgMorphCopyBlock throws one in for some reason.

[AndyAyersMS]

Looked at the flags issue a bit without gaining much insight. The GTF_EXCEPT flag is flipped on for GT_BLK during morph regardless of whether it was set before. So any upstream knowledge about potentially safe memory operations is seemingly just tossed out. And if we don't set the flag then fgDebugCheckFlags gets cranky.

[mikedn]

In the RGB example the jit doesn't try optimizing the box early as it is in a return tree. The importer currently only calls gtFoldExpr on trees feeding conditional branches.

Aha! So the trick I used for string.IsNullOrEmpty works here as well but for completely different reasons:

static bool Test<T>() => default(T) == null ? true : false;

generates

G_M8956_IG02:
       33C0                 xor      eax, eax
G_M8956_IG03:
       C3                   ret

Funny.

[mikedn]

Or in this List<T> code where I extracted the example from:
https://github.com/dotnet/coreclr/blob/8a9c59260b8cc04d266e021229f4b044bf40ad52/src/mscorlib/shared/System/Collections/Generic/List.cs#L193-L198

static bool Test<T>(object value) => (value is T) || (value == null && (default(T) == null ? true : false));

generates

G_M8956_IG02:
       488BC1               mov      rax, rcx
       4885C0               test     rax, rax
       7411                 je       SHORT G_M8956_IG03
       48BA585CC5F0FE7F0000 mov      rdx, 0x7FFEF0C55C58
       483910               cmp      qword ptr [rax], rdx
       7402                 je       SHORT G_M8956_IG03
       33C0                 xor      rax, rax
G_M8956_IG03:
       4885C0               test     rax, rax
       750B                 jne      SHORT G_M8956_IG07
       4885C9               test     rcx, rcx
       7503                 jne      SHORT G_M8956_IG05
       33C0                 xor      eax, eax
G_M8956_IG04:
       C3                   ret
G_M8956_IG05:
       33C0                 xor      eax, eax
G_M8956_IG06:
       C3                   ret
G_M8956_IG07:
       B801000000           mov      eax, 1
G_M8956_IG08:
       C3                   ret

While not ideal at least it no longer does boxing.

[AndyAyersMS]

Seems like importer folding at relational ops is pretty cheap and effective. Will send a PR.

[AndyAyersMS]

Think the more general opt needs more conceptual bake time, so moving out of 2.1.

[AndyAyersMS]

Though I like this proposal I'm still not quite sure how to handle the case where there are multiple uses and the uses have different demands on the box tree. Say for some uses are type tests (and so just want the type handle) and some others are feeding an unboxed entry point (and so just want a local copy of the payload) and some others feed a null test.

Say we refcount and there are N uses. Until we understand all those use contexts we can't decide what shape the box tree should take on. So in addition to the reference count we need to track the kinds of references (basically the BR_xxx values that the uses would request if they were to optimize). And if any context needs the full box, then we can only satisfy the non-destructive uses (type tests and null tests).

Also since optimizations like the Enum.HasFlag couple the fates of two boxes together, box optimizations can potentially be mutually dependent. Though perhaps we can ignore this dependence as the current cases like this are single-use and so can just be handled directly.

[zpodlovics]

@AndyAyersMS Is it possible to also recognize + optimize box/unbox sequence? It seems this IL sequence is common in F#.

dotnet/fsharp#5019 (comment)

[AndyAyersMS]

Adjacent box / unbox.any is handled via a peephole optimization in the jit importer (look at importer.cpp around line 14890).

If there is IL in between or a dup of the box the jit may not optimize.

[AndyAyersMS]

Looks like a longshot for .NET 9 now too, so moving to future.

[AndyAyersMS]

Fixed by #103361