From 8101ace8695021295672d0eae9efc86bbcd9dc39 Mon Sep 17 00:00:00 2001
From: Andy Ayers <andya@microsoft.com>
Date: Fri, 9 Sep 2022 12:14:07 -0700
Subject: [PATCH] JIT: propagate some block assertions after expansion (#75321)

In #74384 I modified morph to generate assertions before a block op
was morphed, but after the source and dest were morphed. This misses
generating some assertions that arise once the expanded form is
known.

We now re-run assertion gen on the expanded tree, when we do an
`OneAsgBlock` expansion. Other cases my also prove profitable.

Closes #75229.
---
 src/coreclr/jit/morphblock.cpp | 29 +++++++++++++++++++++++++----
 1 file changed, 25 insertions(+), 4 deletions(-)

diff --git a/src/coreclr/jit/morphblock.cpp b/src/coreclr/jit/morphblock.cpp
index d3db65e1785c3..9a15673822991 100644
--- a/src/coreclr/jit/morphblock.cpp
+++ b/src/coreclr/jit/morphblock.cpp
@@ -19,7 +19,8 @@ class MorphInitBlockHelper
     virtual void TrySpecialCases();
     virtual void MorphStructCases();
 
-    void PropagateAssertions();
+    void PropagateBlockAssertions();
+    void PropagateExpansionAssertions();
 
     virtual const char* GetHelperName() const
     {
@@ -127,7 +128,7 @@ GenTree* MorphInitBlockHelper::Morph()
 
     PrepareDst();
     PrepareSrc();
-    PropagateAssertions();
+    PropagateBlockAssertions();
     TrySpecialCases();
 
     if (m_transformationDecision == BlockTransformation::Undefined)
@@ -150,6 +151,8 @@ GenTree* MorphInitBlockHelper::Morph()
         }
     }
 
+    PropagateExpansionAssertions();
+
     assert(m_transformationDecision != BlockTransformation::Undefined);
     assert(m_result != nullptr);
 
@@ -278,7 +281,7 @@ void MorphInitBlockHelper::PrepareDst()
 }
 
 //------------------------------------------------------------------------
-// PropagateAssertions: propagate assertions based on the original tree
+// PropagateBlockAssertions: propagate assertions based on the original tree
 //
 // Notes:
 //    Once the init or copy tree is morphed, assertion gen can no
@@ -286,7 +289,7 @@ void MorphInitBlockHelper::PrepareDst()
 //
 //    So we generate assertions based on the original tree.
 //
-void MorphInitBlockHelper::PropagateAssertions()
+void MorphInitBlockHelper::PropagateBlockAssertions()
 {
     if (m_comp->optLocalAssertionProp)
     {
@@ -294,6 +297,24 @@ void MorphInitBlockHelper::PropagateAssertions()
     }
 }
 
+//------------------------------------------------------------------------
+// PropagateExpansionAssertions: propagate assertions based on the
+//   expanded tree
+//
+// Notes:
+//    After the copy/init is expanded, we may see additional expansions
+//    to generate.
+//
+void MorphInitBlockHelper::PropagateExpansionAssertions()
+{
+    // Consider doing this for FieldByField as well
+    //
+    if (m_comp->optLocalAssertionProp && (m_transformationDecision == BlockTransformation::OneAsgBlock))
+    {
+        m_comp->optAssertionGen(m_asg);
+    }
+}
+
 //------------------------------------------------------------------------
 // PrepareSrc: Transform the asg src to an appropriate form and initialize member fields
 //    with information about it.