[Constants, InstCombine] allow RHS (operand 1) identity constants for binops

[spatel]

Summary

This is a follow-up to D48830 to allow more shuffle folding using identity constants. Without this patch, we're only returning identity constants for commutative binops, so we miss opportunities to fold shifts and other binops.

LLVM binary opcode review: there are a total of 18 binops. There are 7 commutative binops (add, mul, and, or, xor, fadd, fmul) which we already fold. We're able to fold 4 more opcodes with this patch (shl, lshr, ashr, fdiv). Loosening the div/rem constraint in the shuffle fold will allow 2 more opcodes (udiv/sdiv). There are no folds for srem/urem/frem AFAIK. We don't bother with sub/fsub with constant operand 1 because those are canonicalized to add/fadd. 7 + 4 + 2 + 3 + 2 = 18.

Note that the more flexible definition for the identity constant shown here almost matches the logic of getDefaultConstantForOpcode() in D28907, so that patch could be reduced assuming this is ok and once we figure out how to handle -0.0 with fadd.

[lebedev.ri]

I don't think this is wrong, but does this propagation of undef not make the new versions more poisonous?

[spatel]

In D48893#1151303, @lebedev.ri wrote:

I don't think this is wrong, but does this propagation of undef not make the new versions more poisonous?

That's a good question. My understanding is that no, there is no more poison here than before because in the original code the shuffle's undef mask element means that lane of the vector is undef, so nothing downstream that might be using that lane can be valid. After this transform, we're performing some binop with an undef operand, so again the result of that lane is undef, and nothing downstream can be any different.

This is ok for everything except integer div/rem because they're the only ops with immediate UB. cc'ing @nlopes @sanjoy @aqjune to see if that reasoning is bogus.

[aqjune]

I think it depends on what shl i32 %x, undef is - undef can be concretized into an integer larger than 31, and according to https://llvm.org/docs/LangRef.html#id134 , it is poison, so conversion like the one in shuffle_select.ll may be invalid.
Defining shl with large second operand as poison explains another kind of optimizations, IMHO: For example, InstCombine optimizes (C2 << X) << C1 into (C2 << C1) << X. Defining shl with large shift as undef does not explain this optimimzation - If C1 = 4, C2 = 1 and X = 33, source is (1 << 33) << 4 = undef << 4 = (undef with low 4 bits zero), but after optimization it is (1 << 4) << 33 = undef. If large shift yields poison, this is explained, because poison << 4 = poison.

[nlopes]

shl %x, undef is poison; we don't want more undefs :)
So this transformation for shifts is not correct. The way to make it correct would be to introduce a poison value and use it in the shuffle instructions instead of undef. I suggest we finally go ahead and do that.

[spatel]

In D48893#1151840, @nlopes wrote:

shl %x, undef is poison; we don't want more undefs :)
So this transformation for shifts is not correct. The way to make it correct would be to introduce a poison value and use it in the shuffle instructions instead of undef. I suggest we finally go ahead and do that.

1. Does the 'undef' in the shuffle mask represent something different than the 'undef' in the shift amount?
2. Are we distinguishing shifts from the other binary opcodes here? Or is this existing transform also incorrect?

%b = mul nsw nuw <4 x i32> %v, <i32 11, i32 12, i32 13, i32 14>
%s = shufflevector <4 x i32> %v, <4 x i32> %b, <4 x i32> <i32 undef, i32 5, i32 2, i32 7>
-->
%s = mul nuw nsw <4 x i32> [[V:%.*]], <i32 undef, i32 12, i32 1, i32 14>

[nlopes]

In D48893#1152112, @spatel wrote:

In D48893#1151840, @nlopes wrote:

shl %x, undef is poison; we don't want more undefs :)
So this transformation for shifts is not correct. The way to make it correct would be to introduce a poison value and use it in the shuffle instructions instead of undef. I suggest we finally go ahead and do that.

1. Does the 'undef' in the shuffle mask represent something different than the 'undef' in the shift amount?

No: undef means any value in both cases.
The deal is that x << y is poison if y >= bitwidth. Since undef can be such a number, x << undef is poison.
Why is it poison? Because different CPU architectures (e.g. x86 and ARM) deal with this case differently, so we don't really want to define it. C/C++ allows to do that.

Shuffle is never poison if its inputs are not poison.

2. Are we distinguishing shifts from the other binary opcodes here? Or is this existing transform also incorrect?

%b = mul nsw nuw <4 x i32> %v, <i32 11, i32 12, i32 13, i32 14>
%s = shufflevector <4 x i32> %v, <4 x i32> %b, <4 x i32> <i32 undef, i32 5, i32 2, i32 7>
-->
%s = mul nuw nsw <4 x i32> [[V:%.*]], <i32 undef, i32 12, i32 1, i32 14>

This transformation is also incorrect, unfortunately. If you drop nsw & nuw, it's correct.

[spatel]

In D48893#1152643, @nlopes wrote:

In D48893#1152112, @spatel wrote:

In D48893#1151840, @nlopes wrote:

shl %x, undef is poison; we don't want more undefs :)
So this transformation for shifts is not correct. The way to make it correct would be to introduce a poison value and use it in the shuffle instructions instead of undef. I suggest we finally go ahead and do that.

1. Does the 'undef' in the shuffle mask represent something different than the 'undef' in the shift amount?

No: undef means any value in both cases.
The deal is that x << y is poison if y >= bitwidth. Since undef can be such a number, x << undef is poison.
Why is it poison? Because different CPU architectures (e.g. x86 and ARM) deal with this case differently, so we don't really want to define it. C/C++ allows to do that.

Shuffle is never poison if its inputs are not poison.

2. Are we distinguishing shifts from the other binary opcodes here? Or is this existing transform also incorrect?

%b = mul nsw nuw <4 x i32> %v, <i32 11, i32 12, i32 13, i32 14>
%s = shufflevector <4 x i32> %v, <4 x i32> %b, <4 x i32> <i32 undef, i32 5, i32 2, i32 7>
-->
%s = mul nuw nsw <4 x i32> [[V:%.*]], <i32 undef, i32 12, i32 1, i32 14>

This transformation is also incorrect, unfortunately. If you drop nsw & nuw, it's correct.

Ok, I think I get it now. Even though the shuffle mask's undef guarantees that any uses are also undef in those lanes, we're creating poison in the binops that didn't exist before. Therefore, other passes may use that to create unsound transforms. I don't think we have vector folds that are that sophisticated yet, but we can't open that door. It's similar to the problem we have with div/rem, but not as blatant. So the solution is to generalize the hack that I added for div/rem (getSafeVectorConstantForIntDivRem), so we can allow these transforms for any poison-producers. And the right way to replace undef with safe constants is almost what we have here: translate to the identity constant for a given binop when possible. Other transforms can turn that back into undef when they determine that it's safe.

[spatel]

I committed the enhancement to getBinOpIdentity() as an NFC preliminary because we can use that in D49047. So this patch will just be a one-liner to use the new call once we have all of the safety checks sorted out.

include/llvm/IR/Constants.h

public:
   static Constant *getExactAShr(Constant *C1, Constant *C2) {
     return getAShr(C1, C2, true);
   }
 
   static Constant *getExactLShr(Constant *C1, Constant *C2) {
     return getLShr(C1, C2, true);
   }
 
-  /// Return the identity for the given binary operation,
-  /// i.e. a constant C such that X op C = X and C op X = X for every X.  It
-  /// returns null if the operator doesn't have an identity.
-  static Constant *getBinOpIdentity(unsigned Opcode, Type *Ty);
+  /// Return the identity constant for a binary opcode.
+  /// The identity constant C is defined as X op C = X and C op X = X for every
+  /// X when the binary operation is commutative. If the binop is not
+  /// commutative, callers can acquire the operand 1 identity constant by
+  /// setting AllowRHSConstant to true. For example, any shift has a zero
+  /// identity constant for operand 1: X shift 0 = X.
+  /// Return nullptr if the operator does not have an identity constant.
+  static Constant *getBinOpIdentity(unsigned Opcode, Type *Ty,
+                                    bool AllowRHSConstant = false);
 
   /// Return the absorbing element for the given binary
   /// operation, i.e. a constant C such that X op C = C and C op X = C for
   /// every X.  For example, this returns zero for integer multiplication.
   /// It returns null if the operator doesn't have an absorbing element.
   static Constant *getBinOpAbsorber(unsigned Opcode, Type *Ty);
 
   /// Transparently provide more efficient getOperand methods.

lib/IR/Constants.cpp

return get(Instruction::LShr, C1, C2,
              isExact ? PossiblyExactOperator::IsExact : 0);
 }
 
 Constant *ConstantExpr::getAShr(Constant *C1, Constant *C2, bool isExact) {
   return get(Instruction::AShr, C1, C2,
              isExact ? PossiblyExactOperator::IsExact : 0);
 }
 
-// FIXME: Add a parameter to specify the operand number for non-commutative ops.
-// For example, the operand 1 identity constant for any shift is the null value
-// because shift-by-0 always returns operand 0.
-Constant *ConstantExpr::getBinOpIdentity(unsigned Opcode, Type *Ty) {
-  switch (Opcode) {
-  default:
-    // Doesn't have an identity.
-    return nullptr;
+Constant *ConstantExpr::getBinOpIdentity(unsigned Opcode, Type *Ty,
+                                         bool AllowRHSConstant) {
+  assert(Instruction::isBinaryOp(Opcode) && "Only binops allowed");
 
-  case Instruction::Add:
-  case Instruction::Or:
-  case Instruction::Xor:
+  // Commutative opcodes: it does not matter if AllowRHSConstant is set.
+  if (Instruction::isCommutative(Opcode)) {
+    switch (Opcode) {
+      case Instruction::Add: // X + 0 = X
+      case Instruction::Or:  // X | 0 = X
+      case Instruction::Xor: // X ^ 0 = X
         return Constant::getNullValue(Ty);
-
-  case Instruction::Mul:
+      case Instruction::Mul: // X * 1 = X
         return ConstantInt::get(Ty, 1);
-
-  case Instruction::And:
+      case Instruction::And: // X & -1 = X
         return Constant::getAllOnesValue(Ty);
-
+      case Instruction::FAdd: // X + -0.0 = X
         // TODO: If the fadd has 'nsz', should we return +0.0?
-  case Instruction::FAdd:
         return ConstantFP::getNegativeZero(Ty);
+      case Instruction::FMul: // X * 1.0 = X
+        return ConstantFP::get(Ty, 1.0);
+      default:
+        llvm_unreachable("Every commutative binop has an identity constant");
+    }
+  }
 
-  case Instruction::FMul:
+  // Non-commutative opcodes: AllowRHSConstant must be set.
+  if (!AllowRHSConstant)
+    return nullptr;
+
+  switch (Opcode) {
+    case Instruction::Sub:  // X - 0 = X
+    case Instruction::Shl:  // X << 0 = X
+    case Instruction::LShr: // X >>u 0 = X
+    case Instruction::AShr: // X >> 0 = X
+    case Instruction::FSub: // X - 0.0 = X
+      return Constant::getNullValue(Ty);
+    case Instruction::SDiv: // X / 1 = X
+    case Instruction::UDiv: // X /u 1 = X
+      return ConstantInt::get(Ty, 1);
+    case Instruction::FDiv: // X / 1.0 = X
       return ConstantFP::get(Ty, 1.0);
+    default:
+      return nullptr;
   }
 }
 
 Constant *ConstantExpr::getBinOpAbsorber(unsigned Opcode, Type *Ty) {
   switch (Opcode) {
   default:
     // Doesn't have an absorber.
     return nullptr;

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

static Instruction *foldSelectShuffleWith1Binop(ShuffleVectorInst &Shuf) {
   // TODO: Allow div/rem by accounting for potential UB due to undef elements.
   if (BO->isIntDivRem())
     return nullptr;
 
   // The identity constant for a binop leaves a variable operand unchanged. For
   // a vector, this is a splat of something like 0, -1, or 1.
   // If there's no identity constant for this binop, we're done.
   BinaryOperator::BinaryOps BOpcode = BO->getOpcode();
-  Constant *IdC = ConstantExpr::getBinOpIdentity(BOpcode, Shuf.getType());
+  Constant *IdC = ConstantExpr::getBinOpIdentity(BOpcode, Shuf.getType(), true);
   if (!IdC)
     return nullptr;
 
   // Shuffle identity constants into the lanes that return the original value.
   // Example: shuf (mul X, {-1,-2,-3,-4}), X, {0,5,6,3} --> mul X, {-1,1,1,-4}
   // Example: shuf X, (add X, {-1,-2,-3,-4}), {0,1,6,7} --> add X, {0,0,-3,-4}
   // The existing binop constant vector remains in the same operand position.
   Constant *Mask = Shuf.getMask();

test/Transforms/InstCombine/shuffle_select.ll

 ;
   %b = mul nsw nuw <4 x i32> %v, <i32 11, i32 12, i32 13, i32 14>
   %s = shufflevector <4 x i32> %v, <4 x i32> %b, <4 x i32> <i32 undef, i32 5, i32 2, i32 7>
   ret <4 x i32> %s
 }
 
 define <4 x i32> @shl(<4 x i32> %v) {
 ; CHECK-LABEL: @shl(
-; CHECK-NEXT:    [[B:%.*]] = shl nuw <4 x i32> [[V:%.*]], <i32 11, i32 12, i32 13, i32 14>
-; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x i32> [[B]], <4 x i32> [[V]], <4 x i32> <i32 undef, i32 5, i32 2, i32 undef>
+; CHECK-NEXT:    [[S:%.*]] = shl nuw <4 x i32> [[V:%.*]], <i32 undef, i32 0, i32 13, i32 undef>
 ; CHECK-NEXT:    ret <4 x i32> [[S]]
 ;
   %b = shl nuw <4 x i32> %v, <i32 11, i32 12, i32 13, i32 14>
   %s = shufflevector <4 x i32> %b, <4 x i32> %v, <4 x i32> <i32 undef, i32 5, i32 2, i32 undef>
   ret <4 x i32> %s
 }
 
 define <4 x i32> @lshr_constant_op0(<4 x i32> %v) {
 ; CHECK-LABEL: @lshr_constant_op0(
-; CHECK-NEXT:    [[B:%.*]] = lshr exact <4 x i32> [[V:%.*]], <i32 11, i32 12, i32 13, i32 14>
-; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x i32> [[V]], <4 x i32> [[B]], <4 x i32> <i32 4, i32 5, i32 2, i32 7>
+; CHECK-NEXT:    [[S:%.*]] = lshr exact <4 x i32> [[V:%.*]], <i32 11, i32 12, i32 0, i32 14>
 ; CHECK-NEXT:    ret <4 x i32> [[S]]
 ;
   %b = lshr exact <4 x i32> %v, <i32 11, i32 12, i32 13, i32 14>
   %s = shufflevector <4 x i32> %v, <4 x i32> %b, <4 x i32> <i32 4, i32 5, i32 2, i32 7>
   ret <4 x i32> %s
 }
 define <4 x i32> @lshr_constant_op1(<4 x i32> %v) {
 ; CHECK-LABEL: @lshr_constant_op1(
 ; CHECK-NEXT:    [[B:%.*]] = lshr exact <4 x i32> <i32 11, i32 12, i32 13, i32 14>, [[V:%.*]]
 ; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x i32> [[V]], <4 x i32> [[B]], <4 x i32> <i32 4, i32 5, i32 2, i32 7>
 ; CHECK-NEXT:    ret <4 x i32> [[S]]
 ;
   %b = lshr exact <4 x i32> <i32 11, i32 12, i32 13, i32 14>, %v
   %s = shufflevector <4 x i32> %v, <4 x i32> %b, <4 x i32> <i32 4, i32 5, i32 2, i32 7>
   ret <4 x i32> %s
 }
 
 ; Try weird types.
 
 define <3 x i32> @ashr(<3 x i32> %v) {
 ; CHECK-LABEL: @ashr(
-; CHECK-NEXT:    [[B:%.*]] = ashr <3 x i32> [[V:%.*]], <i32 11, i32 12, i32 13>
-; CHECK-NEXT:    [[S:%.*]] = shufflevector <3 x i32> [[B]], <3 x i32> [[V]], <3 x i32> <i32 3, i32 1, i32 2>
+; CHECK-NEXT:    [[S:%.*]] = ashr <3 x i32> [[V:%.*]], <i32 0, i32 12, i32 13>
 ; CHECK-NEXT:    ret <3 x i32> [[S]]
 ;
   %b = ashr <3 x i32> %v, <i32 11, i32 12, i32 13>
   %s = shufflevector <3 x i32> %b, <3 x i32> %v, <3 x i32> <i32 3, i32 1, i32 2>
   ret <3 x i32> %s
 }
 
 define <3 x i42> @and(<3 x i42> %v) {

 ;
   %b = fdiv fast <4 x double> <double 41.0, double 42.0, double 43.0, double 44.0>, %v
   %s = shufflevector <4 x double> %v, <4 x double> %b, <4 x i32> <i32 undef, i32 1, i32 6, i32 7>
   ret <4 x double> %s
 }
 
 define <4 x double> @fdiv_constant_op1(<4 x double> %v) {
 ; CHECK-LABEL: @fdiv_constant_op1(
-; CHECK-NEXT:    [[B:%.*]] = fdiv reassoc <4 x double> [[V:%.*]], <double 4.100000e+01, double 4.200000e+01, double 4.300000e+01, double 4.400000e+01>
-; CHECK-NEXT:    [[S:%.*]] = shufflevector <4 x double> [[V]], <4 x double> [[B]], <4 x i32> <i32 undef, i32 1, i32 6, i32 7>
+; CHECK-NEXT:    [[S:%.*]] = fdiv reassoc <4 x double> [[V:%.*]], <double undef, double 1.000000e+00, double 4.300000e+01, double 4.400000e+01>
 ; CHECK-NEXT:    ret <4 x double> [[S]]
 ;
   %b = fdiv reassoc <4 x double> %v, <double 41.0, double 42.0, double 43.0, double 44.0>
   %s = shufflevector <4 x double> %v, <4 x double> %b, <4 x i32> <i32 undef, i32 1, i32 6, i32 7>
   ret <4 x double> %s
 }
 
 define <4 x double> @frem(<4 x double> %v) {