Arm64: Add If conversion pass (#73472)

* Arm64: Add If conversion pass

* Minor review fixups

* Return a PhaseStatus

* Fix formatting

* Check for side effects on NOPs

* Add function block comments for the phase

* Remove creation of AND chains from if conversion pass

* Update middleBlock flow

* Check for order side effects

* Remove COLON_COND check

* Remove flag toggling

* Move the conditional assignment to the JTRUE block

* Fix formatting

* Allow conditions with side effects

* Fix formatting

* Correct all moved SSA statements

* Add size costing check

* Only move middle block ssa defs

* Fix formatting

* Fewer SSA assumptions

* Use implicit func for value numbering

* Update header for gtFoldExprConditional

* Cost based on speed

* Add Stress mode for inner loops

* Rework costings

* Check for invalid VNs

* Ignore compares against zero

* Ensure float compares are contained

* Allow if conversion of test compares

* Do not contain test compares within compare chains

* Add float versions of the JIT/opt/Compares tests

* Fix formatting

* Compare chains use CmpCompares, selects use Compares

* Fix flow checking for empty blocks

* Fix to contexts setting JitStdOutFile

* Fix attr and reg producing in select generation

src/coreclr/jit/codegenarm64.cpp

@@ -4613,6 +4613,9 @@ void CodeGen::genCodeForConditionalCompare(GenTreeOp* tree, GenCondition prevCon
     // Should only be called on contained nodes.
     assert(targetReg == REG_NA);
 
+    // Should not be called for test conditionals (Arm64 does not have a ctst).
+    assert(tree->OperIsCmpCompare());
+
     // For the ccmp flags, invert the condition of the compare.
     insCflags cflags = InsCflagsForCcmp(GenCondition::FromRelop(tree));
 
@@ -4711,7 +4714,7 @@ void CodeGen::genCodeForSelect(GenTreeConditional* tree)
     GenTree*  op2     = tree->gtOp2;
     var_types op1Type = genActualType(op1->TypeGet());
     var_types op2Type = genActualType(op2->TypeGet());
-    emitAttr  cmpSize = EA_ATTR(genTypeSize(op1Type));
+    emitAttr  attr    = emitActualTypeSize(tree->TypeGet());
 
     assert(!op1->isUsedFromMemory());
     assert(genTypeSize(op1Type) == genTypeSize(op2Type));
@@ -4721,10 +4724,20 @@ void CodeGen::genCodeForSelect(GenTreeConditional* tree)
     if (opcond->isContained())
     {
         // Generate the contained condition.
-        bool chain = false;
-        JITDUMP("Generating compare chain:\n");
-        genCodeForContainedCompareChain(opcond, &chain, &prevCond);
-        assert(chain);
+        if (opcond->OperIsCompare())
+        {
+            genCodeForCompare(opcond->AsOp());
+            prevCond = GenCondition::FromRelop(opcond);
+        }
+        else
+        {
+            // Condition is a compare chain. Try to contain it.
+            assert(opcond->OperIs(GT_AND));
+            bool chain = false;
+            JITDUMP("Generating compare chain:\n");
+            genCodeForContainedCompareChain(opcond, &chain, &prevCond);
+            assert(chain);
+        }
     }
     else
     {
@@ -4738,8 +4751,9 @@ void CodeGen::genCodeForSelect(GenTreeConditional* tree)
     regNumber               srcReg2   = genConsumeReg(op2);
     const GenConditionDesc& prevDesc  = GenConditionDesc::Get(prevCond);
 
-    emit->emitIns_R_R_R_COND(INS_csel, cmpSize, targetReg, srcReg1, srcReg2, JumpKindToInsCond(prevDesc.jumpKind1));
+    emit->emitIns_R_R_R_COND(INS_csel, attr, targetReg, srcReg1, srcReg2, JumpKindToInsCond(prevDesc.jumpKind1));
     regSet.verifyRegUsed(targetReg);
+    genProduceReg(tree);
 }
 
 //------------------------------------------------------------------------

src/coreclr/jit/codegenlinear.cpp

@@ -1634,9 +1634,10 @@ void CodeGen::genConsumeRegs(GenTree* tree)
             assert(cast->isContained());
             genConsumeAddress(cast->CastOp());
         }
-        else if (tree->OperIsCmpCompare() || tree->OperIs(GT_AND))
+        else if (tree->OperIsCompare() || tree->OperIs(GT_AND))
         {
-            // Compares and ANDs may be contained in a conditional chain.
+            // Compares can be contained by a SELECT.
+            // ANDs and Cmp Compares may be contained in a chain.
             genConsumeRegs(tree->gtGetOp1());
             genConsumeRegs(tree->gtGetOp2());
         }

src/coreclr/jit/compiler.cpp

@@ -4741,6 +4741,7 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
         bool doCse           = true;
         bool doAssertionProp = true;
         bool doRangeAnalysis = true;
+        bool doIfConversion  = true;
         int  iterations      = 1;
 
 #if defined(OPT_CONFIG)
@@ -4753,6 +4754,7 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
         doCse           = doValueNum;
         doAssertionProp = doValueNum && (JitConfig.JitDoAssertionProp() != 0);
         doRangeAnalysis = doAssertionProp && (JitConfig.JitDoRangeAnalysis() != 0);
+        doIfConversion  = doIfConversion && (JitConfig.JitDoIfConversion() != 0);
 
         if (opts.optRepeat)
         {
@@ -4834,6 +4836,13 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
                 DoPhase(this, PHASE_ASSERTION_PROP_MAIN, &Compiler::optAssertionPropMain);
             }
 
+            if (doIfConversion)
+            {
+                // If conversion
+                //
+                DoPhase(this, PHASE_IF_CONVERSION, &Compiler::optIfConversion);
+            }
+
             if (doRangeAnalysis)
             {
                 // Bounds check elimination via range analysis

src/coreclr/jit/compiler.h

@@ -2392,6 +2392,9 @@ class Compiler
     GenTreeLclVar* gtNewLclVarAddrNode(unsigned lclNum, var_types type = TYP_I_IMPL);
     GenTreeLclFld* gtNewLclFldAddrNode(unsigned lclNum, unsigned lclOffs, var_types type = TYP_I_IMPL);
 
+    GenTreeConditional* gtNewConditionalNode(
+        genTreeOps oper, GenTree* cond, GenTree* op1, GenTree* op2, var_types type);
+
 #ifdef FEATURE_SIMD
     GenTreeSIMD* gtNewSIMDNode(
         var_types type, GenTree* op1, SIMDIntrinsicID simdIntrinsicID, CorInfoType simdBaseJitType, unsigned simdSize);
@@ -2831,6 +2834,7 @@ class Compiler
     GenTree* gtFoldExprSpecial(GenTree* tree);
     GenTree* gtFoldBoxNullable(GenTree* tree);
     GenTree* gtFoldExprCompare(GenTree* tree);
+    GenTree* gtFoldExprConditional(GenTree* tree);
     GenTree* gtCreateHandleCompare(genTreeOps             oper,
                                    GenTree*               op1,
                                    GenTree*               op2,
@@ -4869,6 +4873,9 @@ class Compiler
     // Does value-numbering for a block assignment.
     void fgValueNumberBlockAssignment(GenTree* tree);
 
+    // Does value-numbering for a variable definition that has SSA.
+    void fgValueNumberSsaVarDef(GenTreeLclVarCommon* lcl);
+
     // Does value-numbering for a cast tree.
     void fgValueNumberCastTree(GenTree* tree);
 
@@ -6082,6 +6089,7 @@ class Compiler
     void optEnsureUniqueHead(unsigned loopInd, weight_t ambientWeight);
     PhaseStatus optUnrollLoops(); // Unrolls loops (needs to have cost info)
     void        optRemoveRedundantZeroInits();
+    PhaseStatus optIfConversion(); // If conversion
 
 protected:
     // This enumeration describes what is killed by a call.
@@ -6470,6 +6478,7 @@ class Compiler
     OptInvertCountTreeInfoType optInvertCountTreeInfo(GenTree* tree);
 
     bool optInvertWhileLoop(BasicBlock* block);
+    bool optIfConvert(BasicBlock* block);
 
 private:
     static bool optIterSmallOverflow(int iterAtExit, var_types incrType);
@@ -7378,6 +7387,7 @@ class Compiler
     GenTree* optAssertionProp_Cast(ASSERT_VALARG_TP assertions, GenTreeCast* cast, Statement* stmt);
     GenTree* optAssertionProp_Call(ASSERT_VALARG_TP assertions, GenTreeCall* call, Statement* stmt);
     GenTree* optAssertionProp_RelOp(ASSERT_VALARG_TP assertions, GenTree* tree, Statement* stmt);
+    GenTree* optAssertionProp_ConditionalOp(ASSERT_VALARG_TP assertions, GenTree* tree, Statement* stmt);
     GenTree* optAssertionProp_Comma(ASSERT_VALARG_TP assertions, GenTree* tree, Statement* stmt);
     GenTree* optAssertionProp_BndsChk(ASSERT_VALARG_TP assertions, GenTree* tree, Statement* stmt);
     GenTree* optAssertionPropGlobal_RelOp(ASSERT_VALARG_TP assertions, GenTree* tree, Statement* stmt);
@@ -9576,6 +9586,8 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
         STRESS_MODE(EMITTER)                                                                    \
         STRESS_MODE(CHK_REIMPORT)                                                               \
         STRESS_MODE(GENERIC_CHECK)                                                              \
+        STRESS_MODE(IF_CONVERSION_COST)                                                         \
+        STRESS_MODE(IF_CONVERSION_INNER_LOOPS)                                                  \
         STRESS_MODE(COUNT)
 
     enum                compStressArea

src/coreclr/jit/compphases.h

@@ -87,6 +87,7 @@ CompPhaseNameMacro(PHASE_OPTIMIZE_VALNUM_CSES,       "Optimize Valnum CSEs",
 CompPhaseNameMacro(PHASE_VN_COPY_PROP,               "VN based copy prop",             false, -1, false)
 CompPhaseNameMacro(PHASE_OPTIMIZE_BRANCHES,          "Redundant branch opts",          false, -1, false)
 CompPhaseNameMacro(PHASE_ASSERTION_PROP_MAIN,        "Assertion prop",                 false, -1, false)
+CompPhaseNameMacro(PHASE_IF_CONVERSION,              "If conversion",                  false, -1, false)
 CompPhaseNameMacro(PHASE_OPT_UPDATE_FLOW_GRAPH,      "Update flow graph opt pass",     false, -1, false)
 CompPhaseNameMacro(PHASE_COMPUTE_EDGE_WEIGHTS2,      "Compute edge weights (2, false)",false, -1, false)
 CompPhaseNameMacro(PHASE_INSERT_GC_POLLS,            "Insert GC Polls",                false, -1, true)

src/coreclr/jit/gentree.cpp

@@ -5847,6 +5847,9 @@ unsigned Compiler::gtSetEvalOrder(GenTree* tree)
             level = max(level, lvl2);
             costEx += tree->AsConditional()->gtOp2->GetCostEx();
             costSz += tree->AsConditional()->gtOp2->GetCostSz();
+
+            costEx += 1;
+            costSz += 1;
             break;
 
         default:
@@ -7387,6 +7390,17 @@ GenTreeLclFld* Compiler::gtNewLclFldAddrNode(unsigned lclNum, unsigned lclOffs,
     return node;
 }
 
+GenTreeConditional* Compiler::gtNewConditionalNode(
+    genTreeOps oper, GenTree* cond, GenTree* op1, GenTree* op2, var_types type)
+{
+    assert(GenTree::OperIsConditional(oper));
+    GenTreeConditional* node = new (this, oper) GenTreeConditional(oper, type, cond, op1, op2);
+    node->gtFlags |= (cond->gtFlags & GTF_ALL_EFFECT);
+    node->gtFlags |= (op1->gtFlags & GTF_ALL_EFFECT);
+    node->gtFlags |= (op2->gtFlags & GTF_ALL_EFFECT);
+    return node;
+}
+
 GenTreeLclFld* Compiler::gtNewLclFldNode(unsigned lnum, var_types type, unsigned offset)
 {
     GenTreeLclFld* node = new (this, GT_LCL_FLD) GenTreeLclFld(GT_LCL_FLD, type, lnum, offset);
@@ -12643,6 +12657,10 @@ GenTree* Compiler::gtFoldExpr(GenTree* tree)
 
     if (!(kind & GTK_SMPOP))
     {
+        if (tree->OperIsConditional())
+        {
+            return gtFoldExprConditional(tree);
+        }
         return tree;
     }
 
@@ -12911,6 +12929,118 @@ GenTree* Compiler::gtFoldExprCompare(GenTree* tree)
     return cons;
 }
 
+//------------------------------------------------------------------------
+// gtFoldExprConditional: see if a conditional is foldable
+//
+// Arguments:
+//    tree - condition to examine
+//
+// Returns:
+//    The original call if no folding happened.
+//    An alternative tree if folding happens.
+//
+// Notes:
+//    Supporting foldings are:
+//      SELECT TRUE  X Y  ->  X
+//      SELECT FALSE X Y  ->  Y
+//      SELECT COND  X X  ->  X
+//
+GenTree* Compiler::gtFoldExprConditional(GenTree* tree)
+{
+    GenTree* cond = tree->AsConditional()->gtCond;
+    GenTree* op1  = tree->AsConditional()->gtOp1;
+    GenTree* op2  = tree->AsConditional()->gtOp2;
+
+    assert(tree->OperIsConditional());
+
+    // Check for a constant conditional
+    if (cond->OperIsConst())
+    {
+        // Constant conditions must be folded away.
+
+        JITDUMP("\nFolding conditional op with constant condition:\n");
+        DISPTREE(tree);
+
+        assert(cond->TypeIs(TYP_INT));
+        assert((tree->gtFlags & GTF_SIDE_EFFECT & ~GTF_ASG) == 0);
+        assert((tree->gtFlags & GTF_ORDER_SIDEEFF) == 0);
+
+        GenTree* replacement = nullptr;
+        if (cond->IsIntegralConst(0))
+        {
+            JITDUMP("Bashed to false path:\n");
+            replacement = op2;
+        }
+        else
+        {
+            // Condition should never be a constant other than 0 or 1
+            assert(cond->IsIntegralConst(1));
+            JITDUMP("Bashed to true path:\n");
+            replacement = op1;
+        }
+
+        if (fgGlobalMorph)
+        {
+            fgMorphTreeDone(replacement);
+        }
+        else
+        {
+            replacement->gtNext = tree->gtNext;
+            replacement->gtPrev = tree->gtPrev;
+        }
+        DISPTREE(replacement);
+        JITDUMP("\n");
+
+        // If we bashed to a compare, try to fold that.
+        if (replacement->OperIsCompare())
+        {
+            return gtFoldExprCompare(replacement);
+        }
+
+        return replacement;
+    }
+
+    assert(cond->OperIsCompare());
+
+    if (((tree->gtFlags & GTF_SIDE_EFFECT) != 0) || !GenTree::Compare(op1, op2, true))
+    {
+        // No folding.
+        return tree;
+    }
+
+    // GTF_ORDER_SIDEEFF here may indicate volatile subtrees.
+    // Or it may indicate a non-null assertion prop into an indir subtree.
+    if ((tree->gtFlags & GTF_ORDER_SIDEEFF) != 0)
+    {
+        // If op1 is "volatile" and op2 is not, we can still fold.
+        const bool op1MayBeVolatile = (op1->gtFlags & GTF_ORDER_SIDEEFF) != 0;
+        const bool op2MayBeVolatile = (op2->gtFlags & GTF_ORDER_SIDEEFF) != 0;
+
+        if (!op1MayBeVolatile || op2MayBeVolatile)
+        {
+            // No folding.
+            return tree;
+        }
+    }
+
+    JITDUMP("Bashed to first of two identical paths:\n");
+    GenTree* replacement = op1;
+
+    if (fgGlobalMorph)
+    {
+        fgMorphTreeDone(replacement);
+    }
+    else
+    {
+        replacement->gtNext = tree->gtNext;
+        replacement->gtPrev = tree->gtPrev;
+    }
+    DISPTREE(replacement);
+    JITDUMP("\n");
+
+    return replacement;
+}
+
 //------------------------------------------------------------------------
 // gtCreateHandleCompare: generate a type handle comparison
 //

src/coreclr/jit/jitconfigvalues.h

@@ -426,6 +426,7 @@ CONFIG_INTEGER(JitDoValueNumber, W("JitDoValueNumber"), 1) // Perform value numb
 
 CONFIG_METHODSET(JitOptRepeat, W("JitOptRepeat"))            // Runs optimizer multiple times on the method
 CONFIG_INTEGER(JitOptRepeatCount, W("JitOptRepeatCount"), 2) // Number of times to repeat opts when repeating
+CONFIG_INTEGER(JitDoIfConversion, W("JitDoIfConversion"), 1) // Perform If conversion
 #endif                                                       // defined(OPT_CONFIG)
 
 CONFIG_INTEGER(JitTelemetry, W("JitTelemetry"), 1) // If non-zero, gather JIT telemetry data

src/coreclr/jit/lowerarmarch.cpp

@@ -2327,8 +2327,7 @@ bool Lowering::ContainCheckCompareChain(GenTree* child, GenTree* parent, GenTree
             child->SetContained();
             return true;
         }
-        else if (child->OperIsCmpCompare() && varTypeIsIntegral(child->gtGetOp1()) &&
-                 varTypeIsIntegral(child->gtGetOp2()))
+        else if (child->OperIsCmpCompare())
         {
             child->AsOp()->SetContained();
 
@@ -2422,17 +2421,28 @@ void Lowering::ContainCheckSelect(GenTreeConditional* node)
         return;
     }
 
-    // Check if the compare does not need to be generated into a register.
-    GenTree* startOfChain = nullptr;
-    ContainCheckCompareChain(node->gtCond, node, &startOfChain);
-
-    if (startOfChain != nullptr)
+    if (node->gtCond->OperIsCompare())
+    {
+        // All compare node types (including TEST_) are containable.
+        if (IsSafeToContainMem(node, node->gtCond))
+        {
+            node->gtCond->AsOp()->SetContained();
+        }
+    }
+    else
     {
-        // The earliest node in the chain will be generated as a standard compare.
-        assert(startOfChain->OperIsCmpCompare());
-        startOfChain->AsOp()->gtGetOp1()->ClearContained();
-        startOfChain->AsOp()->gtGetOp2()->ClearContained();
-        ContainCheckCompare(startOfChain->AsOp());
+        // Check for a compare chain and try to contain it.
+        GenTree* startOfChain = nullptr;
+        ContainCheckCompareChain(node->gtCond, node, &startOfChain);
+
+        if (startOfChain != nullptr)
+        {
+            // The earliest node in the chain will be generated as a standard compare.
+            assert(startOfChain->OperIsCmpCompare());
+            startOfChain->AsOp()->gtGetOp1()->ClearContained();
+            startOfChain->AsOp()->gtGetOp2()->ClearContained();
+            ContainCheckCompare(startOfChain->AsOp());
+        }
     }
 }
 

src/coreclr/jit/lsrabuild.cpp

@@ -3226,10 +3226,11 @@ int LinearScan::BuildOperandUses(GenTree* node, regMaskTP candidates)
     }
 #endif // FEATURE_HW_INTRINSICS
 #ifdef TARGET_ARM64
-    if (node->OperIs(GT_MUL) || node->OperIsCmpCompare() || node->OperIs(GT_AND))
+    if (node->OperIs(GT_MUL) || node->OperIsCompare() || node->OperIs(GT_AND))
     {
         // MUL can be contained for madd or msub on arm64.
-        // Compare and AND may be contained due to If Conversion.
+        // Compares can be contained by a SELECT.
+        // ANDs and Cmp Compares may be contained in a chain.
         return BuildBinaryUses(node->AsOp(), candidates);
     }
     if (node->OperIs(GT_NEG, GT_CAST, GT_LSH, GT_RSH, GT_RSZ))