ARM64 - Optimizing a % b operations part 2 (#66407)

* Initial work for ARM64 mod optimization

* Updated comment

* Updated comment

* Updated comment

* Fixing build

* Remove uneeded var

* Use '%' morph logic for both x64/arm64

* Adding back in divisor check for x64

* Formatting

* Update comments

* Update comments

* Fixing

* Updated comment

* Updated comment

* Tweaking x64 transformation logic for the mod opt

* Tweaking x64 transformation logic for the mod opt

* Using IntCon

* Fixed build

* Minor tweak

* Fixing x64 diffs

* Removing flag set

* Optimizing signed 'mod' for ARM64 when op2 is pow2

* Updated comments

* Updated comments

* Updated comments

* Updated comments

* Updated comments

* Updated comments

* Feedback

* Merging and fixing test failure by lowering node rather than marking it as contained

* Fixing build

* Feedback

* minor updates

* minor updates

* minor updates

* Fixing tests

* Fixing tests

* Fixing tests

* Formatting

* Fixing tests

* Using AsIntConCommonet()->IntegralValue()

* Feedback

* Fixing build

* Fixing tests

* Fixing tests

* Fixing tests

* Adding tests

* Formatting

* Renamed GT_CS_NEG_MI to GT_CSNEG_MI

* Updated comments, fixed test

* Formatting

* Update codegenarm64.cpp

* Do not use IsIntegralConstAbsPow2 for the check. Use IsIntegralConstPow2

* Update comment

Author: TIHan

src/coreclr/jit/codegen.h

@@ -1381,6 +1381,7 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     void genCodeForMsub(GenTreeOp* tree);
     void genCodeForBfiz(GenTreeOp* tree);
     void genCodeForAddEx(GenTreeOp* tree);
+    void genCodeForCond(GenTreeOp* tree);
 #endif // TARGET_ARM64
 
 #if defined(FEATURE_EH_FUNCLETS)

src/coreclr/jit/codegenarm64.cpp

@@ -3039,6 +3039,18 @@ void CodeGen::genCodeForNegNot(GenTree* tree)
     regNumber   targetReg = tree->GetRegNum();
     instruction ins       = genGetInsForOper(tree->OperGet(), targetType);
 
+    if ((tree->gtFlags & GTF_SET_FLAGS) != 0)
+    {
+        switch (tree->OperGet())
+        {
+            case GT_NEG:
+                ins = INS_negs;
+                break;
+            default:
+                noway_assert(!"Unexpected UnaryOp with GTF_SET_FLAGS set");
+        }
+    }
+
     // The arithmetic node must be sitting in a register (since it's not contained)
     assert(!tree->isContained());
     // The dst can only be a register.
@@ -10190,4 +10202,40 @@ void CodeGen::genCodeForAddEx(GenTreeOp* tree)
     genProduceReg(tree);
 }
 
+//------------------------------------------------------------------------
+// genCodeForCond: Generates the code sequence for a GenTree node that
+// represents a conditional instruction.
+//
+// Arguments:
+//    tree - conditional op
+//
+void CodeGen::genCodeForCond(GenTreeOp* tree)
+{
+    assert(tree->OperIs(GT_CSNEG_MI));
+    assert(!(tree->gtFlags & GTF_SET_FLAGS) && (tree->gtFlags & GTF_USE_FLAGS));
+    genConsumeOperands(tree);
+
+    instruction ins;
+    insCond     cond;
+    switch (tree->OperGet())
+    {
+        case GT_CSNEG_MI:
+        {
+            ins  = INS_csneg;
+            cond = INS_COND_MI;
+            break;
+        }
+
+        default:
+            unreached();
+    }
+
+    regNumber dstReg = tree->GetRegNum();
+    regNumber op1Reg = tree->gtGetOp1()->GetRegNum();
+    regNumber op2Reg = tree->gtGetOp2()->GetRegNum();
+
+    GetEmitter()->emitIns_R_R_R_COND(ins, emitActualTypeSize(tree), dstReg, op1Reg, op2Reg, cond);
+    genProduceReg(tree);
+}
+
 #endif // TARGET_ARM64

src/coreclr/jit/codegenarmarch.cpp

@@ -329,6 +329,10 @@ void CodeGen::genCodeForTreeNode(GenTree* treeNode)
         case GT_BFIZ:
             genCodeForBfiz(treeNode->AsOp());
             break;
+
+        case GT_CSNEG_MI:
+            genCodeForCond(treeNode->AsOp());
+            break;
 #endif // TARGET_ARM64
 
         case GT_JMP:
@@ -550,14 +554,7 @@ void CodeGen::genCodeForTreeNode(GenTree* treeNode)
 
         default:
         {
-#ifdef DEBUG
-            char message[256];
-            _snprintf_s(message, ArrLen(message), _TRUNCATE, "NYI: Unimplemented node type %s",
-                        GenTree::OpName(treeNode->OperGet()));
-            NYIRAW(message);
-#else
-            NYI("unimplemented node");
-#endif
+            unreached();
         }
         break;
     }

src/coreclr/jit/gentree.h

@@ -2075,6 +2075,8 @@ struct GenTree
 
     inline bool IsIntegralConst() const;
 
+    inline bool IsIntegralConstPow2() const;
+
     inline bool IsIntegralConstUnsignedPow2() const;
 
     inline bool IsIntegralConstAbsPow2() const;
@@ -8454,6 +8456,24 @@ inline bool GenTree::IsIntegralConst() const
 #endif // !TARGET_64BIT
 }
 
+//-------------------------------------------------------------------------
+// IsIntegralConstPow2: Determines whether an integral constant is
+//                      the power of 2.
+//
+// Return Value:
+//     Returns true if the GenTree's integral constant
+//     is the power of 2.
+//
+inline bool GenTree::IsIntegralConstPow2() const
+{
+    if (IsIntegralConst())
+    {
+        return isPow2(AsIntConCommon()->IntegralValue());
+    }
+
+    return false;
+}
+
 //-------------------------------------------------------------------------
 // IsIntegralConstUnsignedPow2: Determines whether the unsigned value of
 //                              an integral constant is the power of 2.

src/coreclr/jit/gtlist.h

@@ -227,6 +227,7 @@ GTNODE(MSUB             , GenTreeOp          ,0,GTK_BINOP|DBK_NOTHIR) // Generat
                                                                       // enabling it for both armarch and xarch for floating-point MSUB "unsafe" math.
 GTNODE(ADDEX,             GenTreeOp          ,0,GTK_BINOP|DBK_NOTHIR) // Add with sign/zero extension.
 GTNODE(BFIZ             , GenTreeOp          ,0,GTK_BINOP|DBK_NOTHIR) // Bitfield Insert in Zero.
+GTNODE(CSNEG_MI         , GenTreeOp          ,0,GTK_BINOP|DBK_NOTHIR) // Conditional select, negate, minus result
 #endif
 
 //-----------------------------------------------------------------------------

src/coreclr/jit/lower.cpp

@@ -5861,6 +5861,10 @@ GenTree* Lowering::LowerConstIntDivOrMod(GenTree* node)
     assert(!"unreachable: integral GT_DIV/GT_MOD should get morphed into helper calls");
 #endif // USE_HELPERS_FOR_INT_DIV
 #if defined(TARGET_ARM64)
+    if (divMod->OperIs(GT_MOD) && divisor->IsIntegralConstPow2())
+    {
+        return LowerModPow2(node);
+    }
     assert(node->OperGet() != GT_MOD);
 #endif // TARGET_ARM64
 

src/coreclr/jit/lower.h

@@ -353,6 +353,7 @@ class Lowering final : public Phase
 #elif defined(TARGET_ARM64)
     bool IsValidConstForMovImm(GenTreeHWIntrinsic* node);
     void LowerHWIntrinsicFusedMultiplyAddScalar(GenTreeHWIntrinsic* node);
+    GenTree* LowerModPow2(GenTree* node);
 #endif // !TARGET_XARCH && !TARGET_ARM64
 
     union VectorConstant {

src/coreclr/jit/lowerarmarch.cpp

@@ -610,6 +610,95 @@ void Lowering::LowerRotate(GenTree* tree)
     ContainCheckShiftRotate(tree->AsOp());
 }
 
+#ifdef TARGET_ARM64
+//------------------------------------------------------------------------
+// LowerModPow2: Lower GT_MOD if the second operand is a constant power of 2.
+//
+// Arguments:
+//    tree - the node to lower
+//
+// Return Value:
+//    A new tree node if it changed.
+//
+// Notes:
+//     {expr} % {cns}
+//     Logically turns into:
+//         let a = {expr}
+//         if a > 0 then (a & ({cns} - 1)) else -(-a & ({cns} - 1))
+//     which then turns into:
+//         and   reg1, reg0, #({cns} - 1)
+//         negs  reg0, reg0
+//         and   reg0, reg0, #({cns} - 1)
+//         csneg reg0, reg1, reg0, mi
+//     TODO: We could do this optimization in morph but we do not have
+//           a conditional select op in HIR. At some point, we may
+//           introduce such an op.
+GenTree* Lowering::LowerModPow2(GenTree* node)
+{
+    assert(node->OperIs(GT_MOD));
+    GenTree* mod      = node;
+    GenTree* dividend = mod->gtGetOp1();
+    GenTree* divisor  = mod->gtGetOp2();
+
+    assert(divisor->IsIntegralConstPow2());
+
+    const var_types type = mod->TypeGet();
+    assert((type == TYP_INT) || (type == TYP_LONG));
+
+    LIR::Use use;
+    if (!BlockRange().TryGetUse(node, &use))
+    {
+        return nullptr;
+    }
+
+    ssize_t cnsValue = static_cast<ssize_t>(divisor->AsIntConCommon()->IntegralValue()) - 1;
+
+    BlockRange().Remove(divisor);
+
+    // We need to use the dividend node multiple times so its value needs to be
+    // computed once and stored in a temp variable.
+    LIR::Use opDividend(BlockRange(), &mod->AsOp()->gtOp1, mod);
+    dividend = ReplaceWithLclVar(opDividend);
+
+    GenTree* dividend2 = comp->gtClone(dividend);
+    BlockRange().InsertAfter(dividend, dividend2);
+
+    GenTreeIntCon* cns = comp->gtNewIconNode(cnsValue, type);
+    BlockRange().InsertAfter(dividend2, cns);
+
+    GenTree* const trueExpr = comp->gtNewOperNode(GT_AND, type, dividend, cns);
+    BlockRange().InsertAfter(cns, trueExpr);
+    LowerNode(trueExpr);
+
+    GenTree* const neg = comp->gtNewOperNode(GT_NEG, type, dividend2);
+    neg->gtFlags |= GTF_SET_FLAGS;
+    BlockRange().InsertAfter(trueExpr, neg);
+
+    GenTreeIntCon* cns2 = comp->gtNewIconNode(cnsValue, type);
+    BlockRange().InsertAfter(neg, cns2);
+
+    GenTree* const falseExpr = comp->gtNewOperNode(GT_AND, type, neg, cns2);
+    BlockRange().InsertAfter(cns2, falseExpr);
+    LowerNode(falseExpr);
+
+    GenTree* const cc = comp->gtNewOperNode(GT_CSNEG_MI, type, trueExpr, falseExpr);
+    cc->gtFlags |= GTF_USE_FLAGS;
+
+    JITDUMP("Lower: optimize X MOD POW2");
+    DISPNODE(mod);
+    JITDUMP("to:\n");
+    DISPNODE(cc);
+
+    BlockRange().InsertBefore(mod, cc);
+    ContainCheckNode(cc);
+    BlockRange().Remove(mod);
+
+    use.ReplaceWith(cc);
+
+    return cc->gtNext;
+}
+#endif
+
 #ifdef FEATURE_HW_INTRINSICS
 
 //----------------------------------------------------------------------------------------------
@@ -1723,7 +1812,7 @@ void Lowering::ContainCheckMul(GenTreeOp* node)
 //
 void Lowering::ContainCheckDivOrMod(GenTreeOp* node)
 {
-    assert(node->OperIs(GT_DIV, GT_UDIV));
+    assert(node->OperIs(GT_DIV, GT_UDIV, GT_MOD));
 
     // ARM doesn't have a div instruction with an immediate operand
 }

src/coreclr/jit/morph.cpp

@@ -11652,7 +11652,10 @@ GenTree* Compiler::fgMorphSmpOp(GenTree* tree, MorphAddrContext* mac)
                 }
                 // ARM64 architecture manual suggests this transformation
                 // for the mod operator.
-                else
+                // However, we do skip this optimization for ARM64 if the second operand
+                // is an integral constant power of 2 because there is an even better
+                // optimization in lowering that is specific for ARM64.
+                else if (!(tree->OperIs(GT_MOD) && op2->IsIntegralConstPow2()))
 #else
                 // XARCH only applies this transformation if we know
                 // that magic division will be used - which is determined

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

@@ -0,0 +1,63 @@
+// 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.CompilerServices;
+
+class Program
+{
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    static uint PerformMod_1(uint i)
+    {
+        return i % 8;
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    static int PerformMod_2(int i)
+    {
+        return i % 16;
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    static int PerformMod_3(int i, int j)
+    {
+        return i % j;
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    static int MSUB(int a, int b, int c)
+    {
+        return a - b * c;
+    }
+
+    static int Main(string[] args)
+    {
+        var result = 100;
+
+        if (PerformMod_1(23) != 7)
+        {
+            result = -1;
+            Console.WriteLine("Failed Mod1!");
+        }
+
+        if (PerformMod_2(-23) != -7)
+        {
+            result = -1;
+            Console.WriteLine("Failed Mod2!");
+        }
+
+        if (PerformMod_3(23, 8) != 7)
+        {
+            result = -1;
+            Console.WriteLine("Failed Mod3!");
+        }
+
+        if (MSUB(3, 7, 8) != -53)
+        {
+            result = -1;
+            Console.WriteLine("Failed MSUB");
+        }
+
+        return result;
+    }
+}

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

@@ -0,0 +1,16 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <OutputType>Exe</OutputType>
+  </PropertyGroup>
+  <PropertyGroup>
+    <DebugType />
+    <Optimize>True</Optimize>
+    <AllowUnsafeBlocks>True</AllowUnsafeBlocks>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="$(MSBuildProjectName).cs" />
+  </ItemGroup>
+  <ItemGroup>
+    <ProjectReference Include="$(TestSourceDir)Common/CoreCLRTestLibrary/CoreCLRTestLibrary.csproj" />
+  </ItemGroup>
+</Project>