[LegalizeVectorOps] Pass the post-UpdateNodeOperands version of the Node to the LowerOperation/PromoteNode/ExpandNode calls

[craig.topper]

Summary

Its possible though very unlikely that the UpdateNodeOperands call that is made after legalizing operands will CSE to an existing node. When this happens we end up with Result pointing to a different node than Op. When this happens, in order to make sure the legalized operands propagate correctly, we need to use Result in the LowerOperation, ExpandNode, or PromoteNode calls made later.

While this fix works, I think this may result in Result being legalized twice because we aren't registering it in the LegalizedNodes map as itself. Maybe another option is to specifically detect this case and recursively call LegalizeOp on the node, then just call TranslateLegalizeResults to register it for Op?

[craig.topper]

Update a debug message as well.

[efriedma]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp

270

Same issue here?

295

Same issue here?

475

Looks like we got this right in the other calls to LowerOperation, just not this one. Odd.

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorOps.cpp

if (LD->getMemoryVT().isVector() && ExtType != ISD::NON_EXTLOAD) {
             Lowered = LegalizeOp(Lowered);
             Changed = true;
           }
           return TranslateLegalizeResults(Op, Lowered);
         }
         LLVM_FALLTHROUGH;
       case TargetLowering::Expand:
         Changed = true;
         return ExpandLoad(Op);
       }
     }
   } else if (Op.getOpcode() == ISD::STORE) {
     StoreSDNode *ST = cast<StoreSDNode>(Op.getNode());
     EVT StVT = ST->getMemoryVT();
     MVT ValVT = ST->getValue().getSimpleValueType();
     if (StVT.isVector() && ST->isTruncatingStore()) {
       LLVM_DEBUG(dbgs() << "\nLegalizing truncating vector store: ";
                  Node->dump(&DAG));
       switch (TLI.getTruncStoreAction(ValVT, StVT)) {
       default: llvm_unreachable("This action is not supported yet!");
       case TargetLowering::Legal:
         return TranslateLegalizeResults(Op, Result);
       case TargetLowering::Custom: {
         SDValue Lowered = TLI.LowerOperation(Result, DAG);
         if (Lowered != Result) {
           // Make sure the new code is also legal.
           Lowered = LegalizeOp(Lowered);
           Changed = true;
         }
         return TranslateLegalizeResults(Op, Lowered);
       }
       case TargetLowering::Expand:
         Changed = true;
         return ExpandStore(Op);
       }
     }
   }
 
   bool HasVectorValueOrOp = false;
   for (auto J = Node->value_begin(), E = Node->value_end(); J != E; ++J)
     HasVectorValueOrOp |= J->isVector();
   for (const SDValue &Op : Node->op_values())
#include "llvm/IR/ConstrainedOps.def"
   case ISD::VECREDUCE_FMUL:
   case ISD::VECREDUCE_FMAX:
   case ISD::VECREDUCE_FMIN:
     Action = TLI.getOperationAction(Node->getOpcode(),
                                     Node->getOperand(0).getValueType());
     break;
   }
 
-  LLVM_DEBUG(dbgs() << "\nLegalizing vector op: "; Node->dump(&DAG));
+  LLVM_DEBUG(dbgs() << "\nLegalizing vector op: "; Result->dump(&DAG));
 
   switch (Action) {
   default: llvm_unreachable("This action is not supported yet!");
   case TargetLowering::Promote:
-    Result = Promote(Op);
+    Result = Promote(Result);
     Changed = true;
     break;
   case TargetLowering::Legal:
     LLVM_DEBUG(dbgs() << "Legal node: nothing to do\n");
     break;
   case TargetLowering::Custom: {
     LLVM_DEBUG(dbgs() << "Trying custom legalization\n");
-    if (SDValue Tmp1 = TLI.LowerOperation(Op, DAG)) {
+    if (SDValue Tmp1 = TLI.LowerOperation(Result, DAG)) {
       LLVM_DEBUG(dbgs() << "Successfully custom legalized node\n");
       Result = Tmp1;
       break;
     }
     LLVM_DEBUG(dbgs() << "Could not custom legalize node\n");
     LLVM_FALLTHROUGH;
   }
   case TargetLowering::Expand:
-    Result = Expand(Op);
+    Result = Expand(Result);
   }
 
   // Make sure that the generated code is itself legal.
   if (Result != Op) {
     Result = LegalizeOp(Result);
     Changed = true;
   }
 

llvm/test/CodeGen/X86/avx512-cmp.ll

; ALL-NEXT:    retq
   br i1 %xor1, label %if.end.i, label %if.then.i
 
 if.then.i:
  ret i32 5
 
 if.end.i:
   ret i32 6
 }
+
+; This test previously caused an infinite loop in legalize vector ops. Due to
+; CSE triggering on the call to UpdateNodeOperands and the resulting node not
+; being passed to LowerOperation. The add is needed to force the zext into a
+; sext on that path. The shuffle keeps the zext alive. The xor somehow
+; influences the zext to be visited before the sext exposing the CSE opportunity
+; for the sext since zext of setcc is custom legalized to a sext and shift.
+define <8 x i32> @legalize_loop(<8 x double> %arg) {
+; KNL-LABEL: legalize_loop:
+; KNL:       ## %bb.0:
+; KNL-NEXT:    vxorpd %xmm1, %xmm1, %xmm1
+; KNL-NEXT:    vcmpnltpd %zmm0, %zmm1, %k1
+; KNL-NEXT:    vpternlogd $255, %zmm0, %zmm0, %zmm0 {%k1} {z}
+; KNL-NEXT:    vpsrld $31, %ymm0, %ymm1
+; KNL-NEXT:    vpshufd {{.*#+}} ymm1 = ymm1[3,2,1,0,7,6,5,4]
+; KNL-NEXT:    vpermq {{.*#+}} ymm1 = ymm1[2,3,0,1]
+; KNL-NEXT:    vpsubd %ymm0, %ymm1, %ymm0
+; KNL-NEXT:    retq
+;
+; SKX-LABEL: legalize_loop:
+; SKX:       ## %bb.0:
+; SKX-NEXT:    vxorpd %xmm1, %xmm1, %xmm1
+; SKX-NEXT:    vcmpnltpd %zmm0, %zmm1, %k0
+; SKX-NEXT:    vpmovm2d %k0, %ymm0
+; SKX-NEXT:    vpsrld $31, %ymm0, %ymm1
+; SKX-NEXT:    vpshufd {{.*#+}} ymm1 = ymm1[3,2,1,0,7,6,5,4]
+; SKX-NEXT:    vpermq {{.*#+}} ymm1 = ymm1[2,3,0,1]
+; SKX-NEXT:    vpsubd %ymm0, %ymm1, %ymm0
+; SKX-NEXT:    retq
+  %tmp = fcmp ogt <8 x double> %arg, zeroinitializer
+  %tmp1 = xor <8 x i1> %tmp, <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true>
+  %tmp2 = zext <8 x i1> %tmp1 to <8 x i32>
+  %tmp3 = shufflevector <8 x i32> %tmp2, <8 x i32> undef, <8 x i32> <i32 7, i32 6, i32 5, i32 4, i32 3, i32 2, i32 1, i32 0>
+  %tmp4 = add <8 x i32> %tmp2, %tmp3
+  ret <8 x i32> %tmp4
+}