[DAG] Fix PR45049: LegalizeTypes crash

[jroelofs]

Summary

Sometimes LegalizeTypes knows about common subexpressions before SelectionDAG
does, leading to accidental SDValue removal before its reference count was
truly zero.

Fixes: https://bugs.llvm.org/show_bug.cgi?id=45049

[bjope]

FYI: Looks like this solves https://bugs.llvm.org/show_bug.cgi?id=44546 as well.

[tra]

Looks OK to me, but I'm not that familiar with the legalizer.

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

190

Nit: for (auto const &I: ReplacedValues)

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h

204–205

Do we need an assert at the beginning of the loop that Old != New? Otherwise we may be removing a valid map entry.

[jroelofs]

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h

204–205

Good idea. I believe that would only happen for an RAUW(N, N), which the DAGUpdateListener call sites guarantee won't happen.

[bjope]

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

190

Don't you need to wrap the whole for loop inside #ifndef NDEBUG. Otherwise I would be a def without any uses in a no-assert build.

Looking at the for loop above there is no such protection for UI (but maybe that is because UI is incremented in the loop).

I wouldn't mind if you added #ifndef NDEBUG around both loops, as they only make sense when the asserts are active.

[jroelofs]

Fix NDEBUG build.

[jroelofs]

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

190

Good catch!

[bjope]

LGTM!

[jroelofs]

https://github.com/llvm/llvm-project/commit/17673ae0b2cbf8d48973b673d413fb8591d8aae7

[jroelofs]

Reverted this in 3ce77142a6452d76d6f97c9a6c2da193e78841ba because one of the new assertions failed in the expensive-checks bots: http://lab.llvm.org:8011/builders/llvm-clang-x86_64-expensive-checks-debian/builds/4971/steps/test-check-all/logs/FAIL%3A%20LLVM%3A%3A2008-12-02-IllegalResultType.ll

This patch removes that check, which apparently wasn't correct.

[jroelofs]

@bjope are you happy with me re-landing this with that new assert dropped? Looks like that case is properly covered earlier in PerformExpensiveChecks() anyway.

[bjope]

In D76994#1975405, @jroelofs wrote:

@bjope are you happy with me re-landing this with that new assert dropped? Looks like that case is properly covered earlier in PerformExpensiveChecks() anyway.

Yes, that is fine with me!

[jroelofs]

re-landed in 41f13f1f64d2074ae7512fb23656c22585e912bd

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h

TableId getTableId(SDValue V) {
     assert(NextValueId != 0 &&
            "Ran out of Ids. Increase id type size or add compactification");
     return NextValueId - 1;
   }
 
   const SDValue &getSDValue(TableId &Id) {
     RemapId(Id);
     assert(Id && "TableId should be non-zero");
-    return IdToValueMap[Id];
+    auto I = IdToValueMap.find(Id);
+    assert(I != IdToValueMap.end() && "cannot find Id in map");
+    return I->second;
   }
 
 public:
   explicit DAGTypeLegalizer(SelectionDAG &dag)
     : TLI(dag.getTargetLoweringInfo()), DAG(dag),
     ValueTypeActions(TLI.getValueTypeActions()) {
     static_assert(MVT::LAST_VALUETYPE <= MVT::MAX_ALLOWED_VALUETYPE,
                   "Too many value types for ValueTypeActions to hold!");
   }
 
   /// This is the main entry point for the type legalizer.  This does a
   /// top-down traversal of the dag, legalizing types as it goes.  Returns
   /// "true" if it made any changes.
   bool run();
 
   void NoteDeletion(SDNode *Old, SDNode *New) {
+    assert(Old != New && "node replaced with self");
     for (unsigned i = 0, e = Old->getNumValues(); i != e; ++i) {
       TableId NewId = getTableId(SDValue(New, i));
       TableId OldId = getTableId(SDValue(Old, i));
 
-      if (OldId != NewId)
+      if (OldId != NewId) {
         ReplacedValues[OldId] = NewId;
 
-      // Delete Node from tables.
-      ValueToIdMap.erase(SDValue(Old, i));
+        // Delete Node from tables.  We cannot do this when OldId == NewId,
+        // because NewId can still have table references to it in
+        // ReplacedValues.
         IdToValueMap.erase(OldId);
         PromotedIntegers.erase(OldId);
         ExpandedIntegers.erase(OldId);
         SoftenedFloats.erase(OldId);
         PromotedFloats.erase(OldId);
         SoftPromotedHalfs.erase(OldId);
         ExpandedFloats.erase(OldId);
         ScalarizedVectors.erase(OldId);
         SplitVectors.erase(OldId);
         WidenedVectors.erase(OldId);
       }
+
+      ValueToIdMap.erase(SDValue(Old, i));
+    }
   }
 
   SelectionDAG &getDAG() const { return DAG; }
 
 private:
   SDNode *AnalyzeNewNode(SDNode *N);
   void AnalyzeNewValue(SDValue &Val);
   void PerformExpensiveChecks();

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

void DAGTypeLegalizer::PerformExpensiveChecks() {
 
   // Checked that NewNodes are only used by other NewNodes.
   for (unsigned i = 0, e = NewNodes.size(); i != e; ++i) {
     SDNode *N = NewNodes[i];
     for (SDNode::use_iterator UI = N->use_begin(), UE = N->use_end();
          UI != UE; ++UI)
       assert(UI->getNodeId() == NewNode && "NewNode used by non-NewNode!");
   }
+
+  // Check that all of the replacements still exist.
+  for (auto const &I : ReplacedValues)
+    assert(IdToValueMap.count(I.second) && "remapped value missing");
 }
 
 /// This is the main entry point for the type legalizer. This does a top-down
 /// traversal of the dag, legalizing types as it goes. Returns "true" if it made
 /// any changes.
 bool DAGTypeLegalizer::run() {
   bool Changed = false;
 

llvm/test/CodeGen/X86/legalize-types-remapid.ll

 ; RUN: llc -mtriple=i386 -mcpu=generic -O0 -o /dev/null %s
 
 @c = global i32 0
 @d = global <2 x i64> zeroinitializer
 
 define void @test() {
 bb1:
   %t0 = load <2 x i64>, <2 x i64>* @d
   %t0.i0 = extractelement <2 x i64> %t0, i32 0
   %t0.i0.cast = bitcast i64 %t0.i0 to <2 x i32>
   %t0.i0.cast.i0 = extractelement <2 x i32> %t0.i0.cast, i32 0
   store volatile i32 %t0.i0.cast.i0, i32* @c
   %t0.i0.cast.i1 = extractelement <2 x i32> %t0.i0.cast, i32 1
   store volatile i32 %t0.i0.cast.i1, i32* @c
   ret void
 }
+
+define void @PR45049() local_unnamed_addr {
+so_basic:
+  %a0 = load i1, i1* undef, align 1
+  %a1 = select i1 %a0, i542 4374501449566023848745004454235242730706338861786424872851541212819905998398751846447026354046107648, i542 0 ; constant is: i542 1 << 331
+  %a00 = zext i1 %a0 to i542
+  %a11 = shl i542 %a00, 331
+  %a2 = shl i542 %a00, 330
+  %a4 = or i542 %a1, %a2
+  %a05 = zext i1 %a0 to i488
+  %a55 = shl i488 %a05, 111
+  store i542 %a4, i542* undef, align 8
+  store i488 %a55, i488* undef, align 8
+  ret void
+}