[InstCombine] reassociate logic ops with constants separated by a zext

[spatel]

Summary

This is a partial implementation of a general fold for associative+commutative operators:
(op (cast (op X, C2)), C1) --> (cast (op X, op (C1, C2)))
(op (cast (op X, C2)), C1) --> (op (cast X), op (C1, C2))

There are 7 associative operators and 13 cast types, so this could potentially go a lot further. But I'd like to make sure this much is implemented properly before extending.

Also, the zext + logic case should be enough to allow D22271 to proceed.

[majnemer]

LGTM

lib/Transforms/InstCombine/InstructionCombining.cpp

if (!FPMO) {
     return;
   }
 
   FastMathFlags FMF = I.getFastMathFlags();
   I.clearSubclassOptionalData();
   I.setFastMathFlags(FMF);
 }
 
+/// Combine constant operands of associative operations either before or after a
+/// cast to eliminate one of the associative operations:
+/// (op (cast (op X, C2)), C1) --> (cast (op X, op (C1, C2)))
+/// (op (cast (op X, C2)), C1) --> (op (cast X), op (C1, C2))
+static bool simplifyAssocCastAssoc(BinaryOperator *BinOp1) {
+  auto *Cast = dyn_cast<CastInst>(BinOp1->getOperand(0));
+  if (!Cast || !Cast->hasOneUse())
+    return false;
+
+  // TODO: Enhance logic for other casts and remove this check.
+  auto CastOpcode = Cast->getOpcode();
+  if (CastOpcode != Instruction::ZExt)
+    return false;
+
+  // TODO: Enhance logic for other BinOps and remove this check.
+  auto AssocOpcode = BinOp1->getOpcode();
+  if (AssocOpcode != Instruction::Xor && AssocOpcode != Instruction::And &&
+      AssocOpcode != Instruction::Or)
+    return false;
+
+  auto *BinOp2 = dyn_cast<BinaryOperator>(Cast->getOperand(0));
+  if (!BinOp2 || !BinOp2->hasOneUse() || BinOp2->getOpcode() != AssocOpcode)
+    return false;
+
+  Constant *C1, *C2;
+  if (!match(BinOp1->getOperand(1), m_Constant(C1)) ||
+      !match(BinOp2->getOperand(1), m_Constant(C2)))
+    return false;
+
+  // TODO: This assumes a zext cast.
+  // Eg, if it was a trunc, we'd cast C1 to the source type because casting C2
+  // to the destination type might lose bits.
+
+  // Fold the constants together in the destination type:
+  // (op (cast (op X, C2)), C1) --> (op (cast X), FoldedC)
+  Type *DestTy = C1->getType();
+  Constant *CastC2 = ConstantExpr::getCast(CastOpcode, C2, DestTy);
+  Constant *FoldedC = ConstantExpr::get(AssocOpcode, C1, CastC2);
+  Cast->setOperand(0, BinOp2->getOperand(0));
+  BinOp1->setOperand(1, FoldedC);
+  return true;
+}
+
 /// This performs a few simplifications for operators that are associative or
 /// commutative:
 ///
 ///  Commutative operators:
 ///
 ///  1. Order operands such that they are listed from right (least complex) to
 ///     left (most complex).  This puts constants before unary operators before
 ///     binary operators.
if (I.isAssociative()) {
           Changed = true;
           ++NumReassoc;
           continue;
         }
       }
     }
 
     if (I.isAssociative() && I.isCommutative()) {
+      if (simplifyAssocCastAssoc(&I)) {
+        Changed = true;
+        ++NumReassoc;
+        continue;
+      }
+
       // Transform: "(A op B) op C" ==> "(C op A) op B" if "C op A" simplifies.
       if (Op0 && Op0->getOpcode() == Opcode) {
         Value *A = Op0->getOperand(0);
         Value *B = Op0->getOperand(1);
         Value *C = I.getOperand(1);
 
         // Does "C op A" simplify?
         if (Value *V = SimplifyBinOp(Opcode, C, A, DL)) {

test/Transforms/InstCombine/assoc-cast-assoc.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 define i5 @XorZextXor(i3 %a) {
 ; CHECK-LABEL: @XorZextXor(
-; CHECK-NEXT:    [[OP1:%.*]] = xor i3 %a, 3
-; CHECK-NEXT:    [[CAST:%.*]] = zext i3 [[OP1]] to i5
-; CHECK-NEXT:    [[OP2:%.*]] = xor i5 [[CAST]], 12
+; CHECK-NEXT:    [[CAST:%.*]] = zext i3 %a to i5
+; CHECK-NEXT:    [[OP2:%.*]] = xor i5 [[CAST]], 15
 ; CHECK-NEXT:    ret i5 [[OP2]]
 ;
   %op1 = xor i3 %a, 3
   %cast = zext i3 %op1 to i5
   %op2 = xor i5 %cast, 12
   ret i5 %op2
 }
 
 define <2 x i32> @XorZextXorVec(<2 x i1> %a) {
 ; CHECK-LABEL: @XorZextXorVec(
-; CHECK-NEXT:    [[OP1:%.*]] = xor <2 x i1> %a, <i1 true, i1 false>
-; CHECK-NEXT:    [[CAST:%.*]] = zext <2 x i1> [[OP1]] to <2 x i32>
-; CHECK-NEXT:    [[OP2:%.*]] = xor <2 x i32> [[CAST]], <i32 3, i32 1>
+; CHECK-NEXT:    [[CAST:%.*]] = zext <2 x i1> %a to <2 x i32>
+; CHECK-NEXT:    [[OP2:%.*]] = xor <2 x i32> [[CAST]], <i32 2, i32 1>
 ; CHECK-NEXT:    ret <2 x i32> [[OP2]]
 ;
   %op1 = xor <2 x i1> %a, <i1 true, i1 false>
   %cast = zext <2 x i1> %op1 to <2 x i32>
   %op2 = xor <2 x i32> %cast, <i32 3, i32 1>
   ret <2 x i32> %op2
 }
 
 define i5 @OrZextOr(i3 %a) {
 ; CHECK-LABEL: @OrZextOr(
-; CHECK-NEXT:    [[OP1:%.*]] = or i3 %a, 3
-; CHECK-NEXT:    [[CAST:%.*]] = zext i3 [[OP1]] to i5
-; CHECK-NEXT:    [[OP2:%.*]] = or i5 [[CAST]], 8
+; CHECK-NEXT:    [[CAST:%.*]] = zext i3 %a to i5
+; CHECK-NEXT:    [[OP2:%.*]] = or i5 [[CAST]], 11
 ; CHECK-NEXT:    ret i5 [[OP2]]
 ;
   %op1 = or i3 %a, 3
   %cast = zext i3 %op1 to i5
   %op2 = or i5 %cast, 8
   ret i5 %op2
 }
 
 define <2 x i32> @OrZextOrVec(<2 x i2> %a) {
 ; CHECK-LABEL: @OrZextOrVec(
-; CHECK-NEXT:    [[OP1:%.*]] = or <2 x i2> %a, <i2 -2, i2 0>
-; CHECK-NEXT:    [[CAST:%.*]] = zext <2 x i2> [[OP1]] to <2 x i32>
-; CHECK-NEXT:    [[OP2:%.*]] = or <2 x i32> [[CAST]], <i32 1, i32 5>
+; CHECK-NEXT:    [[CAST:%.*]] = zext <2 x i2> %a to <2 x i32>
+; CHECK-NEXT:    [[OP2:%.*]] = or <2 x i32> [[CAST]], <i32 3, i32 5>
 ; CHECK-NEXT:    ret <2 x i32> [[OP2]]
 ;
   %op1 = or <2 x i2> %a, <i2 2, i2 0>
   %cast = zext <2 x i2> %op1 to <2 x i32>
   %op2 = or <2 x i32> %cast, <i32 1, i32 5>
   ret <2 x i32> %op2
 }
 
 ; Unlike the rest, this case is handled by SimplifyDemandedBits / ShrinkDemandedConstant.
 
 define i5 @AndZextAnd(i3 %a) {
 ; CHECK-LABEL: @AndZextAnd(
 ; CHECK-NEXT:    [[CAST:%.*]] = zext i3 %a to i5
 ; CHECK-NEXT:    [[OP2:%.*]] = and i5 [[CAST]], 2
 ; CHECK-NEXT:    ret i5 [[OP2]]
 ;
   %op1 = and i3 %a, 3
   %cast = zext i3 %op1 to i5
   %op2 = and i5 %cast, 14
   ret i5 %op2
 }
 
 define <2 x i32> @AndZextAndVec(<2 x i8> %a) {
 ; CHECK-LABEL: @AndZextAndVec(
-; CHECK-NEXT:    [[OP1:%.*]] = and <2 x i8> %a, <i8 7, i8 0>
-; CHECK-NEXT:    [[CAST:%.*]] = zext <2 x i8> [[OP1]] to <2 x i32>
-; CHECK-NEXT:    [[OP2:%.*]] = and <2 x i32> [[CAST]], <i32 261, i32 1>
+; CHECK-NEXT:    [[CAST:%.*]] = zext <2 x i8> %a to <2 x i32>
+; CHECK-NEXT:    [[OP2:%.*]] = and <2 x i32> [[CAST]], <i32 5, i32 0>
 ; CHECK-NEXT:    ret <2 x i32> [[OP2]]
 ;
   %op1 = and <2 x i8> %a, <i8 7, i8 0>
   %cast = zext <2 x i8> %op1 to <2 x i32>
   %op2 = and <2 x i32> %cast, <i32 261, i32 1>
   ret <2 x i32> %op2
 }