Add support for checking and dropping poison-generating flags

Author: dc03-work

llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h

@@ -34,6 +34,83 @@ class GenericMachineInstr : public MachineInstr {
   static bool classof(const MachineInstr *MI) {
     return isPreISelGenericOpcode(MI->getOpcode());
   }
+
+  bool hasPoisonGeneratingFlags() const {
+    switch (getOpcode()) {
+    case TargetOpcode::G_ADD:
+    case TargetOpcode::G_SUB:
+    case TargetOpcode::G_MUL:
+    case TargetOpcode::G_SHL:
+    case TargetOpcode::G_TRUNC:
+      return getFlag(NoUWrap) || getFlag(NoSWrap);
+
+    case TargetOpcode::G_UDIV:
+    case TargetOpcode::G_SDIV:
+    case TargetOpcode::G_ASHR:
+    case TargetOpcode::G_LSHR:
+      return getFlag(IsExact);
+
+    case TargetOpcode::G_OR:
+      return getFlag(Disjoint);
+
+    case TargetOpcode::G_UITOFP:
+    case TargetOpcode::G_ZEXT:
+      return getFlag(NonNeg);
+
+    case TargetOpcode::G_FNEG:
+    case TargetOpcode::G_FADD:
+    case TargetOpcode::G_FSUB:
+    case TargetOpcode::G_FMUL:
+    case TargetOpcode::G_FDIV:
+    case TargetOpcode::G_FREM:
+    case TargetOpcode::G_FCMP:
+      return getFlag(FmNoNans) || getFlag(FmNoInfs);
+
+    default:
+      return false;
+    }
+  }
+
+  void dropPoisonGeneratingFlags() {
+    switch (getOpcode()) {
+    case TargetOpcode::G_ADD:
+    case TargetOpcode::G_SUB:
+    case TargetOpcode::G_MUL:
+    case TargetOpcode::G_SHL:
+    case TargetOpcode::G_TRUNC:
+      clearFlag(NoUWrap);
+      clearFlag(NoSWrap);
+      break;
+
+    case TargetOpcode::G_UDIV:
+    case TargetOpcode::G_SDIV:
+    case TargetOpcode::G_ASHR:
+    case TargetOpcode::G_LSHR:
+      clearFlag(IsExact);
+      break;
+
+    case TargetOpcode::G_OR:
+      clearFlag(Disjoint);
+      break;
+
+    case TargetOpcode::G_UITOFP:
+    case TargetOpcode::G_ZEXT:
+      clearFlag(NonNeg);
+      break;
+
+    case TargetOpcode::G_FNEG:
+    case TargetOpcode::G_FADD:
+    case TargetOpcode::G_FSUB:
+    case TargetOpcode::G_FMUL:
+    case TargetOpcode::G_FDIV:
+    case TargetOpcode::G_FREM:
+    case TargetOpcode::G_FCMP:
+      clearFlag(FmNoNans);
+      clearFlag(FmNoInfs);
+      break;
+    }
+    assert(!hasPoisonGeneratingFlags());
+  }
 };
 
 /// Provides common memory operand functionality.

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

@@ -234,12 +234,13 @@ bool CombinerHelper::matchFreezeOfSingleMaybePoisonOperand(
 
   MachineInstr *OrigDef = MRI.getUniqueVRegDef(OrigOp);
   // Avoid trying to fold G_PHI, G_UNMERGE_VALUES, G_FREEZE (the latter is
-  // handled by idempotent_prop)
+  // handled by idempotent_prop).
   if (!OrigDef || OrigDef->isPHI() || isa<GUnmerge>(OrigDef) ||
       isa<GFreeze>(OrigDef))
     return false;
 
-  if (canCreateUndefOrPoison(OrigOp, MRI))
+  if (canCreateUndefOrPoison(OrigOp, MRI,
+                             /*ConsiderFlagsAndMetadata*/ false))
     return false;
 
   std::optional<MachineOperand> MaybePoisonOperand;
@@ -258,6 +259,8 @@ bool CombinerHelper::matchFreezeOfSingleMaybePoisonOperand(
       return false;
   }
 
+  cast<GenericMachineInstr>(OrigDef)->dropPoisonGeneratingFlags();
+
   // Eliminate freeze if all operands are guaranteed non-poison.
   if (!MaybePoisonOperand) {
     MatchInfo = [=](MachineIRBuilder &B) { MRI.replaceRegWith(DstOp, OrigOp); };

llvm/lib/CodeGen/GlobalISel/Utils.cpp

@@ -1733,6 +1733,15 @@ static bool canCreateUndefOrPoison(Register Reg, const MachineRegisterInfo &MRI,
                                    UndefPoisonKind Kind) {
   MachineInstr *RegDef = MRI.getVRegDef(Reg);
 
+  if (auto *GMI = dyn_cast<GenericMachineInstr>(RegDef)) {
+    if (ConsiderFlagsAndMetadata && includesPoison(Kind) &&
+        GMI->hasPoisonGeneratingFlags())
+      return true;
+  }
+  // Conservatively return true.
+  else
+    return true;
+
   switch (RegDef->getOpcode()) {
   case TargetOpcode::G_FREEZE:
     return false;