Segregate merged returns by constant value

[JosephTremoulet]

The JIT enforces a limit on the number of epilogs emitted in any given
method. Change the logic so that when this merging kicks in, returns of
constant values are given merged return blocks distinct from each other
and from the general return block, as long as we can do so without going
over the limit. This particularly helps avoid redundancy (while still
keeping method size down) in methods with a large number of constant
returns but only a few distinct constants, which is true of many
predicate methods with bool return type.

This is the compiler portion of #13466 and dotnet/corefx#23395.

[JosephTremoulet]

@AndyAyersMS, @briansull PTAL
/cc @dotnet/jit-contrib

jit-diff results (--frameworks --tests, Windows x64 release)

Summary:
(Note: Lower is better)
Total bytes of diff: -25461 (-0.02 % of base)
    diff is an improvement.
Total byte diff includes 0 bytes from reconciling methods
        Base had    0 unique methods,        0 unique bytes
        Diff had    0 unique methods,        0 unique bytes
Top file regressions by size (bytes):
        1997 : System.Private.CoreLib.dasm (0.06 % of base)
         325 : System.Net.Http.dasm (0.10 % of base)
         166 : System.Collections.dasm (0.14 % of base)
         100 : JIT\Methodical\Overflow\FloatInfinitiesToInt_do\FloatInfinitiesToInt_do.dasm (2.50 % of base)
         100 : JIT\Methodical\Overflow\FloatInfinitiesToInt_ro\FloatInfinitiesToInt_ro.dasm (2.50 % of base)
Top file improvements by size (bytes):
      -15115 : JIT\opt\virtualstubdispatch\bigvtbl\bigvtbl_cs_do\bigvtbl_cs_do.dasm (-4.19 % of base)
      -15115 : JIT\opt\virtualstubdispatch\bigvtbl\bigvtbl_cs_ro\bigvtbl_cs_ro.dasm (-4.19 % of base)
       -1302 : Microsoft.CodeAnalysis.CSharp.dasm (-0.06 % of base)
        -916 : JIT\CodeGenBringUpTests\struct16args\struct16args.dasm (-14.01 % of base)
        -822 : Microsoft.CodeAnalysis.dasm (-0.10 % of base)
370 total files with size differences (87 improved, 283 regressed), 6174 unchanged.
Top method regessions by size (bytes):
        1253 : System.Private.CoreLib.dasm - Comparer`1:System.Collections.IComparer.Compare(ref,ref):int:this (39 methods)
         225 : System.Private.CoreLib.dasm - ValueTuple`8:GetHashCode():int:this (15 methods)
         225 : System.Private.CoreLib.dasm - ValueTuple`8:GetHashCodeCore(ref):int:this (15 methods)
         208 : System.Private.CoreLib.dasm - ValueTuple`8:System.Collections.IStructuralComparable.CompareTo(ref,ref):int:this (15 methods)
         131 : System.Private.CoreLib.dasm - SignatureTypeExtensions:TryResolve(ref,ref):ref
Top method improvements by size (bytes):
      -15115 : JIT\opt\virtualstubdispatch\bigvtbl\bigvtbl_cs_do\bigvtbl_cs_do.dasm - CTest:Main():int
      -15115 : JIT\opt\virtualstubdispatch\bigvtbl\bigvtbl_cs_ro\bigvtbl_cs_ro.dasm - CTest:Main():int
        -335 : Microsoft.CodeAnalysis.dasm - ILOpCodeExtensions:StackPopCount(char):int
        -231 : Microsoft.CodeAnalysis.dasm - ILOpCodeExtensions:StackPushCount(char):int
        -225 : Microsoft.CodeAnalysis.dasm - OneToOneUnicodeComparer:Equals(ref,ref):bool:this
1507 total methods with size differences (533 improved, 974 regressed), 313995 unchanged.

[AndyAyersMS]

I'll write up some more detailed comments later, but here are some initial thoughts:

• Feels like there ought to be some kind of profitability metric. Not sure what it is, but we are trading off a potential extra taken branch vs a duplicated epilog, so seems like it should determine when there is really a branch cost to be paid and have some idea of epilog cost. Though in all likelihood it is too early to know the latter.
• Apropos of the above (and assuming I actually understand the change), it seems odd to only try merging when there are too many epilogs, rather than to always try merging.
• The merging is done greedily, so the code is not guaranteed to maximize the effectiveness of merging from a static standpoint (say a method has 100 returns, and the first 4 are distinct, and the last 96 all return the same value). Might be interesting to get a histogram of the static return distribution in methods; if methods with > 4 return opcodes are super rare then the greedy strategy is probably just fine.
• I am tempted to say that you should immediately generalize this to handle simple leaf expressions, but don't know if that is common enough to matter. But having a GenTree* in your table would give you a natural null case.
• I am curious how this impacts the inliner/optimizer ability to handle cases where these methods are inlined into predicate contexts, eg if (foo(...)). Seems like it might make it easier for a simple post-inline peephole to see a join/fork block and decide to duplicate it to allow constant prop to take out the if.

[mikedn]

Feels like there ought to be some kind of profitability metric. Not sure what it is, but we are trading off a potential extra taken branch vs a duplicated epilog, so seems like it should determine when there is really a branch cost to be paid and have some idea of epilog cost. Though in all likelihood it is too early to know the latter.

Wouldn't be better to merge all return blocks and then duplicate blocks as needed somewhere after lowering/lsra where information about the epilog size is available? I've seen cases where the epilog is pretty large and the cost of the jmp saved is insignificant (cold code paths).

[JosephTremoulet]

I should have mentioned that this isn't an attempt at optimal epilog duplication/merging, just an incremental step that improves on the current state and in particular is sufficient to get the motivating cases from bullets numbered "2" in https://github.com/dotnet/coreclr/issues/13466#issuecomment-323831444 and https://github.com/dotnet/corefx/issues/23395#issuecomment-324364657. Diffs tend to look like this:

diff --git "a/.\\before\\IsSubsetOf.dasm" "b/.\\after\\IsSubsetOf.dasm"
index 24b3c4c..3f1ec29 100644
--- "a/.\\before\\IsSubsetOf.dasm"
+++ "b/.\\after\\IsSubsetOf.dasm"
@@ -5,96 +5,104 @@
 ; partially interruptible
 ; Final local variable assignments
 ;
-;  V00 this         [V00,T01] (  9,  6   )     ref  ->  rsi         this class-hnd
-;  V01 arg1         [V01,T02] (  5,  4   )     ref  ->  rdi         class-hnd
-;  V02 loc0         [V02,T03] (  6,  3   )     ref  ->  rbp         class-hnd
-;  V03 loc1         [V03    ] (  3,  1.50)  struct ( 8) [rsp+0x28]   do-not-enreg[SF] must-init
-;  V04 tmp0         [V04,T04] (  3,  3   )     int  ->  rbx        
+;  V00 this         [V00,T00] (  9,  6   )     ref  ->  rsi         this class-hnd
+;  V01 arg1         [V01,T01] (  5,  4   )     ref  ->  rdi         class-hnd
+;  V02 loc0         [V02,T03] (  6,  3   )     ref  ->  rbx         class-hnd
+;  V03 loc1         [V03    ] (  3,  1.50)  struct ( 8) [rsp+0x20]   do-not-enreg[SF] must-init
+;  V04 tmp0         [V04,T04] (  3,  3   )     int  ->  rdi        
 ;  V05 tmp1         [V05,T06] (  3,  0   )     ref  ->  rdi         class-hnd exact
-;  V06 tmp2         [V06,T00] (  6,  7   )     int  ->  rbx        
-;  V07 tmp3         [V07    ] (  2,  1   )     int  ->  [rsp+0x28]   do-not-enreg[] V03.UniqueCount(offs=0x00) P-DEP
-;  V08 tmp4         [V08    ] (  2,  1   )     int  ->  [rsp+0x2C]   do-not-enreg[] V03.UnfoundCount(offs=0x04) P-DEP
+;  V06 tmp2         [V06,T02] (  3,  3   )     int  ->  rbx        
+;  V07 tmp3         [V07    ] (  2,  1   )     int  ->  [rsp+0x20]   do-not-enreg[] V03.UniqueCount(offs=0x00) P-DEP
+;  V08 tmp4         [V08    ] (  2,  1   )     int  ->  [rsp+0x24]   do-not-enreg[] V03.UnfoundCount(offs=0x04) P-DEP
 ;  V09 tmp5         [V09,T07] (  2,  0   )     ref  ->  rdx        
 ;  V10 tmp6         [V10,T05] (  2,  2   )    long  ->  rcx        
 ;  V11 OutArgs      [V11    ] (  1,  1   )  lclBlk (32) [rsp+0x00]  
 ;
-; Lcl frame size = 56
+; Lcl frame size = 48
 G_M12846_IG01:
        push     rdi
        push     rsi
-       push     rbp
        push     rbx
-       sub      rsp, 56
+       sub      rsp, 48
        xor      rax, rax
-       mov      qword ptr [rsp+28H], rax
-       mov      qword ptr [rsp+30H], rcx
+       mov      qword ptr [rsp+20H], rax
+       mov      qword ptr [rsp+28H], rcx
        mov      rsi, rcx
        mov      rdi, rdx
 G_M12846_IG02:
        test     rdi, rdi
-       je       G_M12846_IG10
+       je       G_M12846_IG13
 G_M12846_IG03:
        mov      rcx, rsi
        call     [SortedSet`1:get_Count():int:this]
        test     eax, eax
-       jne      SHORT G_M12846_IG04
-       mov      ebx, 1
-       jmp      G_M12846_IG08
+       je       G_M12846_IG07
 G_M12846_IG04:
        mov      rcx, qword ptr [rsi]
        call     [CORINFO_HELP_READYTORUN_GENERIC_HANDLE]
        mov      rcx, rax
        mov      rdx, rdi
        call     [CORINFO_HELP_ISINSTANCEOFANY]
-       mov      rbp, rax
-       test     rbp, rbp
+       mov      rbx, rax
+       test     rbx, rbx
        je       SHORT G_M12846_IG06
        mov      rcx, rsi
-       mov      rdx, rbp
+       mov      rdx, rbx
        call     [SortedSet`1:HasEqualComparer(ref):bool:this]
        test     al, al
        je       SHORT G_M12846_IG06
        mov      rcx, rsi
        call     [SortedSet`1:get_Count():int:this]
-       mov      ebx, eax
-       mov      rcx, rbp
+       mov      edi, eax
+       mov      rcx, rbx
        call     [SortedSet`1:get_Count():int:this]
-       cmp      eax, ebx
-       jge      SHORT G_M12846_IG05
-       xor      ebx, ebx
-       jmp      SHORT G_M12846_IG08
+       cmp      eax, edi
+       jl       SHORT G_M12846_IG09
 G_M12846_IG05:
        mov      rcx, rsi
-       mov      rdx, rbp
+       mov      rdx, rbx
        call     [SortedSet`1:IsSubsetOfSortedSetWithSameComparer(ref):bool:this]
        movzx    rbx, al
-       jmp      SHORT G_M12846_IG08
+       jmp      SHORT G_M12846_IG11
 G_M12846_IG06:
        mov      rcx, rsi
        mov      rdx, rdi
        xor      r8d, r8d
        call     [SortedSet`1:CheckUniqueAndUnfoundElements(ref,bool):struct:this]
-       mov      qword ptr [rsp+28H], rax
+       mov      qword ptr [rsp+20H], rax
        mov      rcx, rsi
        call     [SortedSet`1:get_Count():int:this]
-       cmp      eax, dword ptr [rsp+28H]
-       jne      SHORT G_M12846_IG07
-       cmp      dword ptr [rsp+2CH], 0
+       cmp      eax, dword ptr [rsp+20H]
+       jne      SHORT G_M12846_IG09
+       cmp      dword ptr [rsp+24H], 0
        setge    bl
        movzx    rbx, bl
-       jmp      SHORT G_M12846_IG08
+       jmp      SHORT G_M12846_IG11
 G_M12846_IG07:
-       xor      ebx, ebx
+       mov      eax, 1
 G_M12846_IG08:
-       movzx    rax, bl
-G_M12846_IG09:
-       add      rsp, 56
+       add      rsp, 48
        pop      rbx
-       pop      rbp
        pop      rsi
        pop      rdi
        ret      
+G_M12846_IG09:
+       xor      eax, eax
 G_M12846_IG10:
+       add      rsp, 48
+       pop      rbx
+       pop      rsi
+       pop      rdi
+       ret      
+G_M12846_IG11:
+       movzx    rax, bl
+G_M12846_IG12:
+       add      rsp, 48
+       pop      rbx
+       pop      rsi
+       pop      rdi
+       ret      
+G_M12846_IG13:
        call     [CORINFO_HELP_READYTORUN_NEW]
        mov      rdi, rax
        mov      ecx, 0x849
@@ -105,5 +113,5 @@ G_M12846_IG10:
        mov      rcx, rdi
        call     CORINFO_HELP_THROW
        int3     
-; Total bytes of code 250, prolog size 20 for method SortedSet`1:IsSubsetOf(ref):bool:this
+; Total bytes of code 259, prolog size 19 for method SortedSet`1:IsSubsetOf(ref):bool:this

We're not too graceful with the handling of the temp generated for the common merge case, especially for bool methods, so as you can see the dynamic path here replaces (not taken) jcc + mov reg, imm + jmp + movzx reg, reg + epilog with (taken) jcc + mov reg, imm + epilog, and this is a fairly common pattern in the diffs.

In light of that, regarding some of the specific points:

it seems odd to only try merging when there are too many epilogs, rather than to always try merging

I agree, and likewise that it's odd to have the arbitrary limit of 4 applied to all targets; my thinking was just to untangle some of the cases that are currently hurting us from merging, and to leave this existing oddity alone (for now).

Feels like there ought to be some kind of profitability metric. Not sure what it is

My thinking was that the current metric is effectively "5 is too many, fewer is ok", and that I'm effectively preserving that. Having a smarter metric (e.g. that includes epilog size as @mikedn suggests) certainly sounds desirable, but I'm not sure what data points are available to inform it without making this a significantly larger change...

The merging is done greedily
I am tempted to say that you should immediately generalize this to handle simple leaf expressions

Both of these are getting at TP/CQ tradeoff, and yes what I've done is biased toward TP in that sense. Should be straightforward enough to compute that histogram you suggest to get a better sense...

[AndyAyersMS]

Note: if you are tempted to go to a two-pass approach to try and prioritize merge candidates when there are lots of returns -- we've already counted up returns during the IL prescan (or at least done the equivalent), so you could presumably avoid being clever about merging if there are 4 or fewer return points. That probably amortizes the TP cost for the cases where there is no benefit sufficiently that you can spend more TP on getting the best results for cases where there is a large benefit.

[briansull]

If this runs after the ValueNumbering phase I would suggest using the ValueNumber as the key instead of a constant.
That way a Select function could return argA or argB in several places in the method and those returns could be combined as with constants.

[JosephTremoulet]

This runs in fgAddInternal, well before value numbering.

[JosephTremoulet]

Histograms as promised:

--frameworks:

# distinct constants	# methods	cumulative % methods
1	369	39.26%
2	379	79.57%
3	96	89.79%
4	14	91.28%
5	14	92.77%
6	18	94.68%
7	3	95.00%
8	10	96.06%
9	3	96.38%
10	1	96.49%
11	1	96.60%
12	2	96.81%
13	2	97.02%
14	6	97.66%
15	4	98.09%
16	4	98.51%
17	3	98.83%
18	4	99.26%
22	3	99.57%
29	1	99.68%
34	1	99.79%
36	1	99.89%
82	1	100.00%

--benchmarks:

# distinct constants	# methods	cumulative % methods
1	201	29.30%
2	235	63.56%
3	3	63.99%
4	2	64.29%
5	85	76.68%
6	32	81.34%
7	8	82.51%
57	6	83.38%
58	8	84.55%
68	8	85.71%
78	16	88.05%
81	8	89.21%
93	16	91.55%
111	16	93.88%
284	42	100.00%

--frameworks --tests:

# distinct constants	# methods	cumulative % methods
1	570	35.06%
2	614	72.82%
3	99	78.91%
4	16	79.89%
5	99	85.98%
6	50	89.05%
7	11	89.73%
8	10	90.34%
9	3	90.53%
10	1	90.59%
11	1	90.65%
12	2	90.77%
13	2	90.90%
14	6	91.27%
15	4	91.51%
16	4	91.76%
17	3	91.94%
18	4	92.19%
22	3	92.37%
29	1	92.44%
34	1	92.50%
36	1	92.56%
57	6	92.93%
58	8	93.42%
68	8	93.91%
78	16	94.90%
81	8	95.39%
82	1	95.45%
93	16	96.43%
111	16	97.42%
284	42	100.00%

Data points are the number of distinct return constants, limited to the set of methods where discretionary return merging kicked in (more than four return sites, and none of the sync flag etc. that force us to merge down to a single return) and at least one return site returned a constant.

[JosephTremoulet]

Since that showed "5 constants" and "6 constants" as potentially interesting, I decided to see what the distributions looked like in that subset, to get an idea of how skewed they are to one or another of the returned values (the idea being that if it's flat, we're faced with an arbitrary choice anyway, assuming we're keeping our limit of 4 epilogs for the time being). Here's what that looked like:

--frameworks:

histogram of constants	# methods	cumulative % methods
(1, 1, 1, 1, 1, 0)	7	21.88%
(1, 1, 1, 1, 1, 1)	9	50.00%
(2, 1, 1, 1, 1, 0)	3	59.38%
(2, 1, 1, 1, 1, 1)	5	75.00%
(2, 2, 1, 1, 1, 0)	2	81.25%
(2, 2, 2, 1, 1, 1)	1	84.38%
(2, 2, 2, 2, 2, 0)	1	87.50%
(3, 2, 2, 2, 1, 1)	1	90.63%
(3, 3, 3, 3, 1, 1)	1	93.75%
(5, 1, 1, 1, 1, 1)	1	96.88%
(20, 19, 11, 3, 3, 0)	1	100.00%

--benchmarks:

histogram of constants	# methods	cumulative % methods
(1, 1, 1, 1, 1, 0)	68	58.12%
(1, 1, 1, 1, 1, 1)	32	85.47%
(2, 1, 1, 1, 1, 0)	1	86.32%
(3, 3, 3, 3, 1, 0)	16	100.00%

--frameworks --tests:

histogram of constants	# methods	cumulative % methods
(1, 1, 1, 1, 1, 0)	75	50.34%
(1, 1, 1, 1, 1, 1)	41	77.85%
(2, 1, 1, 1, 1, 0)	4	80.54%
(2, 1, 1, 1, 1, 1)	5	83.89%
(2, 2, 1, 1, 1, 0)	2	85.23%
(2, 2, 2, 1, 1, 1)	1	85.91%
(2, 2, 2, 2, 2, 0)	1	86.58%
(3, 2, 2, 2, 1, 1)	1	87.25%
(3, 3, 3, 3, 1, 0)	16	97.99%
(3, 3, 3, 3, 1, 1)	1	98.66%
(5, 1, 1, 1, 1, 1)	1	99.33%
(20, 19, 11, 3, 3, 0)	1	100.00%

[AndyAyersMS]

This is from a "universe" of 339K or so methods, right? So roughly 0.2% of all methods in this set get into the merging logic, which seems plausible, and only some subset of those could potentially be further improved.

I suppose we could spin this different ways:

• these cases are rare enough that it's not worth complicating things further for them
• these cases are rare enough that we can justify extra TP getting better results for them

[JosephTremoulet]

I've filed #13814 to track the opportunity to implement a more robust epilog generation strategy; @AndyAyersMS, @mikedn, @briansull, I think it captures the points you've all made about that, please update if not.

I still think there's merit in this change, since it is a fairly small modification of the current algorithm with minimal compile-time overhead, caps codesize growth at 4 epilogs in the small percentage of functions with constant returns, and (most importantly) improves codegen enough that people can stop writing goto returnFalse-type workarounds like we see here, here, here, and here.

I am curious how this impacts the inliner/optimizer ability to handle cases where these methods are inlined

This is part of fgAddInternal, which runs in global morph just after fgInline, so I don't believe it has any effect on inlinee IR; also, since its function is to merge returns (and this change just makes that less aggressive), the inlinee IR should already have "separate" returns (that then join together at the merge point the inliner creates for returning to the caller). Separating out the inliner-created merge point for these methods could certainly be interesting, but I think would have a home in the inliner or a general tail duplication optimization.

I suppose we could spin this different ways:

• these cases are rare enough that it's not worth complicating things further for them
• these cases are rare enough that we can justify extra TP getting better results for them

My inclination is the former -- the goal here, quoting https://github.com/dotnet/coreclr/issues/13466#issuecomment-323831444, is "a very narrow pattern match [that] could cheaply identify and help with search-loop predicate methods, which the motivating cases here have been", and I think it makes more sense to tackle #13814 than to embellish this much beyond what's here.

[briansull]

I'm also in the former camp:

• these cases are rare enough that it's not worth complicating things further for them

[AndyAyersMS]

Logic looks good, just some notes on comments.

[AndyAyersMS]

Maybe add a comment here about what it means when returnBlock is nullptr.

Also should this method (and others in the class) have the "standard" jit method header comments?

[JosephTremoulet]

Sure; done.

[AndyAyersMS]

I'm probably missing something here, but the comment seems stale.

If we found a constant return for this return point then we did rewrite the return. Maybe the first part should say "we didn't rewrite the return for cases where the return will use the common return block"?

And the second part, seems like merge always "succeeds" unless it is in the eager create mode and so decrementing is something that is always done, whether or not rewriting happened.

[JosephTremoulet]

I mentioned "for the non-const cases" in the first clause, but yeah the way the method flows now there's not a convenient place for a codepath representing that distinction, so I've moved the comments about this up to the method headers.

[AndyAyersMS]

Maybe it's a parsing ambiguity -- is a "non-const" case describing the input or the output? I've been thinking it applied to the input....

[JosephTremoulet]

Updated to incorporate @AndyAyersMS's feedback and to fix the issue that popped up in CI where I'm not creating the return statement for genReturnBB before rather than after the other code in fgAddInternal, so need to use fgInsertStatementNearEnd rather than fgInsertStatementAtEnd in a couple places there.

[AndyAyersMS]

Thanks for the updates.

[JosephTremoulet]

Pushed one more update: pulled the new class definition out of the enclosing method, to help the method's readability.