[llvm] [RISCV] Reverse (add x, (zext c)) back to (select c, (add x, 1), x) (PR #87236)

[Luke Lau via llvm-commits]

================
@@ -52,12 +52,58 @@ class RISCVCodeGenPrepare : public FunctionPass,
   }
 
   bool visitInstruction(Instruction &I) { return false; }
+  bool visitBinaryOperator(BinaryOperator &BO);
   bool visitAnd(BinaryOperator &BO);
   bool visitIntrinsicInst(IntrinsicInst &I);
 };
 
 } // end anonymous namespace
 
+/// InstCombine will canonicalize selects of binary ops where the identity is
+/// zero to zexts:
+///
+/// select c, (add x, 1), x -> add x, (zext c)
+///
+/// On RISC-V though, a zext of an i1 vector will be lowered as a vmv.v.i and a
+/// vmerge.vim:
+///
+///       vmv.v.i v12, 0
+///       vmerge.vim      v9, v12, 1, v0
+///       vadd.vv v8, v8, v9
+///
+/// Reverse this transform so that we pull the select outside of the binary op,
+/// which allows us to fold it into a masked op:
+///
+///       vadd.vi v8, v8, 1, v0.t
+bool RISCVCodeGenPrepare::visitBinaryOperator(BinaryOperator &BO) {
+  if (!BO.getType()->isVectorTy())
+    return false;
+
+  // TODO: We could allow sub if we did a non-commutative match
+  Constant *Identity = ConstantExpr::getIdentity(&BO, BO.getType());
+  if (!Identity || !Identity->isZeroValue())
+    return false;
+
+  using namespace PatternMatch;
+
+  Value *Mask, *RHS;
+  if (!match(&BO, m_c_BinOp(m_OneUse(m_ZExt(m_Value(Mask))), m_Value(RHS))))
+    return false;
+
+  if (!Mask->getType()->isIntOrIntVectorTy(1))
+    return false;
+
+  IRBuilder<> Builder(&BO);
+  Value *Splat = ConstantInt::get(BO.getType(), 1);
+  Value *BinOp = Builder.CreateBinOp(BO.getOpcode(), RHS, Splat);
----------------
lukel97 wrote:

Doesn't look like it's needed for correctness, but probably better to copy them to retain the information: https://alive2.llvm.org/ce/z/9PvEdE

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