JIT generates suboptimal assembly code when using pattern matching

[h3xds1nz]

JIT generates a suboptimal assembly code when I replace plain if statements with pattern-matching. Now I do understand it generates a slightly different IL code, however, when running with optimizations, the generated assembly in the second case is just wasteful.

Do note that if you change the backing enum type to (U)Int64, everything is perfect in both cases. However, as we know, the standard backing type is Int32, and rarely anyone ever modifies that.

Code

C# code

public static unsafe class Extensions
{
    [JitGeneric(typeof(TestClass.TEST_ENUM))]
    public static bool HasThisFlagIFStatements<T>(this T @enum, T flag) where T : unmanaged, Enum
    {
        if (sizeof(T) == sizeof(Byte))
            return (Unsafe.As<T, Byte>(ref @enum) & Unsafe.As<T, Byte>(ref flag)) == Unsafe.As<T, Byte>(ref flag);

        if (sizeof(T) == sizeof(UInt16))
            return (Unsafe.As<T, UInt16>(ref @enum) & Unsafe.As<T, UInt16>(ref flag)) == Unsafe.As<T, UInt16>(ref flag);

        if (sizeof(T) == sizeof(UInt32))
            return (Unsafe.As<T, UInt32>(ref @enum) & Unsafe.As<T, UInt32>(ref flag)) == Unsafe.As<T, UInt32>(ref flag);

        if (sizeof(T) == sizeof(UInt64))
            return (Unsafe.As<T, UInt64>(ref @enum) & Unsafe.As<T, UInt64>(ref flag)) == Unsafe.As<T, UInt64>(ref flag);

        throw new InvalidOperationException("Target enum is not backed by Byte/UInt16/UInt32/UInt64");
    }

    [JitGeneric(typeof(TestClass.TEST_ENUM))]
    public static bool HasThisFlagPatternMatching<T>(this T @enum, T flags) where T : unmanaged, Enum => sizeof(T) switch
    {
        sizeof(Byte) => (Unsafe.As<T, Byte>(ref @enum) & Unsafe.As<T, Byte>(ref flags)) == Unsafe.As<T, Byte>(ref flags),
        sizeof(UInt16) => (Unsafe.As<T, UInt16>(ref @enum) & Unsafe.As<T, UInt16>(ref flags)) == Unsafe.As<T, UInt16>(ref flags),
        sizeof(UInt32) => (Unsafe.As<T, UInt32>(ref @enum) & Unsafe.As<T, UInt32>(ref flags)) == Unsafe.As<T, UInt32>(ref flags),
        sizeof(UInt64) => (Unsafe.As<T, UInt64>(ref @enum) & Unsafe.As<T, UInt64>(ref flags)) == Unsafe.As<T, UInt64>(ref flags),
        _ => throw new InvalidOperationException("Target enum is not backed by Byte/UInt16/UInt32/UInt64")
    };
}

public class TestClass {     
    [Flags]
    public enum TEST_ENUM : UInt32
    {
        BIT_1 = 0x01,
        BIT_2 = 0x02,
        BIT_3 = 0x04
    } 
}

Generated assembly when using Byte/UInt16/UInt32 as backing for enum (same when in-lined into a different method)

Extensions.HasThisFlagIFStatements[[TestClass+TEST_ENUM, _]](TEST_ENUM, TEST_ENUM)
    L0000: and ecx, edx
    L0002: xor eax, eax
    L0004: cmp ecx, edx
    L0006: sete al
    L0009: ret

Extensions.HasThisFlagPatternMatching[[TestClass+TEST_ENUM, _]](TEST_ENUM, TEST_ENUM)
    L0000: mov [rsp+8], ecx
    L0004: mov [rsp+0x10], edx
    L0008: mov eax, [rsp+8]
    L000c: and eax, [rsp+0x10]
    L0010: cmp eax, [rsp+0x10]
    L0014: sete al
    L0017: movzx eax, al
    L001a: ret

Generated assembly when using UInt64 as backing for enum (expected)

Extensions.HasThisFlagIFStatements[[TestClass+TEST_ENUM, _]](TEST_ENUM, TEST_ENUM)
    L0000: and rcx, rdx
    L0003: xor eax, eax
    L0005: cmp rcx, rdx
    L0008: sete al
    L000b: ret

Extensions.HasThisFlagPatternMatching[[TestClass+TEST_ENUM, _]](TEST_ENUM, TEST_ENUM)
    L0000: and rcx, rdx
    L0003: xor eax, eax
    L0005: cmp rcx, rdx
    L0008: sete al
    L000b: ret

Configuration

.NET 8.0.5 x64 Win10

SharpLab

SharpLab demo

[dotnet-policy-service[bot]]

Tagging subscribers to this area: @JulieLeeMSFT, @jakobbotsch See info in area-owners.md if you want to be subscribed.

[huoyaoyuan]

Replacing Unsafe.As with Unsafe.BitCast will result in expected codegen:

    [JitGeneric(typeof(TestClass.TEST_ENUM))]
    public static bool HasThisFlagPatternMatching<T>(this T @enum, T flags) where T : unmanaged, Enum => sizeof(T) switch
    {
        sizeof(Byte) => (Unsafe.BitCast<T, Byte>(@enum) & Unsafe.BitCast<T, Byte>(flags)) == Unsafe.BitCast<T, Byte>(flags),
        sizeof(UInt16) => (Unsafe.BitCast<T, UInt16>(@enum) & Unsafe.BitCast<T, UInt16>(flags)) == Unsafe.BitCast<T, UInt16>(flags),
        sizeof(UInt32) => (Unsafe.BitCast<T, UInt32>(@enum) & Unsafe.BitCast<T, UInt32>(flags)) == Unsafe.BitCast<T, UInt32>(flags),
        sizeof(UInt64) => (Unsafe.BitCast<T, UInt64>(@enum) & Unsafe.BitCast<T, UInt64>(flags)) == Unsafe.BitCast<T, UInt64>(flags),
        _ => throw new InvalidOperationException("Target enum is not backed by Byte/UInt16/UInt32/UInt64")
    };

Extensions.HasThisFlagPatternMatching[[TestClass+TEST_ENUM, _]](TEST_ENUM, TEST_ENUM)
    L0000: and ecx, edx
    L0002: xor eax, eax
    L0004: cmp ecx, edx
    L0006: sete al
    L0009: ret

This is the particular reason for creating Unsafe.BitCast, because Unsafe.As will take the address of variables and make it unfriendly for optimization.

Also note that in recent versions of .NET, Enum.HasFlag has been optimized to not box:

    [JitGeneric(typeof(TestClass.TEST_ENUM))]
    public static bool HasFlag<T>(this T @enum, T flag) where T : unmanaged, Enum => @enum.HasFlag(flag);

Extensions.HasFlag[[TestClass+TEST_ENUM, _]](TEST_ENUM, TEST_ENUM)
    L0000: and ecx, edx
    L0002: xor eax, eax
    L0004: cmp ecx, edx
    L0006: sete al
    L0009: ret

[h3xds1nz]

@huoyaoyuan I actually had no idea about Unsafe.BitCast, completely missed this new API in .NET 8. - Thank you for bringing that one to my attention.

I'm aware that Enum.HasFlag has been rewritten but the reason for my endeavours were actually different bit operations, I've just used this particular compare for demonstration purposes.

Nonetheless, I believe/hope the codegen from switch pattern(s) might still be optimized in this case.

[huoyaoyuan]

It's not specific to switch expression-pattern. Unsafe.As can result in suboptimal codegen under many conditions. Unsafe.BitCast is the suggested approach, especially for types smaller than pointer size.

[jakobbotsch]

The problem is that the JIT doesn't manage to fold the switch expression away early enough. This means we end up with Unsafe.As<uint, ulong>(ref @enum) under an unreachable branch. This forces address exposure because it accesses beyond the size of @enum. The switch gets folded away later, but at that point the damage of address exposure has already occurred.

On the other hand, with BitCast<uint, ulong>(x) we just give up on importing it and leave it as a call under the unreachable branch. The call just takes the original value without address exposing it, and the switch then gets removed later.

We would need some early additional form of propagation to address a limitation like this. I would be partial to having some form of propagation pass happen after inlining since that would also help with some other limitations we've had around delegate/function pointer inlining.

Going to put this in future since I don't see a simple fix and there is a straightforward workaround using BitCast.

[h3xds1nz]

Thank you for the thorough explanation. That has shed light for me on why does it happen :)

Yeah, I'm fine with there being a workaround for this case in .NET 8.