[llvm] [InstCombine] Restrict `foldBitCeil` to power-of-two integer widths (PR #173849)

[Yunbo Ni via llvm-commits]

https://github.com/cardigan1008 created https://github.com/llvm/llvm-project/pull/173849

The masking rewrite in `foldBitCeil` assumes a power-of-two bitwidth.

For non-power-of-two integer types, `(-ctlz) & (BitWidth - 1)` is not equivalent to `BitWidth - ctlz` and can miscompile.

This patch restricts the transform to power-of-two bitwidths. 

Alive2 proof: https://alive2.llvm.org/ce/z/i2E6zT

Fixes #173787 


>From 71ea73f7b3820cc416eff1668675a85991648b18 Mon Sep 17 00:00:00 2001
From: cardigan1008 <ybni at cse.cuhk.edu.hk>
Date: Mon, 29 Dec 2025 16:41:17 +0800
Subject: [PATCH 1/2] [InstCombine][NFC] Add pre-commit test case.

---
 llvm/test/Transforms/InstCombine/bit_ceil.ll | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/llvm/test/Transforms/InstCombine/bit_ceil.ll b/llvm/test/Transforms/InstCombine/bit_ceil.ll
index 09f90ee05735d..caf1944a05d36 100644
--- a/llvm/test/Transforms/InstCombine/bit_ceil.ll
+++ b/llvm/test/Transforms/InstCombine/bit_ceil.ll
@@ -337,6 +337,23 @@ define i32 @test_drop_range_attr(i32 %x) {
   ret i32 %sel
 }
 
+define i33 @test_bit_ceil_i33_non_pow2(i33 %x) {
+; CHECK-LABEL: @test_bit_ceil_i33_non_pow2(
+; CHECK-NEXT:    [[CTLZ:%.*]] = call range(i33 0, 34) i33 @llvm.ctlz.i33(i33 [[X:%.*]], i1 false)
+; CHECK-NEXT:    [[TMP1:%.*]] = sub nsw i33 0, [[CTLZ]]
+; CHECK-NEXT:    [[TMP2:%.*]] = and i33 [[TMP1]], 32
+; CHECK-NEXT:    [[SEL:%.*]] = shl nuw i33 1, [[TMP2]]
+; CHECK-NEXT:    ret i33 [[SEL]]
+;
+  %ctlz = call i33 @llvm.ctlz.i33(i33 %x, i1 false)
+  %sub = sub i33 33, %ctlz
+  %shl = shl i33 1, %sub
+  %dec = add i33 %x, -1
+  %ult = icmp ult i33 %dec, -2
+  %sel = select i1 %ult, i33 %shl, i33 1
+  ret i33 %sel
+}
+
 define i32 @bit_ceil_plus_nsw(i32 %x) {
 ; CHECK-LABEL: @bit_ceil_plus_nsw(
 ; CHECK-NEXT:  entry:
@@ -378,5 +395,6 @@ entry:
 }
 
 declare i32 @llvm.ctlz.i32(i32, i1 immarg)
+declare i33 @llvm.ctlz.i33(i33, i1 immarg)
 declare i64 @llvm.ctlz.i64(i64, i1 immarg)
 declare <4 x i32> @llvm.ctlz.v4i32(<4 x i32>, i1)

>From 2883bdae45149c1edc14c0345f8edbbd2874d35f Mon Sep 17 00:00:00 2001
From: cardigan1008 <ybni at cse.cuhk.edu.hk>
Date: Mon, 29 Dec 2025 17:15:13 +0800
Subject: [PATCH 2/2] [InstCombine] Disable non-power-of-two integer widths in
 `foldBitCeil`.

---
 llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp | 3 +++
 llvm/test/Transforms/InstCombine/bit_ceil.ll          | 8 +++++---
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 67d1845832725..b7ab642cde6ea 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -3871,6 +3871,9 @@ static Instruction *foldBitCeil(SelectInst &SI, IRBuilderBase &Builder,
                                    ShouldDropNoWrap))
     return nullptr;
 
+  if (!isPowerOf2_32(BitWidth))
+    return nullptr;
+
   if (ShouldDropNoWrap) {
     cast<Instruction>(CtlzOp)->setHasNoUnsignedWrap(false);
     cast<Instruction>(CtlzOp)->setHasNoSignedWrap(false);
diff --git a/llvm/test/Transforms/InstCombine/bit_ceil.ll b/llvm/test/Transforms/InstCombine/bit_ceil.ll
index caf1944a05d36..edf1c176ff12c 100644
--- a/llvm/test/Transforms/InstCombine/bit_ceil.ll
+++ b/llvm/test/Transforms/InstCombine/bit_ceil.ll
@@ -340,10 +340,12 @@ define i32 @test_drop_range_attr(i32 %x) {
 define i33 @test_bit_ceil_i33_non_pow2(i33 %x) {
 ; CHECK-LABEL: @test_bit_ceil_i33_non_pow2(
 ; CHECK-NEXT:    [[CTLZ:%.*]] = call range(i33 0, 34) i33 @llvm.ctlz.i33(i33 [[X:%.*]], i1 false)
-; CHECK-NEXT:    [[TMP1:%.*]] = sub nsw i33 0, [[CTLZ]]
-; CHECK-NEXT:    [[TMP2:%.*]] = and i33 [[TMP1]], 32
+; CHECK-NEXT:    [[TMP2:%.*]] = sub nuw nsw i33 33, [[CTLZ]]
 ; CHECK-NEXT:    [[SEL:%.*]] = shl nuw i33 1, [[TMP2]]
-; CHECK-NEXT:    ret i33 [[SEL]]
+; CHECK-NEXT:    [[DEC:%.*]] = add i33 [[X]], -1
+; CHECK-NEXT:    [[ULT:%.*]] = icmp ult i33 [[DEC]], -2
+; CHECK-NEXT:    [[SEL1:%.*]] = select i1 [[ULT]], i33 [[SEL]], i33 1
+; CHECK-NEXT:    ret i33 [[SEL1]]
 ;
   %ctlz = call i33 @llvm.ctlz.i33(i33 %x, i1 false)
   %sub = sub i33 33, %ctlz