[llvm] [InstCombine] Extend bitmask mul combine to handle independent operands (PR #142503)

[Jeffrey Byrnes via llvm-commits]

================
@@ -3659,6 +3665,88 @@ static std::optional<DecomposedBitMaskMul> matchBitmaskMul(Value *V) {
   return std::nullopt;
 }
 
+/// (A & N) * C + (A & M) * C -> (A & (N + M)) & C
+/// This also accepts the equivalent select form of (A & N) * C
+/// expressions i.e. !(A & N) ? 0 : N * C)
+static Value *foldBitmaskMul(Value *Op0, Value *Op1,
+                             InstCombiner::BuilderTy &Builder) {
+  auto Decomp1 = matchBitmaskMul(Op1);
+
+  if (Decomp1) {
+    auto Decomp0 = matchBitmaskMul(Op0);
+
+    if (Decomp0) {
+      // If we have independent operands in the BitmaskMul chain, then just
+      // reassociate to encourage combining in future iterations.
+
+      if (Decomp0->isCombineableWith(*Decomp1)) {
+        auto NewAnd = Builder.CreateAnd(
+            Decomp0->X, ConstantInt::get(Decomp0->X->getType(),
+                                         (Decomp0->Mask + Decomp1->Mask)));
+
+        auto Res = Builder.CreateMul(
+            NewAnd, ConstantInt::get(NewAnd->getType(), Decomp1->Factor), "",
+            Decomp0->NUW && Decomp1->NUW, Decomp0->NSW && Decomp1->NSW);
+        return Res;
+      }
+    }
+  }
+
+  return nullptr;
+}
+
+Value *InstCombinerImpl::foldDisjointOr(Value *LHS, Value *RHS) {
+  if (Value *Res = foldBitmaskMul(LHS, RHS, Builder)) {
+    return Res;
+  }
+
+  return nullptr;
+}
+
+Value *InstCombinerImpl::reassociateDisjointOr(Value *LHS, Value *RHS) {
+
+  Value *X, *Y;
+  if (match(RHS, m_OneUse(m_DisjointOr(m_Value(X), m_Value(Y))))) {
+    if (Value *Res = foldDisjointOr(LHS, X))
+      return Builder.CreateOr(Res, Y, "", /*IsDisjoint=*/true);
+    if (Value *Res = foldDisjointOr(LHS, Y))
+      return Builder.CreateOr(Res, X, "", /*IsDisjoint=*/true);
+  }
+
+  if (match(LHS, m_OneUse(m_DisjointOr(m_Value(X), m_Value(Y))))) {
+    if (Value *Res = foldDisjointOr(X, RHS))
+      return Builder.CreateOr(Res, Y, "", /*IsDisjoint=*/true);
+    if (Value *Res = foldDisjointOr(Y, RHS))
+      return Builder.CreateOr(Res, X, "", /*IsDisjoint=*/true);
+  }
+
+  Value *X1, *Y1;
+  if (match(LHS, m_OneUse(m_DisjointOr(m_Value(X), m_Value(Y)))) &&
+      (match(RHS, m_OneUse(m_DisjointOr(m_Value(X1), m_Value(Y1)))))) {
----------------
jrbyrnes wrote:

My original implementation captured this case. `matchCombinedBitmaskMul` would take an operand from an or-disjoint and try to `matchBitmaskMul` on it. If there was no such match, and the operand was itself an or-disjoint, it would try to find a `matchBitmaskMul` in either of the child operands. Thus, we could end up combining two grandchildren from the original or-disjoint.

However, rerunning the motivating workload through its default pipeline indicates that this set of permutations / reassociations is not necessary, and we provide the optimization with the standard reassociations.

I think it's fine to just include the standard reassociations, and potentially extend something if we find we have left some on the table.

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