match correct mask only

Author: macsencasaus

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -1350,37 +1350,38 @@ Value *InstCombinerImpl::SimplifySelectsFeedingBinaryOp(BinaryOperator &I,
   };
 
   // Special case for reconstructing across a select:
-  // (Cond ? V1 : (X & Mask)) op
+  // (Cond ? V1 : (X & Mask)) |
   // zext (Cond ? V2 : trunc X)
-  // -> (Cond ? (V1 op zext V2) : ((X & Mask) op zext trunc X))
+  // -> (Cond ? (V1 | zext V2) : X)
   auto foldReconstruction = [&](Value *V1, Value *Masked,
                                 Value *ZExtSel) -> Value * {
     if (Opcode != Instruction::Or)
       return nullptr;
 
     Value *X;
-    if (!match(Masked, m_OneUse(m_And(m_Value(X), m_Constant()))))
+    const APInt *C;
+    if (!match(Masked, m_OneUse(m_And(m_Value(X), m_APInt(C)))))
       return nullptr;
 
     Value *V2, *Trunc;
     if (!match(ZExtSel, m_ZExt(m_OneUse(m_Select(m_Specific(Cond), m_Value(V2),
                                                  m_Value(Trunc))))))
       return nullptr;
 
+    if (*C != APInt::getBitsSetFrom(X->getType()->getScalarSizeInBits(),
+                                    Trunc->getType()->getScalarSizeInBits())) {
+      return nullptr;
+    }
+
     if (!match(Trunc, m_Trunc(m_Specific(X))))
       return nullptr;
 
     Value *ZExtTrue = Builder.CreateZExt(V2, V1->getType());
-    Value *True;
-    if (!(True = simplifyBinOp(Opcode, V1, ZExtTrue, FMF, Q)))
-      True = Builder.CreateOr(V1, ZExtTrue);
-
-    Value *ZExtFalse = Builder.CreateZExt(Trunc, V1->getType());
-    Value *False;
-    if (!(False = simplifyBinOp(Opcode, Masked, ZExtFalse, FMF, Q)))
-      False = Builder.CreateOr(Masked, ZExtFalse);
+    Value *True = simplifyBinOp(Opcode, V1, ZExtTrue, FMF, Q);
+    if (!True)
+      return nullptr;
 
-    return Builder.CreateSelect(Cond, True, False, I.getName());
+    return Builder.CreateSelect(Cond, True, X, I.getName());
   };
 
   if (LHSIsSelect && RHSIsSelect && A == D) {

llvm/test/Transforms/InstCombine/select-reconstruction.ll

@@ -37,14 +37,18 @@ define i40 @select_reconstruction_any_cmp_val(i40 %arg0, i8 %arg1) {
   ret i40 %recomb
 }
 
+; negative test
 define i40 @select_reconstruction_257_mask(i40 %arg0) {
 ; CHECK-LABEL: define i40 @select_reconstruction_257_mask(
 ; CHECK-SAME: i40 [[ARG0:%.*]]) {
 ; CHECK-NEXT:    [[TMP1:%.*]] = trunc i40 [[ARG0]] to i8
 ; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i8 [[TMP1]], 2
 ; CHECK-NEXT:    [[TMP3:%.*]] = and i40 [[ARG0]], -257
+; CHECK-NEXT:    [[SELECT_LOW:%.*]] = select i1 [[TMP2]], i8 0, i8 [[TMP1]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = select i1 [[TMP2]], i40 0, i40 [[TMP3]]
-; CHECK-NEXT:    ret i40 [[TMP4]]
+; CHECK-NEXT:    [[ZEXT_LOW:%.*]] = zext i8 [[SELECT_LOW]] to i40
+; CHECK-NEXT:    [[RECOMB:%.*]] = or disjoint i40 [[TMP4]], [[ZEXT_LOW]]
+; CHECK-NEXT:    ret i40 [[RECOMB]]
 ;
   %low = trunc i40 %arg0 to i8
   %is_low_two = icmp eq i8 %low, 2