[InstCombine] Propagate disjoint flags during shl-binop transform (llvm#91333)

Alive2 Proof:
https://alive2.llvm.org/ce/z/D2jHrn

Author: AZero13

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

@@ -1174,7 +1174,11 @@ Instruction *InstCombinerImpl::visitShl(BinaryOperator &I) {
         // X & (CC << C)
         Value *M = Builder.CreateAnd(X, ConstantInt::get(Ty, CC->shl(*C)),
                                      X->getName() + ".mask");
-        return BinaryOperator::Create(Op0BO->getOpcode(), M, YS);
+        auto *NewOp = BinaryOperator::Create(Op0BO->getOpcode(), M, YS);
+        if (auto *Disjoint = dyn_cast<PossiblyDisjointInst>(Op0BO);
+            Disjoint && Disjoint->isDisjoint())
+          cast<PossiblyDisjointInst>(NewOp)->setIsDisjoint(true);
+        return NewOp;
       }
     }
 

llvm/test/Transforms/InstCombine/shl-bo.ll

@@ -294,6 +294,39 @@ define i8 @lshr_and_or(i8 %a, i8 %y)  {
   ret i8 %l
 }
 
+define i8 @lshr_and_or_disjoint(i8 %a, i8 %y)  {
+; CHECK-LABEL: @lshr_and_or_disjoint(
+; CHECK-NEXT:    [[X:%.*]] = srem i8 [[A:%.*]], 42
+; CHECK-NEXT:    [[B1:%.*]] = shl i8 [[X]], 2
+; CHECK-NEXT:    [[Y_MASK:%.*]] = and i8 [[Y:%.*]], 52
+; CHECK-NEXT:    [[L:%.*]] = or disjoint i8 [[Y_MASK]], [[B1]]
+; CHECK-NEXT:    ret i8 [[L]]
+;
+  %x = srem i8 %a, 42 ; thwart complexity-based canonicalization
+  %r = lshr i8 %y, 2
+  %m = and i8 %r, 13
+  %b = or disjoint i8 %x, %m
+  %l = shl i8 %b, 2
+  ret i8 %l
+}
+
+define i8 @ashr_and_or_disjoint(i8 %a, i8 %y)  {
+; CHECK-LABEL: @ashr_and_or_disjoint(
+; CHECK-NEXT:    [[X:%.*]] = srem i8 [[A:%.*]], 42
+; CHECK-NEXT:    [[B1:%.*]] = shl i8 [[X]], 2
+; CHECK-NEXT:    [[Y_MASK:%.*]] = and i8 [[Y:%.*]], 52
+; CHECK-NEXT:    [[L:%.*]] = or disjoint i8 [[Y_MASK]], [[B1]]
+; CHECK-NEXT:    ret i8 [[L]]
+;
+  %x = srem i8 %a, 42 ; thwart complexity-based canonicalization
+  %r = ashr i8 %y, 2
+  %m = and i8 %r, 13
+  %b = or disjoint i8 %x, %m
+  %l = shl i8 %b, 2
+  ret i8 %l
+}
+
+
 define <2 x i8> @lshr_and_or_commute_splat(<2 x i8> %a, <2 x i8> %y)  {
 ; CHECK-LABEL: @lshr_and_or_commute_splat(
 ; CHECK-NEXT:    [[X:%.*]] = srem <2 x i8> [[A:%.*]], <i8 42, i8 42>
@@ -614,8 +647,8 @@ define <8 x i16> @test_FoldShiftByConstant_CreateSHL2(<8 x i16> %in) {
 
 define <16 x i8> @test_FoldShiftByConstant_CreateAnd(<16 x i8> %in0) {
 ; CHECK-LABEL: @test_FoldShiftByConstant_CreateAnd(
-; CHECK-NEXT:    [[TMP1:%.*]] = mul <16 x i8> [[IN0:%.*]], <i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33>
-; CHECK-NEXT:    [[VSHL_N:%.*]] = and <16 x i8> [[TMP1]], <i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32>
+; CHECK-NEXT:    [[VSRA_N2:%.*]] = mul <16 x i8> [[IN0:%.*]], <i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33, i8 33>
+; CHECK-NEXT:    [[VSHL_N:%.*]] = and <16 x i8> [[VSRA_N2]], <i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32, i8 -32>
 ; CHECK-NEXT:    ret <16 x i8> [[VSHL_N]]
 ;
   %vsra_n = ashr <16 x i8> %in0, <i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5>