[InstCombine] reverse canonicalization of binops to allow more shuffle folding

[spatel]

Summary

This is a superset of D48485 that includes transforming an 'or' back into an 'add' in addition to the shl+mul case shown there.

Here, we need value tracking to determine that the 'or' can be reversed into an 'add'.

IIUC, there are only a handful of opcodes where we'll try this. To repeat from the other review - I think we're distinguishing this pattern from the scenario where we have eliminated a binop completely and need to materialize a 'ghost' op with an identity constant.

[lebedev.ri]

I'm guessing this superseded D48485, which should be closed?
Also, is D48678 supposed to go on top of this, or trunk?

[spatel]

In D48662#1146575, @lebedev.ri wrote:

I'm guessing this superseded D48485, which should be closed?
Also, is D48678 supposed to go on top of this, or trunk?

Each patch can be applied to trunk as-is. I posted this one to show one way that D48485 could be extended, so yes this one would obsolete D48485 if we want to make a larger change all at once. But my preference is to make the smaller change first to reduce risk. Ie, I think D48485 should be reviewed/committed, then we can make some NFC changes in preparation for this patch (and I need to add more tests). If we take that route, this patch will shrink to just adding one more 'case' to the switch.

D48678 is also against trunk, so it can be reviewed as a stand-alone patch. I wanted to demonstrate that the 2-variable case is a small add-on either way. Given that, maybe it's easier to start there because that part won't have any follow-ups AFAIK.

[lebedev.ri]

In D48662#1146670, @spatel wrote:

In D48662#1146575, @lebedev.ri wrote:

I'm guessing this superseded D48485, which should be closed?
Also, is D48678 supposed to go on top of this, or trunk?

Each patch can be applied to trunk as-is. I posted this one to show one way that D48485 could be extended, so yes this one would obsolete D48485 if we want to make a larger change all at once. But my preference is to make the smaller change first to reduce risk. Ie, I think D48485 should be reviewed/committed, then we can make some NFC changes in preparation for this patch (and I need to add more tests). If we take that route, this patch will shrink to just adding one more 'case' to the switch.

Then i would suggest to do so now:

1. Update D48485 with this patch minus the Instruction::Or handling.
2. Update this D48662 - on top of D48485 (so the diff is relative), add Instruction::Or handling.

Will simplify reviewing, and signal the review order :)

D48678 is also against trunk, so it can be reviewed as a stand-alone patch. I wanted to demonstrate that the 2-variable case is a small add-on either way. Given that, maybe it's easier to start there because that part won't have any follow-ups AFAIK.

Yeah, that one looks really simple. I'll try to review that one..

[lebedev.ri]

Please rebase once D48678 lands too.

[spatel]

Patch updated:

1. Rebased now that D48485 / D48678 are committed.
2. Added more tests with rL335984.

[RKSimon]

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1155 ↗

(On Diff #153502)

Better to do this as an Optional<BinopElts> ?

[lebedev.ri]

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1158 ↗

(On Diff #153502)

// shl X, C --> mul X, (1 << C)

1215–1216 ↗

(On Diff #153502)

1. Why are you checking isa<Constant>()? Right now that can be an assertion. Future-proofing?
2. How about just adding BinopElts::operator bool() const and doing

if(BinopElts AltB0 = getAlternateBinop(B0, DL))

Or Optional<> like @RKSimon suggests.
?

[spatel]

Patch updated - no functional changes from the last rev, but:

1. Add a "operator bool()" to make the code cleaner.
2. Assert that operand 1 of the alternate binop is a constant (might change as we add more cases).
3. Improve formula in code comment with parens.

[lebedev.ri]

Other than the flags question, i think this is ok.

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1181 ↗

(On Diff #153526)

return {}; doesn't work?
Elsewhere too, i don't think you need to specify BinopElts.

test/Transforms/InstCombine/shuffle_select.ll

600 ↗

(On Diff #153526)

Too bad alive does not work with vectors :/
I'm not sure i understand why we drop nuw but keep nsw.
If anything, i'd guess the or is implicit nuw,
so i would be less suprized if we dropped nsw instead.

[spatel]

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1181 ↗

(On Diff #153526)

That works. I wasn't sure if there's some LLVM pref/reason to do it this way (example: return SDValue() in codegen).

test/Transforms/InstCombine/shuffle_select.ll

600 ↗

(On Diff #153526)

This transform drops the flags since 'or' has no wrapping.
We add back the 'nsw' in visitAdd, but it doesn't find 'nuw' is possible (a shortcoming in computeOverflowForUnsignedAdd)...it should be both I think.

[spatel]

Patch updated:
Use less wordy {} struct initialization for return vals.

[RKSimon]

LGTM - thanks

lib/Transforms/InstCombine/InstCombineVectorOps.cpp

1181 ↗

(On Diff #153526)

Hmm, still think using Optional<> is cleaner, but this is OK.

test/Transforms/InstCombine/shuffle_select.ll

600 ↗

(On Diff #153526)

Please can you raise this as a poor codegen bug?

[spatel]

test/Transforms/InstCombine/shuffle_select.ll

600 ↗

(On Diff #153526)

https://bugs.llvm.org/show_bug.cgi?id=38021

llvm/trunk/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

static bool isShuffleExtractingFromLHS(ShuffleVectorInst &SVI,
   if (BegIdx > EndIdx || EndIdx >= LHSElems || EndIdx - BegIdx != MaskElems - 1)
     return false;
   for (unsigned I = 0; I != MaskElems; ++I)
     if (static_cast<unsigned>(Mask[I]) != BegIdx + I)
       return false;
   return true;
 }
 
+/// These are the ingredients in an alternate form binary operator as described
+/// below.
+struct BinopElts {
+  BinaryOperator::BinaryOps Opcode;
+  Value *Op0;
+  Value *Op1;
+  BinopElts(BinaryOperator::BinaryOps Opc = (BinaryOperator::BinaryOps)0,
+            Value *V0 = nullptr, Value *V1 = nullptr) :
+      Opcode(Opc), Op0(V0), Op1(V1) {}
+  operator bool() const { return Opcode != 0; }
+};
+
+/// Binops may be transformed into binops with different opcodes and operands.
+/// Reverse the usual canonicalization to enable folds with the non-canonical
+/// form of the binop. If a transform is possible, return the elements of the
+/// new binop. If not, return invalid elements.
+static BinopElts getAlternateBinop(BinaryOperator *BO, const DataLayout &DL) {
+  Value *BO0 = BO->getOperand(0), *BO1 = BO->getOperand(1);
+  Type *Ty = BO->getType();
+  switch (BO->getOpcode()) {
+    case Instruction::Shl: {
+      // shl X, C --> mul X, (1 << C)
+      Constant *C;
+      if (match(BO1, m_Constant(C))) {
+        Constant *ShlOne = ConstantExpr::getShl(ConstantInt::get(Ty, 1), C);
+        return { Instruction::Mul, BO0, ShlOne };
+      }
+      break;
+    }
+    case Instruction::Or: {
+      // or X, C --> add X, C (when X and C have no common bits set)
+      const APInt *C;
+      if (match(BO1, m_APInt(C)) && MaskedValueIsZero(BO0, *C, DL))
+        return { Instruction::Add, BO0, BO1 };
+      break;
+    }
+    default:
+      break;
+  }
+  return {};
+}
+
+/// Try to fold shuffles that are the equivalent of a vector select.
 static Instruction *foldSelectShuffle(ShuffleVectorInst &Shuf,
-                                      InstCombiner::BuilderTy &Builder) {
-  // Folds under here require the equivalent of a vector select.
+                                      InstCombiner::BuilderTy &Builder,
+                                      const DataLayout &DL) {
   if (!Shuf.isSelect())
     return nullptr;
 
   BinaryOperator *B0, *B1;
   if (!match(Shuf.getOperand(0), m_BinOp(B0)) ||
       !match(Shuf.getOperand(1), m_BinOp(B1)))
     return nullptr;
 
static Instruction *foldSelectShuffle(ShuffleVectorInst &Shuf,
   else
     return nullptr;
 
   // We need matching binops to fold the lanes together.
   BinaryOperator::BinaryOps Opc0 = B0->getOpcode();
   BinaryOperator::BinaryOps Opc1 = B1->getOpcode();
   bool DropNSW = false;
   if (ConstantsAreOp1 && Opc0 != Opc1) {
-    // If we have multiply and shift-left-by-constant, convert the shift:
-    // shl X, C --> mul X, 1 << C
     // TODO: We drop "nsw" if shift is converted into multiply because it may
     // not be correct when the shift amount is BitWidth - 1. We could examine
     // each vector element to determine if it is safe to keep that flag.
-    if (Opc0 == Instruction::Mul && Opc1 == Instruction::Shl) {
-      C1 = ConstantExpr::getShl(ConstantInt::get(C1->getType(), 1), C1);
-      Opc1 = Instruction::Mul;
-      DropNSW = true;
-    } else if (Opc0 == Instruction::Shl && Opc1 == Instruction::Mul) {
-      C0 = ConstantExpr::getShl(ConstantInt::get(C0->getType(), 1), C0);
-      Opc0 = Instruction::Mul;
+    if (Opc0 == Instruction::Shl || Opc1 == Instruction::Shl)
       DropNSW = true;
+    if (BinopElts AltB0 = getAlternateBinop(B0, DL)) {
+      assert(isa<Constant>(AltB0.Op1) && "Expecting constant with alt binop");
+      Opc0 = AltB0.Opcode;
+      C0 = cast<Constant>(AltB0.Op1);
+    } else if (BinopElts AltB1 = getAlternateBinop(B1, DL)) {
+      assert(isa<Constant>(AltB1.Op1) && "Expecting constant with alt binop");
+      Opc1 = AltB1.Opcode;
+      C1 = cast<Constant>(AltB1.Op1);
     }
   }
 
   if (Opc0 != Opc1)
     return nullptr;
 
   // The opcodes must be the same. Use a new name to make that clear.
   BinaryOperator::BinaryOps BOpc = Opc0;
Instruction *InstCombiner::visitShuffleVectorInst(ShuffleVectorInst &SVI) {
   Value *RHS = SVI.getOperand(1);
   SmallVector<int, 16> Mask = SVI.getShuffleMask();
   Type *Int32Ty = Type::getInt32Ty(SVI.getContext());
 
   if (auto *V = SimplifyShuffleVectorInst(
           LHS, RHS, SVI.getMask(), SVI.getType(), SQ.getWithInstruction(&SVI)))
     return replaceInstUsesWith(SVI, V);
 
-  if (Instruction *I = foldSelectShuffle(SVI, Builder))
+  if (Instruction *I = foldSelectShuffle(SVI, Builder, DL))
     return I;
 
   bool MadeChange = false;
   unsigned VWidth = SVI.getType()->getVectorNumElements();
 
   APInt UndefElts(VWidth, 0);
   APInt AllOnesEltMask(APInt::getAllOnesValue(VWidth));
   if (Value *V = SimplifyDemandedVectorElts(&SVI, AllOnesEltMask, UndefElts)) {

llvm/trunk/test/Transforms/InstCombine/shuffle_select.ll

 
 ; Or with constant can be converted to add to enable the fold.
 ; The 'shl' is here to allow analysis to determine that the 'or' can be transformed to 'add'.
 ; TODO: The 'or' constant is limited to a splat.
 
 define <4 x i32> @add_or(<4 x i32> %v) {
 ; CHECK-LABEL: @add_or(
 ; CHECK-NEXT:    [[V0:%.*]] = shl <4 x i32> [[V:%.*]], <i32 5, i32 5, i32 5, i32 5>
-; CHECK-NEXT:    [[T1:%.*]] = add <4 x i32> [[V0]], <i32 undef, i32 undef, i32 65536, i32 65537>
-; CHECK-NEXT:    [[T2:%.*]] = or <4 x i32> [[V0]], <i32 31, i32 31, i32 undef, i32 undef>
-; CHECK-NEXT:    [[T3:%.*]] = shufflevector <4 x i32> [[T1]], <4 x i32> [[T2]], <4 x i32> <i32 4, i32 5, i32 2, i32 3>
+; CHECK-NEXT:    [[T3:%.*]] = add <4 x i32> [[V0]], <i32 31, i32 31, i32 65536, i32 65537>
 ; CHECK-NEXT:    ret <4 x i32> [[T3]]
 ;
   %v0 = shl <4 x i32> %v, <i32 5, i32 5, i32 5, i32 5>                   ; clear the bottom bits
   %t1 = add <4 x i32> %v0, <i32 65534, i32 65535, i32 65536, i32 65537>  ; this can't be converted to 'or'
   %t2 = or <4 x i32> %v0, <i32 31, i32 31, i32 31, i32 31>               ; set the bottom bits
   %t3 = shufflevector <4 x i32> %t1, <4 x i32> %t2, <4 x i32> <i32 4, i32 5, i32 2, i32 3>
   ret <4 x i32> %t3
 }
 
 ; Try with 'or' as operand 0 of the shuffle.
 
 define <4 x i8> @or_add(<4 x i8> %v) {
 ; CHECK-LABEL: @or_add(
 ; CHECK-NEXT:    [[V0:%.*]] = lshr <4 x i8> [[V:%.*]], <i8 3, i8 3, i8 3, i8 3>
-; CHECK-NEXT:    [[T1:%.*]] = or <4 x i8> [[V0]], <i8 undef, i8 undef, i8 -64, i8 -64>
-; CHECK-NEXT:    [[T2:%.*]] = add nuw nsw <4 x i8> [[V0]], <i8 1, i8 2, i8 undef, i8 undef>
-; CHECK-NEXT:    [[T3:%.*]] = shufflevector <4 x i8> [[T1]], <4 x i8> [[T2]], <4 x i32> <i32 4, i32 5, i32 2, i32 3>
+; CHECK-NEXT:    [[T3:%.*]] = add nsw <4 x i8> [[V0]], <i8 1, i8 2, i8 -64, i8 -64>
 ; CHECK-NEXT:    ret <4 x i8> [[T3]]
 ;
   %v0 = lshr <4 x i8> %v, <i8 3, i8 3, i8 3, i8 3>          ; clear the top bits
   %t1 = or <4 x i8> %v0, <i8 192, i8 192, i8 192, i8 192>   ; set some top bits
   %t2 = add nsw nuw <4 x i8> %v0, <i8 1, i8 2, i8 3, i8 4>  ; this can't be converted to 'or'
   %t3 = shufflevector <4 x i8> %t1, <4 x i8> %t2, <4 x i32> <i32 4, i32 5, i32 2, i32 3>
   ret <4 x i8> %t3
 }
 
+; Negative test: not all 'or' insts can be converted to 'add'.
+
 define <4 x i8> @or_add_not_enough_masking(<4 x i8> %v) {
 ; CHECK-LABEL: @or_add_not_enough_masking(
 ; CHECK-NEXT:    [[V0:%.*]] = lshr <4 x i8> [[V:%.*]], <i8 1, i8 1, i8 1, i8 1>
 ; CHECK-NEXT:    [[T1:%.*]] = or <4 x i8> [[V0]], <i8 undef, i8 undef, i8 -64, i8 -64>
 ; CHECK-NEXT:    [[T2:%.*]] = add nuw nsw <4 x i8> [[V0]], <i8 1, i8 2, i8 undef, i8 undef>
 ; CHECK-NEXT:    [[T3:%.*]] = shufflevector <4 x i8> [[T1]], <4 x i8> [[T2]], <4 x i32> <i32 4, i32 5, i32 2, i32 3>
 ; CHECK-NEXT:    ret <4 x i8> [[T3]]
 ;
   %v0 = lshr <4 x i8> %v, <i8 1, i8 1, i8 1, i8 1>          ; clear not enough top bits
   %t1 = or <4 x i8> %v0, <i8 192, i8 192, i8 192, i8 192>   ; set some top bits
   %t2 = add nsw nuw <4 x i8> %v0, <i8 1, i8 2, i8 3, i8 4>  ; this can't be converted to 'or'
   %t3 = shufflevector <4 x i8> %t1, <4 x i8> %t2, <4 x i32> <i32 4, i32 5, i32 2, i32 3>
   ret <4 x i8> %t3
 }
 
 ; Try with 2 variable inputs.
 
 define <4 x i32> @add_or_2_vars(<4 x i32> %v, <4 x i32> %v1) {
 ; CHECK-LABEL: @add_or_2_vars(
 ; CHECK-NEXT:    [[V0:%.*]] = shl <4 x i32> [[V:%.*]], <i32 5, i32 5, i32 5, i32 5>
-; CHECK-NEXT:    [[T1:%.*]] = add <4 x i32> [[V1:%.*]], <i32 undef, i32 undef, i32 65536, i32 65537>
-; CHECK-NEXT:    [[T2:%.*]] = or <4 x i32> [[V0]], <i32 31, i32 31, i32 undef, i32 undef>
-; CHECK-NEXT:    [[T3:%.*]] = shufflevector <4 x i32> [[T1]], <4 x i32> [[T2]], <4 x i32> <i32 4, i32 5, i32 2, i32 3>
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i32> [[V1:%.*]], <4 x i32> [[V0]], <4 x i32> <i32 4, i32 5, i32 2, i32 3>
+; CHECK-NEXT:    [[T3:%.*]] = add <4 x i32> [[TMP1]], <i32 31, i32 31, i32 65536, i32 65537>
 ; CHECK-NEXT:    ret <4 x i32> [[T3]]
 ;
   %v0 = shl <4 x i32> %v, <i32 5, i32 5, i32 5, i32 5>                   ; clear the bottom bits
   %t1 = add <4 x i32> %v1, <i32 65534, i32 65535, i32 65536, i32 65537>  ; this can't be converted to 'or'
   %t2 = or <4 x i32> %v0, <i32 31, i32 31, i32 31, i32 31>               ; set the bottom bits
   %t3 = shufflevector <4 x i32> %t1, <4 x i32> %t2, <4 x i32> <i32 4, i32 5, i32 2, i32 3>
   ret <4 x i32> %t3
 }
 
 define <4 x i8> @or_add_2_vars(<4 x i8> %v, <4 x i8> %v1) {
 ; CHECK-LABEL: @or_add_2_vars(
 ; CHECK-NEXT:    [[V0:%.*]] = lshr <4 x i8> [[V:%.*]], <i8 3, i8 3, i8 3, i8 3>
-; CHECK-NEXT:    [[T1:%.*]] = or <4 x i8> [[V0]], <i8 undef, i8 undef, i8 -64, i8 -64>
-; CHECK-NEXT:    [[T2:%.*]] = add nuw nsw <4 x i8> [[V1:%.*]], <i8 1, i8 2, i8 undef, i8 undef>
-; CHECK-NEXT:    [[T3:%.*]] = shufflevector <4 x i8> [[T1]], <4 x i8> [[T2]], <4 x i32> <i32 4, i32 5, i32 2, i32 3>
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i8> [[V0]], <4 x i8> [[V1:%.*]], <4 x i32> <i32 4, i32 5, i32 2, i32 3>
+; CHECK-NEXT:    [[T3:%.*]] = add <4 x i8> [[TMP1]], <i8 1, i8 2, i8 -64, i8 -64>
 ; CHECK-NEXT:    ret <4 x i8> [[T3]]
 ;
   %v0 = lshr <4 x i8> %v, <i8 3, i8 3, i8 3, i8 3>          ; clear the top bits
   %t1 = or <4 x i8> %v0, <i8 192, i8 192, i8 192, i8 192>   ; set some top bits
   %t2 = add nsw nuw <4 x i8> %v1, <i8 1, i8 2, i8 3, i8 4>  ; this can't be converted to 'or'
   %t3 = shufflevector <4 x i8> %t1, <4 x i8> %t2, <4 x i32> <i32 4, i32 5, i32 2, i32 3>
   ret <4 x i8> %t3
 }