Fix too wide memory accesses for nullchecks (#64683)

src/coreclr/jit/codegenarmarch.cpp

@@ -1508,7 +1508,7 @@ void CodeGen::genCodeForNullCheck(GenTreeIndir* tree)
     genConsumeRegs(op1);
     regNumber targetReg = REG_ZR;
 
-    GetEmitter()->emitInsLoadStoreOp(INS_ldr, EA_4BYTE, targetReg, tree);
+    GetEmitter()->emitInsLoadStoreOp(ins_Load(tree->TypeGet()), emitActualTypeSize(tree), targetReg, tree);
 #endif
 }
 

src/coreclr/jit/codegenxarch.cpp

@@ -3967,7 +3967,7 @@ void CodeGen::genCodeForNullCheck(GenTreeIndir* tree)
 
     assert(tree->gtOp1->isUsedFromReg());
     regNumber reg = genConsumeReg(tree->gtOp1);
-    GetEmitter()->emitIns_AR_R(INS_cmp, EA_4BYTE, reg, reg, 0);
+    GetEmitter()->emitIns_AR_R(INS_cmp, emitTypeSize(tree), reg, reg, 0);
 }
 
 //------------------------------------------------------------------------

src/coreclr/jit/compiler.cpp

@@ -9736,6 +9736,29 @@ bool Compiler::lvaIsOSRLocal(unsigned varNum)
     return false;
 }
 
+//------------------------------------------------------------------------------
+// gtTypeForNullCheck: helper to get the most optimal and correct type for nullcheck
+//
+// Arguments:
+//    tree - the node for nullcheck;
+//
+var_types Compiler::gtTypeForNullCheck(GenTree* tree)
+{
+    if (varTypeIsIntegralOrI(tree))
+    {
+#if defined(TARGET_XARCH)
+        // Just an optimization for XARCH - smaller mov
+        if (varTypeIsLong(tree))
+        {
+            return TYP_INT;
+        }
+#endif
+        return tree->TypeGet();
+    }
+    // for the rest: probe a single byte to avoid potential AVEs
+    return TYP_BYTE;
+}
+
 //------------------------------------------------------------------------------
 // gtChangeOperToNullCheck: helper to change tree oper to a NULLCHECK.
 //
@@ -9752,7 +9775,7 @@ void Compiler::gtChangeOperToNullCheck(GenTree* tree, BasicBlock* block)
 {
     assert(tree->OperIs(GT_FIELD, GT_IND, GT_OBJ, GT_BLK));
     tree->ChangeOper(GT_NULLCHECK);
-    tree->ChangeType(TYP_INT);
+    tree->ChangeType(gtTypeForNullCheck(tree));
     block->bbFlags |= BBF_HAS_NULLCHECK;
     optMethodFlags |= OMF_HAS_NULLCHECK;
 }

src/coreclr/jit/compiler.h

@@ -3361,6 +3361,7 @@ class Compiler
 
     GenTree* gtNewNullCheck(GenTree* addr, BasicBlock* basicBlock);
 
+    var_types gtTypeForNullCheck(GenTree* tree);
     void gtChangeOperToNullCheck(GenTree* tree, BasicBlock* block);
 
     static fgArgTabEntry* gtArgEntryByArgNum(GenTreeCall* call, unsigned argNum);

src/coreclr/jit/emitarm64.cpp

@@ -6993,13 +6993,6 @@ void emitter::emitIns_R_R_R_Ext(instruction ins,
         shiftAmount = insOptsLSL(opt) ? scale : 0;
     }
 
-    // If target reg is ZR - it means we're doing an implicit nullcheck
-    // where target type was ignored and set to TYP_INT.
-    if ((reg1 == REG_ZR) && (shiftAmount > 0))
-    {
-        shiftAmount = scale;
-    }
-
     assert((shiftAmount == scale) || (shiftAmount == 0));
 
     reg2 = encodingSPtoZR(reg2);

src/coreclr/jit/lower.cpp

@@ -7193,11 +7193,14 @@ void Lowering::TransformUnusedIndirection(GenTreeIndir* ind, Compiler* comp, Bas
     // - On ARM64, always use GT_NULLCHECK for a dead indirection.
     // - On ARM, always use GT_IND.
     // - On XARCH, use GT_IND if we have a contained address, and GT_NULLCHECK otherwise.
-    // In all cases, change the type to TYP_INT.
+    // In all cases we try to preserve the original type and never make it wider to avoid AVEs.
+    // For structs we conservatively lower it to BYTE. For 8-byte primitives we lower it to TYP_INT
+    // on XARCH as an optimization.
     //
     assert(ind->OperIs(GT_NULLCHECK, GT_IND, GT_BLK, GT_OBJ));
 
-    ind->gtType = TYP_INT;
+    ind->ChangeType(comp->gtTypeForNullCheck(ind));
+
 #ifdef TARGET_ARM64
     bool useNullCheck = true;
 #elif TARGET_ARM

src/coreclr/jit/lsraxarch.cpp

@@ -610,6 +610,15 @@ int LinearScan::BuildNode(GenTree* tree)
         case GT_NULLCHECK:
         {
             assert(dstCount == 0);
+#ifdef TARGET_X86
+            if (varTypeIsByte(tree))
+            {
+                // on X86 we have to use byte-able regs for byte-wide loads
+                BuildUse(tree->gtGetOp1(), RBM_BYTE_REGS);
+                srcCount = 1;
+                break;
+            }
+#endif
             // If we have a contained address on a nullcheck, we transform it to
             // an unused GT_IND, since we require a target register.
             BuildUse(tree->gtGetOp1());

src/tests/JIT/Regression/JitBlue/Runtime_64657/Runtime_64657.cs

@@ -0,0 +1,113 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Runtime.InteropServices;
+using System.Runtime.CompilerServices;
+
+public unsafe class Runtime_64657
+{
+    [DllImport("kernel32")]
+    public static extern byte* VirtualAlloc(IntPtr lpAddress, nuint dwSize, uint flAllocationType, uint flProtect);
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    static void Validate<T>(T* c, int x) where T : unmanaged
+    {
+        // this nullcheck should not read more than requested
+        T implicitNullcheck = c[x];
+    }
+
+    public static int Main()
+    {
+        if (!OperatingSystem.IsWindows())
+            return 100; // VirtualAlloc is only for Windows
+
+        uint length = (uint)Environment.SystemPageSize;
+        byte* ptr = VirtualAlloc(IntPtr.Zero, length, 0x1000 | 0x2000 /* reserve commit */, 0x04 /*readonly guard*/);
+
+        Validate((byte*)(ptr + length - sizeof(byte)), 0);
+        Validate((sbyte*)(ptr + length - sizeof(sbyte)), 0);
+        Validate((bool*)(ptr + length - sizeof(bool)), 0);
+        Validate((ushort*)(ptr + length - sizeof(ushort)), 0);
+        Validate((short*)(ptr + length - sizeof(short)), 0);
+        Validate((uint*)(ptr + length - sizeof(uint)), 0);
+        Validate((int*)(ptr + length - sizeof(int)), 0);
+        Validate((ulong*)(ptr + length - sizeof(ulong)), 0);
+        Validate((long*)(ptr + length - sizeof(long)), 0);
+        Validate((nint*)(ptr + length - sizeof(nint)), 0);
+        Validate((nuint*)(ptr + length - sizeof(nuint)), 0);
+
+        Validate((S1*)(ptr + length - sizeof(S1)), 0);
+        Validate((S2*)(ptr + length - sizeof(S2)), 0);
+        Validate((S3*)(ptr + length - sizeof(S3)), 0);
+        Validate((S4*)(ptr + length - sizeof(S4)), 0);
+
+        TestStructures();
+
+        return 100;
+    }
+
+    private static void TestStructures()
+    {
+        S1 s1 = new S1();
+        TestS1_1(ref s1);
+        TestS1_2(ref s1);
+
+        S2 s2 = new S2();
+        TestS2_1(ref s2);
+        TestS2_2(ref s2);
+
+        S3 s3 = new S3();
+        TestS3_1(ref s3);
+        TestS3_2(ref s3);
+
+        S4 s4 = new S4();
+        TestS4_1(ref s4);
+        TestS4_2(ref s4);
+
+        S5 s5 = new S5 { a1 = "1", a2 = "2" };
+        TestS5_1(ref s5);
+        TestS5_2(ref s5);
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)] static void TestS1_1(ref S1 s) { var _ = s.a1; }
+    [MethodImpl(MethodImplOptions.NoInlining)] static void TestS1_2(ref S1 s) { var _ = s.a2; }
+    [MethodImpl(MethodImplOptions.NoInlining)] static void TestS2_1(ref S2 s) { var _ = s.a1; }
+    [MethodImpl(MethodImplOptions.NoInlining)] static void TestS2_2(ref S2 s) { var _ = s.a2; }
+    [MethodImpl(MethodImplOptions.NoInlining)] static void TestS3_1(ref S3 s) { var _ = s.a1; }
+    [MethodImpl(MethodImplOptions.NoInlining)] static void TestS3_2(ref S3 s) { var _ = s.a2; }
+    [MethodImpl(MethodImplOptions.NoInlining)] static void TestS4_1(ref S4 s) { var _ = s.a1; }
+    [MethodImpl(MethodImplOptions.NoInlining)] static void TestS4_2(ref S4 s) { var _ = s.a2; }
+    [MethodImpl(MethodImplOptions.NoInlining)] static void TestS5_1(ref S5 s) { var _ = s.a1; }
+    [MethodImpl(MethodImplOptions.NoInlining)] static void TestS5_2(ref S5 s) { var _ = s.a2; }
+
+    public struct S1
+    {
+        public byte a1;
+        public byte a2;
+    }
+
+    public struct S2
+    {
+        public short a1;
+        public short a2;
+    }
+
+    public struct S3
+    {
+        public int a1;
+        public int a2;
+    }
+
+    public struct S4
+    {
+        public long a1;
+        public long a2;
+    }
+
+    public struct S5
+    {
+        public string a1;
+        public string a2;
+    }
+}

src/tests/JIT/Regression/JitBlue/Runtime_64657/Runtime_64657.csproj

@@ -0,0 +1,10 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <OutputType>Exe</OutputType>
+    <Optimize>True</Optimize>
+    <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="$(MSBuildProjectName).cs" />
+  </ItemGroup>
+</Project>