[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());
----------------
lukel97 wrote:

I've left the sub case for later, since I believe InstCombine canonicalizes `sub x, (zext c)` to `add x, (sext c)` anyway. Definitely something that the loop vectorizer could be emitting though

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