[llvm-branch-commits] [llvm] [AMDGPU] Fold fmul with zero checks to fmul.legacy (PR #229357)

[Steffen Larsen via llvm-branch-commits]

================
@@ -1578,11 +1581,85 @@ bool AMDGPUCodeGenPrepareImpl::visitLoadInst(LoadInst &I) {
   return false;
 }
 
+// select (X == 0 || Y == 0), +0.0, (X * Y) -> fmul_legacy(X, Y)
+bool AMDGPUCodeGenPrepareImpl::foldSelectToFMulLegacy(SelectInst &I) {
+  if (!I.getType()->isFloatTy())
+    return false;
+
+  // Inverted form: select (X != 0 && Y != 0), X * Y, +0.0.
+  bool Inverted = match(I.getFalseValue(), m_PosZeroFP());
+  if (!Inverted && !match(I.getTrueValue(), m_PosZeroFP()))
+    return false;
+
+  Value *X, *Y;
+  if (!match(Inverted ? I.getTrueValue() : I.getFalseValue(),
+             m_OneUse(m_FMul(m_Value(X), m_Value(Y)))))
+    return false;
+
+  SmallVector<Value *, 4> Roots[2] = {{I.getCondition()}, {X, Y}};
+  SmallPtrSet<Value *, 4> Leaves[2];
+  for (auto [Worklist, Found] : zip(Roots, Leaves)) {
+    SmallPtrSet<Value *, 8> Visited;
+
+    while (!Worklist.empty()) {
+      Value *Cur = Worklist.pop_back_val();
+      if (!Visited.insert(Cur).second)
+        continue;
+
+      Value *Src, *LHS, *RHS;
+      if (Cur->getType()->isIntegerTy(1)) {
+        if (Inverted ? match(Cur, m_And(m_Value(LHS), m_Value(RHS)))
+                     : match(Cur, m_Or(m_Value(LHS), m_Value(RHS)))) {
+          Worklist.append({LHS, RHS});
+          continue;
+        }
+
+        if (!match(Cur, m_SpecificFCmp(Inverted ? FCmpInst::FCMP_UNE
+                                                : FCmpInst::FCMP_OEQ,
+                                       m_Value(Src), m_AnyZeroFP())))
+          return false;
+
+        Worklist.push_back(Src);
+        continue;
+      }
----------------
steffenlarsen wrote:

This seems like something that should only ever be applied to `I.getCondition()`, is that right? If so, could we do it as pre-processing of `Roots[0]`?

https://github.com/llvm/llvm-project/pull/229357