From 7cdc60c3db1ed96f3d976ad913709c6c83776f3b Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper@intel.com>
Date: Fri, 3 Jan 2020 11:52:18 -0800
Subject: [PATCH] [LegalizeVectorOps] Pass the post-UpdateNodeOperands version
 of Op to ExpandLoad/ExpandStore

UpdateNodeOperands might CSE to another existing node. So we should make sure we're legalizing that node otherwise we might fail to hook up the operands properly. I've moved the result registration up to the caller to avoid having to pass both Result and Op into the functions where it might be confusing which is which.

This address 2 other issues pointed out in D71861.

Differential Revision: https://reviews.llvm.org/D72021
---
 .../SelectionDAG/LegalizeVectorOps.cpp        | 25 +++++++++++--------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp
index 5d76aeb07995d..c4287b09780c6 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp
@@ -130,7 +130,7 @@ class VectorLegalizer {
   /// supported by the target.
   SDValue ExpandVSELECT(SDValue Op);
   SDValue ExpandSELECT(SDValue Op);
-  SDValue ExpandLoad(SDValue Op);
+  std::pair<SDValue, SDValue> ExpandLoad(SDValue Op);
   SDValue ExpandStore(SDValue Op);
   SDValue ExpandFNEG(SDValue Op);
   SDValue ExpandFSUB(SDValue Op);
@@ -265,9 +265,13 @@ SDValue VectorLegalizer::LegalizeOp(SDValue Op) {
           return TranslateLegalizeResults(Op, Lowered);
         }
         LLVM_FALLTHROUGH;
-      case TargetLowering::Expand:
+      case TargetLowering::Expand: {
         Changed = true;
-        return ExpandLoad(Op);
+        std::pair<SDValue, SDValue> Tmp = ExpandLoad(Result);
+        AddLegalizedOperand(Op.getValue(0), Tmp.first);
+        AddLegalizedOperand(Op.getValue(1), Tmp.second);
+        return Op.getResNo() ? Tmp.first : Tmp.second;
+      }
       }
     }
   } else if (Op.getOpcode() == ISD::STORE) {
@@ -290,9 +294,12 @@ SDValue VectorLegalizer::LegalizeOp(SDValue Op) {
         }
         return TranslateLegalizeResults(Op, Lowered);
       }
-      case TargetLowering::Expand:
+      case TargetLowering::Expand: {
         Changed = true;
-        return ExpandStore(Op);
+        SDValue Chain = ExpandStore(Result);
+        AddLegalizedOperand(Op, Chain);
+        return Chain;
+      }
       }
     }
   }
@@ -633,7 +640,7 @@ SDValue VectorLegalizer::PromoteFP_TO_INT(SDValue Op) {
   return Promoted;
 }
 
-SDValue VectorLegalizer::ExpandLoad(SDValue Op) {
+std::pair<SDValue, SDValue> VectorLegalizer::ExpandLoad(SDValue Op) {
   LoadSDNode *LD = cast<LoadSDNode>(Op.getNode());
 
   EVT SrcVT = LD->getMemoryVT();
@@ -760,16 +767,12 @@ SDValue VectorLegalizer::ExpandLoad(SDValue Op) {
     std::tie(Value, NewChain) = TLI.scalarizeVectorLoad(LD, DAG);
   }
 
-  AddLegalizedOperand(Op.getValue(0), Value);
-  AddLegalizedOperand(Op.getValue(1), NewChain);
-
-  return (Op.getResNo() ? NewChain : Value);
+  return std::make_pair(Value, NewChain);
 }
 
 SDValue VectorLegalizer::ExpandStore(SDValue Op) {
   StoreSDNode *ST = cast<StoreSDNode>(Op.getNode());
   SDValue TF = TLI.scalarizeVectorStore(ST, DAG);
-  AddLegalizedOperand(Op, TF);
   return TF;
 }