[WebAssembly] Disable SimplifyDemandedVectorElts after legalization

[aheejin]

Summary

This fixes a reported bug that caused an infinite loop during the
SelectionDAG optimization phase in ISel, by creating an overridable hook
in TargetLowering that allows us to bail out from running
SimplifyDemandedVectorElts.

[tlively]

LGTM! This seems fairly reasonable as a solution, all things considered. Have you been able to check whether this has any broad impacts on the quality of our codegen? My guess is that it doesn't since it seems that any code that would have had the vector elements simplified previously would have ended up in that infinite loop, but we went a long time without ever seeing that bug in the wild.

[aheejin]

I wasn't able to check performances for real-world programs myself, but there's no changes in our existing tests. I don't think it will have a huge impact because I made the condition fairly specific: we bail out only when it is a build_vector and the DAG is already type-legalized and legalized, meaning it does so only in the very last stage before instruction selection. Also I'm not sure if there's any other way around this at this point.

llvm/include/llvm/CodeGen/TargetLowering.h

bool SimplifyDemandedVectorElts(SDValue Op, const APInt &DemandedEltMask,
                                   TargetLoweringOpt &TLO, unsigned Depth = 0,
                                   bool AssumeSingleUse = false) const;
 
   /// Helper wrapper around SimplifyDemandedVectorElts.
   /// Adds Op back to the worklist upon success.
   bool SimplifyDemandedVectorElts(SDValue Op, const APInt &DemandedElts,
                                   DAGCombinerInfo &DCI) const;
 
+  /// Return true if the target supports simplifying demanded vector elements by
+  /// converting them to undefs.
+  virtual bool
+  shouldSimplifyDemandedVectorElts(SDValue Op,
+                                   const TargetLoweringOpt &TLO) const {
+    return true;
+  }
+
   /// Determine which of the bits specified in Mask are known to be either zero
   /// or one and return them in the KnownZero/KnownOne bitsets. The DemandedElts
   /// argument allows us to only collect the known bits that are shared by the
   /// requested vector elements.
   virtual void computeKnownBitsForTargetNode(const SDValue Op,
                                              KnownBits &Known,
                                              const APInt &DemandedElts,
                                              const SelectionDAG &DAG,

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

bool TargetLowering::SimplifyDemandedVectorElts(
   EVT VT = Op.getValueType();
   unsigned Opcode = Op.getOpcode();
   APInt DemandedElts = OriginalDemandedElts;
   unsigned NumElts = DemandedElts.getBitWidth();
   assert(VT.isVector() && "Expected vector op");
 
   KnownUndef = KnownZero = APInt::getZero(NumElts);
 
+  const TargetLowering &TLI = TLO.DAG.getTargetLoweringInfo();
+  if (!TLI.shouldSimplifyDemandedVectorElts(Op, TLO))
+    return false;
+
   // TODO: For now we assume we know nothing about scalable vectors.
   if (VT.isScalableVector())
     return false;
 
   assert(VT.getVectorNumElements() == NumElts &&
          "Mask size mismatches value type element count!");
 
   // Undef operand.

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.h

private:
 
   void ReplaceNodeResults(SDNode *N, SmallVectorImpl<SDValue> &Results,
                           SelectionDAG &DAG) const override;
 
   const char *getClearCacheBuiltinName() const override {
     report_fatal_error("llvm.clear_cache is not supported on wasm");
   }
 
+  bool
+  shouldSimplifyDemandedVectorElts(SDValue Op,
+                                   const TargetLoweringOpt &TLO) const override;
+
   // Custom lowering hooks.
   SDValue LowerOperation(SDValue Op, SelectionDAG &DAG) const override;
   SDValue LowerFrameIndex(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerRETURNADDR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerFRAMEADDR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerGlobalAddress(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerGlobalTLSAddress(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerExternalSymbol(SDValue Op, SelectionDAG &DAG) const;

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

if (VT.isFixedLengthVector()) {
     if (EltVT == MVT::i8 || EltVT == MVT::i16 || EltVT == MVT::i32 ||
         EltVT == MVT::i64 || EltVT == MVT::f32 || EltVT == MVT::f64)
       return TypeWidenVector;
   }
 
   return TargetLoweringBase::getPreferredVectorAction(VT);
 }
 
+bool WebAssemblyTargetLowering::shouldSimplifyDemandedVectorElts(
+    SDValue Op, const TargetLoweringOpt &TLO) const {
+  // ISel process runs DAGCombiner after legalization; this step is called
+  // SelectionDAG optimization phase. This post-legalization combining process
+  // runs DAGCombiner on each node, and if there was a change to be made,
+  // re-runs legalization again on it and its user nodes to make sure
+  // everythiing is in a legalized state.
+  //
+  // The legalization calls lowering routines, and we do our custom lowering for
+  // build_vectors (LowerBUILD_VECTOR), which converts undef vector elements
+  // into zeros. But there is a set of routines in DAGCombiner that turns unused
+  // (= not demanded) nodes into undef, among which SimplifyDemandedVectorElts
+  // turns unused vector elements into undefs. But this routine does not work
+  // with our custom LowerBUILD_VECTOR, which turns undefs into zeros. This
+  // combination can result in a infinite loop, in which undefs are converted to
+  // zeros in legalization and back to undefs in combining.
+  //
+  // So after DAG is legalized, we prevent SimplifyDemandedVectorElts from
+  // running for build_vectors.
+  if (Op.getOpcode() == ISD::BUILD_VECTOR && TLO.LegalOps && TLO.LegalTys)
+    return false;
+  return true;
+}
+
 //===----------------------------------------------------------------------===//
 // WebAssembly Lowering private implementation.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Lowering Code
 //===----------------------------------------------------------------------===//
 

llvm/test/CodeGen/WebAssembly/simd-simplify-demanded-vector-elts.ll

+; RUN: llc < %s -mattr=+simd128 -verify-machineinstrs
+
+target triple = "wasm32-unknown-unknown"
+
+; After DAG legalization, in SelectionDAG optimization phase, ISel runs
+; DAGCombiner on each node, among which SimplifyDemandedVectorElts turns unused
+; vector elements into undefs. And in order to make sure the DAG is in a
+; legalized state, it runs legalization again, which runs our custom
+; LowerBUILD_VECTOR, which converts undefs into zeros, causing an infinite loop.
+; We prevent this from happening by creating a custom hook , which allows us to
+; bail out of SimplifyDemandedVectorElts after legalization.
+
+; This is a reduced test case from a bug reproducer reported. This should not
+; hang.
+define void @test(i8 %0) {
+  %2 = insertelement <4 x i8> <i8 -1, i8 -1, i8 -1, i8 poison>, i8 %0, i64 3
+  %3 = zext <4 x i8> %2 to <4 x i32>
+  %4 = mul nuw nsw <4 x i32> %3, <i32 257, i32 257, i32 257, i32 257>
+  %5 = add nuw nsw <4 x i32> %4, <i32 1, i32 1, i32 1, i32 1>
+  %6 = lshr <4 x i32> %5, <i32 1, i32 1, i32 1, i32 1>
+  %7 = mul nuw nsw <4 x i32> %6, <i32 20000, i32 20000, i32 20000, i32 20000>
+  %8 = add nuw nsw <4 x i32> %7, <i32 32768, i32 32768, i32 32768, i32 32768>
+  %9 = and <4 x i32> %8, <i32 2147418112, i32 2147418112, i32 2147418112, i32 2147418112>
+  %10 = sub nsw <4 x i32> <i32 655360000, i32 655360000, i32 655360000, i32 655360000>, %9
+  %11 = ashr exact <4 x i32> %10, <i32 16, i32 16, i32 16, i32 16>
+  %12 = trunc <4 x i32> %11 to <4 x i16>
+  store <4 x i16> %12, <4 x i16>* undef, align 4
+  ret void
+}