[InstCombine] Fix buggy `(mul X, Y)` -> `(shl X, Log2(Y))` transform PR62175

[goldstein.w.n]

Summary

Bug was because we recognized patterns like (shl 4, Z) as a power of
2 we could take Log2 of (2 + Z), but doing (shl X, (2 + Z)) can
cause a poison shift.

https://alive2.llvm.org/ce/z/yuJm_k

The fix is to verify that Log2(Y) will be a non-poisonous shift
amount. We can do this with:

`nsw` flag:
    - https://alive2.llvm.org/ce/z/yyyJBr
    - https://alive2.llvm.org/ce/z/YgubD_
`nuw` flag:
    - https://alive2.llvm.org/ce/z/-4mpyV
    - https://alive2.llvm.org/ce/z/a6ik6r
Prove `Y != 0`:
    - https://alive2.llvm.org/ce/z/ced4su
    - https://alive2.llvm.org/ce/z/X-JJHb

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

189

Replace RequireNoOverflow with an AssumeNonZero flag. This is what justifies it for division.

1227

Don't we always infer nuw for shl of one?

[goldstein.w.n]

RequireNoOverflow -> AssumeNonZero + just use nuw for 1 << shl

[goldstein.w.n]

done

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1220–1226

Would be a bit cleaner imho.

1252

This isn't right: The result being non-zero doesn't imply the arguments being non-zero here, see https://alive2.llvm.org/ce/z/AtZmi_. For the sake of simplicity, just always pass AssumeNonZero=false here.

[goldstein.w.n]

Fix bug in umin/umax case

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1252

This isn't right: The result being non-zero doesn't imply the arguments being non-zero here, see https://alive2.llvm.org/ce/z/AtZmi_. For the sake of simplicity, just always pass AssumeNonZero=false here.

Good catch. Thats been buggy since llvm15: https://godbolt.org/z/ehen9hz6G
We probably need to backport?

Should I do the following so backporting is easier?

1. Revert D146347
2. Add AssumeNonZero with false in this umin/umax case
3. Repost D146347 with AssumeNonZero changes?

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1252

This isn't right: The result being non-zero doesn't imply the arguments being non-zero here, see https://alive2.llvm.org/ce/z/AtZmi_. For the sake of simplicity, just always pass AssumeNonZero=false here.

Good catch. Thats been buggy since llvm15: https://godbolt.org/z/ehen9hz6G
We probably need to backport?

Should I do the following so backporting is easier?

1. Revert D146347
2. Add AssumeNonZero with false in this umin/umax case
3. Repost D146347 with AssumeNonZero changes?

The rationale being then the patch for Step 2 can just be backported w.o any conflicts.

[nikic]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1252

No strong opinion either way. The conflict resolution here looks pretty trivial.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1252

No strong opinion either way. The conflict resolution here looks pretty trivial.

Okay, I'll push as is.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1252

This isn't right: The result being non-zero doesn't imply the arguments being non-zero here, see https://alive2.llvm.org/ce/z/AtZmi_. For the sake of simplicity, just always pass AssumeNonZero=false here.

Created issue: https://github.com/llvm/llvm-project/issues/62221

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

if (match(Y, m_OneUse(m_Not(m_OneUse(m_Shl(m_AllOnes(), m_Value(Z))))))) {
     Value *Shl = Builder.CreateShl(FrX, Z, "mulshl");
     return Builder.CreateSub(Shl, FrX, Mul.getName());
   }
 
   return nullptr;
 }
 
 static Value *takeLog2(IRBuilderBase &Builder, Value *Op, unsigned Depth,
-                       bool DoFold);
+                       bool AssumeNonZero, bool DoFold);
 
 Instruction *InstCombinerImpl::visitMul(BinaryOperator &I) {
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
   if (Value *V =
           simplifyMulInst(Op0, Op1, I.hasNoSignedWrap(), I.hasNoUnsignedWrap(),
                           SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);
 
if (match(&I, m_CombineOr(m_c_Mul(m_SMax(m_Value(X), m_Value(Y)),
                                     m_c_UMin(m_Deferred(X), m_Deferred(Y))))))
     return BinaryOperator::CreateWithCopiedFlags(Instruction::Mul, X, Y, &I);
 
   // (mul Op0 Op1):
   //    if Log2(Op0) folds away ->
   //        (shl Op1, Log2(Op0))
   //    if Log2(Op1) folds away ->
   //        (shl Op0, Log2(Op1))
-  if (takeLog2(Builder, Op0, /*Depth*/ 0, /*DoFold*/ false)) {
-    Value *Res = takeLog2(Builder, Op0, /*Depth*/ 0, /*DoFold*/ true);
+  if (takeLog2(Builder, Op0, /*Depth*/ 0, /*AssumeNonZero*/ false,
+               /*DoFold*/ false)) {
+    Value *Res = takeLog2(Builder, Op0, /*Depth*/ 0, /*AssumeNonZero*/ false,
+                          /*DoFold*/ true);
     BinaryOperator *Shl = BinaryOperator::CreateShl(Op1, Res);
     // We can only propegate nuw flag.
     Shl->setHasNoUnsignedWrap(HasNUW);
     return Shl;
   }
-  if (takeLog2(Builder, Op1, /*Depth*/ 0, /*DoFold*/ false)) {
-    Value *Res = takeLog2(Builder, Op1, /*Depth*/ 0, /*DoFold*/ true);
+  if (takeLog2(Builder, Op1, /*Depth*/ 0, /*AssumeNonZero*/ false,
+               /*DoFold*/ false)) {
+    Value *Res = takeLog2(Builder, Op1, /*Depth*/ 0, /*AssumeNonZero*/ false,
+                          /*DoFold*/ true);
     BinaryOperator *Shl = BinaryOperator::CreateShl(Op0, Res);
     // We can only propegate nuw flag.
     Shl->setHasNoUnsignedWrap(HasNUW);
     return Shl;
   }
 
   bool Changed = false;
   if (!HasNSW && willNotOverflowSignedMul(Op0, Op1, I)) {

 }
 
 static const unsigned MaxDepth = 6;
 
 // Take the exact integer log2 of the value. If DoFold is true, create the
 // actual instructions, otherwise return a non-null dummy value. Return nullptr
 // on failure.
 static Value *takeLog2(IRBuilderBase &Builder, Value *Op, unsigned Depth,
-                       bool DoFold) {
+                       bool AssumeNonZero, bool DoFold) {
   auto IfFold = [DoFold](function_ref<Value *()> Fn) {
     if (!DoFold)
       return reinterpret_cast<Value *>(-1);
     return Fn();
   };
 
   // FIXME: assert that Op1 isn't/doesn't contain undef.
 
static Value *takeLog2(IRBuilderBase &Builder, Value *Op, unsigned Depth,
   // The remaining tests are all recursive, so bail out if we hit the limit.
   if (Depth++ == MaxDepth)
     return nullptr;
 
   // log2(zext X) -> zext log2(X)
   // FIXME: Require one use?
   Value *X, *Y;
   if (match(Op, m_ZExt(m_Value(X))))
-    if (Value *LogX = takeLog2(Builder, X, Depth, DoFold))
+    if (Value *LogX = takeLog2(Builder, X, Depth, AssumeNonZero, DoFold))
       return IfFold([&]() { return Builder.CreateZExt(LogX, Op->getType()); });
 
   // log2(X << Y) -> log2(X) + Y
   // FIXME: Require one use unless X is 1?
-  if (match(Op, m_Shl(m_Value(X), m_Value(Y))))
-    if (Value *LogX = takeLog2(Builder, X, Depth, DoFold))
+  if (match(Op, m_Shl(m_Value(X), m_Value(Y)))) {
+    bool Usable = AssumeNonZero;
+    if (!AssumeNonZero) {
+      BinaryOperator *BO = cast<BinaryOperator>(Op);
+      // nuw will be set if the `shl` is trivially non-zero.
+      Usable = BO->hasNoUnsignedWrap() || BO->hasNoSignedWrap();
+    }
+    if (Usable)
+      if (Value *LogX = takeLog2(Builder, X, Depth, AssumeNonZero, DoFold))
         return IfFold([&]() { return Builder.CreateAdd(LogX, Y); });
+  }
 
   // log2(Cond ? X : Y) -> Cond ? log2(X) : log2(Y)
   // FIXME: missed optimization: if one of the hands of select is/contains
   //        undef, just directly pick the other one.
   // FIXME: can both hands contain undef?
   // FIXME: Require one use?
   if (SelectInst *SI = dyn_cast<SelectInst>(Op))
-    if (Value *LogX = takeLog2(Builder, SI->getOperand(1), Depth, DoFold))
-      if (Value *LogY = takeLog2(Builder, SI->getOperand(2), Depth, DoFold))
+    if (Value *LogX = takeLog2(Builder, SI->getOperand(1), Depth,
+                               AssumeNonZero, DoFold))
+      if (Value *LogY = takeLog2(Builder, SI->getOperand(2), Depth,
+                                 AssumeNonZero, DoFold))
         return IfFold([&]() {
           return Builder.CreateSelect(SI->getOperand(0), LogX, LogY);
         });
 
   // log2(umin(X, Y)) -> umin(log2(X), log2(Y))
   // log2(umax(X, Y)) -> umax(log2(X), log2(Y))
   auto *MinMax = dyn_cast<MinMaxIntrinsic>(Op);
   if (MinMax && MinMax->hasOneUse() && !MinMax->isSigned())
-    if (Value *LogX = takeLog2(Builder, MinMax->getLHS(), Depth, DoFold))
-      if (Value *LogY = takeLog2(Builder, MinMax->getRHS(), Depth, DoFold))
+    if (Value *LogX = takeLog2(Builder, MinMax->getLHS(), Depth,
+                               AssumeNonZero, DoFold))
+      if (Value *LogY = takeLog2(Builder, MinMax->getRHS(), Depth,
+                                 AssumeNonZero, DoFold))
         return IfFold([&]() {
-          return Builder.CreateBinaryIntrinsic(
-              MinMax->getIntrinsicID(), LogX, LogY);
+          return Builder.CreateBinaryIntrinsic(MinMax->getIntrinsicID(), LogX,
+                                               LogY);
         });
 
   return nullptr;
 }
 
 /// If we have zero-extended operands of an unsigned div or rem, we may be able
 /// to narrow the operation (sink the zext below the math).
 static Instruction *narrowUDivURem(BinaryOperator &I,
if (match(Op0, m_LShr(m_NUWMul(m_Specific(Op1), m_Value(A)), m_Value(B))) ||
       match(Op0, m_LShr(m_NUWMul(m_Value(A), m_Specific(Op1)), m_Value(B)))) {
     Instruction *Lshr = BinaryOperator::CreateLShr(A, B);
     if (I.isExact() && cast<PossiblyExactOperator>(Op0)->isExact())
       Lshr->setIsExact();
     return Lshr;
   }
 
   // Op1 udiv Op2 -> Op1 lshr log2(Op2), if log2() folds away.
-  if (takeLog2(Builder, Op1, /*Depth*/0, /*DoFold*/false)) {
-    Value *Res = takeLog2(Builder, Op1, /*Depth*/0, /*DoFold*/true);
+  if (takeLog2(Builder, Op1, /*Depth*/ 0, /*AssumeNonZero*/ true,
+               /*DoFold*/ false)) {
+    Value *Res = takeLog2(Builder, Op1, /*Depth*/ 0,
+                          /*AssumeNonZero*/ true, /*DoFold*/ true);
     return replaceInstUsesWith(
         I, Builder.CreateLShr(Op0, Res, I.getName(), I.isExact()));
   }
 
   return nullptr;
 }
 
 Instruction *InstCombinerImpl::visitSDiv(BinaryOperator &I) {

llvm/test/Transforms/InstCombine/mul-pow2.ll

 ; CHECK-NEXT:    [[R:%.*]] = shl <2 x i8> [[X]], [[TMP1]]
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %x = mul <2 x i8> %xx, %xx
   %s = select i1 %c, <2 x i8> <i8 8, i8 16>, <2 x i8> <i8 4, i8 1>
   %r = mul <2 x i8> %x, %s
   ret <2 x i8> %r
 }
+
+
+define i8 @shl_add_log_may_cause_poison_pr62175_fail(i8 %x, i8 %y) {
+; CHECK-LABEL: @shl_add_log_may_cause_poison_pr62175_fail(
+; CHECK-NEXT:    [[SHL:%.*]] = shl i8 4, [[X:%.*]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul i8 [[SHL]], [[Y:%.*]]
+; CHECK-NEXT:    ret i8 [[MUL]]
+;
+  %shl = shl i8 4, %x
+  %mul = mul i8 %y, %shl
+  ret i8 %mul
+}
+
+define i8 @shl_add_log_may_cause_poison_pr62175_with_nuw(i8 %x, i8 %y) {
+; CHECK-LABEL: @shl_add_log_may_cause_poison_pr62175_with_nuw(
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[X:%.*]], 2
+; CHECK-NEXT:    [[MUL:%.*]] = shl i8 [[Y:%.*]], [[TMP1]]
+; CHECK-NEXT:    ret i8 [[MUL]]
+;
+  %shl = shl nuw i8 4, %x
+  %mul = mul i8 %y, %shl
+  ret i8 %mul
+}
+
+define i8 @shl_add_log_may_cause_poison_pr62175_with_nsw(i8 %x, i8 %y) {
+; CHECK-LABEL: @shl_add_log_may_cause_poison_pr62175_with_nsw(
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[X:%.*]], 2
+; CHECK-NEXT:    [[MUL:%.*]] = shl i8 [[Y:%.*]], [[TMP1]]
+; CHECK-NEXT:    ret i8 [[MUL]]
+;
+  %shl = shl nsw i8 4, %x
+  %mul = mul i8 %y, %shl
+  ret i8 %mul
+}