[X86DomainReassignment] Don't compare stack-allocated values by address

[loladiro]

Summary

The Closure allocated in the main loop is allocated on the stack. However,
later in the code its address is taken (and used for comparisons). This
obviously doesn't work. In fact, the Closure will get the same stack address
during every loop iteration, rendering the check that intended to identify
Closure conflicts entirely ineffective. Fix this bug by giving every Closure
a unique ID and using that for comparison. Alternatively, we could heap
allocate the closure object.

Fixes PR37396
Fixes JuliaLang/julia#27032

[craig.topper]

lib/Target/X86/X86DomainReassignment.cpp

355

Won't this be unused in release builds? Will that cause an error?

357

Capitalize variable name

[loladiro]

lib/Target/X86/X86DomainReassignment.cpp

355

Fair point, I'll annotate this this with LLVM_DUMP_METHOD.

[loladiro]

Add LLVM_DUMP_METHOD to the added dump() method.

[craig.topper]

LGTM

lib/Target/X86/X86DomainReassignment.cpp

 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/Printable.h"
 #include <bitset>
 
 using namespace llvm;
 
 namespace llvm {
 void initializeX86DomainReassignmentPass(PassRegistry &);
 }
 
private:
   DenseSet<unsigned> Edges;
 
   /// Instructions in the closure.
   SmallVector<MachineInstr *, 8> Instrs;
 
   /// Domains which this closure can legally be reassigned to.
   std::bitset<NumDomains> LegalDstDomains;
 
+  /// An ID to uniquely identify this closure, even when it gets
+  /// moved around
+  unsigned ID;
+
 public:
-  Closure(std::initializer_list<RegDomain> LegalDstDomainList) {
+  Closure(unsigned ID, std::initializer_list<RegDomain> LegalDstDomainList) : ID(ID) {
     for (RegDomain D : LegalDstDomainList)
       LegalDstDomains.set(D);
   }
 
   /// Mark this closure as illegal for reassignment to all domains.
   void setAllIllegal() { LegalDstDomains.reset(); }
 
   /// \returns true if this closure has domains which are legal to reassign to.
public:
   void addInstruction(MachineInstr *I) {
     Instrs.push_back(I);
   }
 
   ArrayRef<MachineInstr *> instructions() const {
     return Instrs;
   }
 
+  LLVM_DUMP_METHOD void dump(const MachineRegisterInfo *MRI) const {
+    dbgs() << "Registers: ";
+    bool First = true;
+    for (unsigned Reg : Edges) {
+      if (!First)
+        dbgs() << ", ";
+      First = false;
+      dbgs() << printReg(Reg, MRI->getTargetRegisterInfo(), 0, MRI);
+    }
+    dbgs() << "\n" << "Instructions:";
+    for (MachineInstr *MI : Instrs) {
+      dbgs() << "\n  ";
+      MI->print(dbgs());
+    }
+    dbgs() << "\n";
+  }
+
+  unsigned getID() const {
+    return ID;
+  }
+
 };
 
 class X86DomainReassignment : public MachineFunctionPass {
   const X86Subtarget *STI;
   MachineRegisterInfo *MRI;
   const X86InstrInfo *TII;
 
   /// All edges that are included in some closure
   DenseSet<unsigned> EnclosedEdges;
 
   /// All instructions that are included in some closure.
-  DenseMap<MachineInstr *, Closure *> EnclosedInstrs;
+  DenseMap<MachineInstr *, unsigned> EnclosedInstrs;
 
 public:
   static char ID;
 
   X86DomainReassignment() : MachineFunctionPass(ID) {
     initializeX86DomainReassignmentPass(*PassRegistry::getPassRegistry());
   }
 
if (Domain != RD)
     return;
 
   Worklist.push_back(Reg);
 }
 
 void X86DomainReassignment::encloseInstr(Closure &C, MachineInstr *MI) {
   auto I = EnclosedInstrs.find(MI);
   if (I != EnclosedInstrs.end()) {
-    if (I->second != &C)
+    if (I->second != C.getID())
       // Instruction already belongs to another closure, avoid conflicts between
       // closure and mark this closure as illegal.
       C.setAllIllegal();
     return;
   }
 
-  EnclosedInstrs[MI] = &C;
+  EnclosedInstrs[MI] = C.getID();
   C.addInstruction(MI);
 
   // Mark closure as illegal for reassignment to domains, if there is no
   // converter for the instruction or if the converter cannot convert the
   // instruction.
   for (int i = 0; i != NumDomains; ++i) {
     if (C.isLegal((RegDomain)i)) {
       InstrConverterBase *IC = Converters.lookup({i, MI->getOpcode()});
bool X86DomainReassignment::runOnMachineFunction(MachineFunction &MF) {
   bool Changed = false;
 
   EnclosedEdges.clear();
   EnclosedInstrs.clear();
 
   std::vector<Closure> Closures;
 
   // Go over all virtual registers and calculate a closure.
+  unsigned ClosureID = 0;
   for (unsigned Idx = 0; Idx < MRI->getNumVirtRegs(); ++Idx) {
     unsigned Reg = TargetRegisterInfo::index2VirtReg(Idx);
 
     // GPR only current source domain supported.
     if (!isGPR(MRI->getRegClass(Reg)))
       continue;
 
     // Register already in closure.
     if (EnclosedEdges.count(Reg))
       continue;
 
     // Calculate closure starting with Reg.
-    Closure C({MaskDomain});
+    Closure C(ClosureID++, {MaskDomain});
     buildClosure(C, Reg);
 
     // Collect all closures that can potentially be converted.
     if (!C.empty() && C.isLegal(MaskDomain))
       Closures.push_back(std::move(C));
   }
 
-  for (Closure &C : Closures)
+  for (Closure &C : Closures) {
+    DEBUG(C.dump(MRI));
     if (isReassignmentProfitable(C, MaskDomain)) {
       reassign(C, MaskDomain);
       ++NumClosuresConverted;
       Changed = true;
     }
+  }
 
   DeleteContainerSeconds(Converters);
 
   LLVM_DEBUG(
       dbgs() << "***** Machine Function after Domain Reassignment *****\n");
   LLVM_DEBUG(MF.print(dbgs()));
 
   return Changed;

test/CodeGen/X86/domain-reassignment-test.ll

+; RUN: llc -mcpu=skylake-avx512 -mtriple=x86_64-unknown-linux-gnu %s -o - | FileCheck %s
+; RUN: llc -mcpu=skylake-avx512 -mtriple=x86_64-unknown-linux-gnu %s -o - | llvm-mc
+
+; Check that the X86 domain reassignment pass doesn't introduce an illegal
+; test instruction. See PR37396
+define void @japi1_foo2_34617() {
+pass2:
+  br label %if5
+
+L174:
+  %tmp = icmp sgt <2 x i64> undef, zeroinitializer
+  %tmp1 = icmp sle <2 x i64> undef, undef
+  %tmp2 = and <2 x i1> %tmp, %tmp1
+  %tmp3 = extractelement <2 x i1> %tmp2, i32 0
+  %tmp4 = extractelement <2 x i1> %tmp2, i32 1
+  %tmp106 = and i1 %tmp4, %tmp3
+  %tmp107 = zext i1 %tmp106 to i8
+  %tmp108 = and i8 %tmp122, %tmp107
+  %tmp109 = icmp eq i8 %tmp108, 0
+; CHECK-NOT: testb  {{%k[0-7]}}
+  br i1 %tmp109, label %L188, label %L190
+
+if5:
+  %b.055 = phi i8 [ 1, %pass2 ], [ %tmp122, %if5 ]
+  %tmp118 = icmp sgt i64 undef, 0
+  %tmp119 = icmp sle i64 undef, undef
+  %tmp120 = and i1 %tmp118, %tmp119
+  %tmp121 = zext i1 %tmp120 to i8
+  %tmp122 = and i8 %b.055, %tmp121
+  br i1 undef, label %L174, label %if5
+
+L188:
+  unreachable
+
+L190:
+  ret void
+}