[ValueTracking] Let ComputeKnownSignBits handle (shl (zext X), C)

[bjope]

Add simple support for looking through a zext when doing ComputeKnownSignBits for shl. This is valid for the case when all extended bits are shifted out, because then the number of sign bits can be found by analysing the zext operand.

The solution here is simple as it only handle a single zext (not passing remaining left shift amount during recursion). It could be possible to generalize this in the future by for example passing an 'OffsetFromMSB' parameter to ComputeNumSignBitsImpl, telling it to calculate number of sign bits starting at some offset from the most significant bit.

[llvmbot]

@llvm/pr-subscribers-llvm-analysis

Author: Björn Pettersson (bjope)

Changes

Add simple support for looking through a zext when doing ComputeKnownSignBits for shl. This is valid for the case when all extended bits are shifted out, because then the number of sign bits can be found by analysing the zext operand.

The solution here is simple as it only handle a single zext (not passing remaining left shift amount during recursion). It could be possible to generalize this in the future by for example passing an 'OffsetFromMSB' parameter to ComputeNumSignBitsImpl, telling it to calculate number of sign bits starting at some offset from the most significant bit.

---

Full diff: https://github.com/llvm/llvm-project/pull/97693.diff

2 Files Affected:

• (modified) llvm/lib/Analysis/ValueTracking.cpp (+13-3)
• (modified) llvm/unittests/Analysis/ValueTrackingTest.cpp (+30)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 5476dc5d851829..7d229f58d3f6a9 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -3757,11 +3757,21 @@ static unsigned ComputeNumSignBitsImpl(const Value *V,
     }
     case Instruction::Shl: {
       const APInt *ShAmt;
+      Value *X = nullptr;
       if (match(U->getOperand(1), m_APInt(ShAmt))) {
         // shl destroys sign bits.
-        Tmp = ComputeNumSignBits(U->getOperand(0), Depth + 1, Q);
-        if (ShAmt->uge(TyBits) ||   // Bad shift.
-            ShAmt->uge(Tmp)) break; // Shifted all sign bits out.
+        if (ShAmt->uge(TyBits))
+          break; // Bad shift.
+        // We can look through a zext (more or less treating it as a sext) if
+        // all extended bits are shifted out.
+        if (match(U->getOperand(0), m_ZExt(m_Value(X))) &&
+            ShAmt->uge(TyBits - X->getType()->getScalarSizeInBits())) {
+          Tmp = ComputeNumSignBits(X, Depth + 1, Q);
+          Tmp += TyBits - X->getType()->getScalarSizeInBits();
+        } else
+          Tmp = ComputeNumSignBits(U->getOperand(0), Depth + 1, Q);
+        if (ShAmt->uge(Tmp))
+          break; // Shifted all sign bits out.
         Tmp2 = ShAmt->getZExtValue();
         return Tmp - Tmp2;
       }
diff --git a/llvm/unittests/Analysis/ValueTrackingTest.cpp b/llvm/unittests/Analysis/ValueTrackingTest.cpp
index a30db468c77291..e850338ff8ff1e 100644
--- a/llvm/unittests/Analysis/ValueTrackingTest.cpp
+++ b/llvm/unittests/Analysis/ValueTrackingTest.cpp
@@ -697,6 +697,36 @@ TEST_F(ValueTrackingTest, ComputeNumSignBits_PR32045) {
   EXPECT_EQ(ComputeNumSignBits(A, M->getDataLayout()), 32u);
 }
 
+TEST_F(ValueTrackingTest, ComputeNumSignBits_shl_ext1) {
+  parseAssembly("define i32 @test(i8 %a) {\n"
+                "  %b = ashr i8 %a, 4\n"
+                "  %c = zext i8 %b to i32\n"
+                "  %A = shl i32 %c, 24\n"
+                "  ret i32 %A\n"
+                "}\n");
+  EXPECT_EQ(ComputeNumSignBits(A, M->getDataLayout()), 5u);
+}
+
+TEST_F(ValueTrackingTest, ComputeNumSignBits_shl_ext2) {
+  parseAssembly("define i32 @test(i8 %a) {\n"
+                "  %b = ashr i8 %a, 4\n"
+                "  %c = zext i8 %b to i32\n"
+                "  %A = shl i32 %c, 26\n"
+                "  ret i32 %A\n"
+                "}\n");
+  EXPECT_EQ(ComputeNumSignBits(A, M->getDataLayout()), 3u);
+}
+
+TEST_F(ValueTrackingTest, ComputeNumSignBits_shl_ext3) {
+  parseAssembly("define i32 @test(i8 %a) {\n"
+                "  %b = ashr i8 %a, 4\n"
+                "  %c = zext i8 %b to i32\n"
+                "  %A = shl i32 %c, 30\n"
+                "  ret i32 %A\n"
+                "}\n");
+  EXPECT_EQ(ComputeNumSignBits(A, M->getDataLayout()), 1u);
+}
+
 // No guarantees for canonical IR in this analysis, so this just bails out.
 TEST_F(ValueTrackingTest, ComputeNumSignBits_Shuffle) {
   parseAssembly(

[bjope]

SelectionDAG variant is here: #97695

[nikic]

This looks reasonable to me. Can you replace the unit tests with lit tests please?

[bjope]

This looks reasonable to me. Can you replace the unit tests with lit tests please?

I don't really know how to make good lit tests for ComputeNumSignBits.
All I've got is something like this

; RUN: opt -passes=instcombine -S < %s | FileCheck %s

define i16 @numsignbits_shl_zext(i8 %a) {
  %nsb6 = ashr i8 %a, 6
  %zext = zext i8 %nsb6 to i16
  %nsb4 = shl i16 %zext, 10
  %or = or i16 %nsb4, 1
  ; With four sign bits this should be nsw.
  %shl = add i16 %nsb4, 4096
  ret i16 %shl
}

to show that the add will get a nsw. But it far from verifies that ComputeNumSignBits actually is returning the correct value. Isn't it much easier to verify the correct behavior via the unit test?
Maybe I should add a lit test like the one above as a complement, to show that it is possible to trigger this via a lit test as well (even though it doesn't validate that the calculation in ComputeNumSignBits is correct).

[goldsteinn]

This looks reasonable to me. Can you replace the unit tests with lit tests please?

I don't really know how to make good lit tests for ComputeNumSignBits. All I've got is something like this

; RUN: opt -passes=instcombine -S < %s | FileCheck %s

define i16 @numsignbits_shl_zext(i8 %a) {
  %nsb6 = ashr i8 %a, 6
  %zext = zext i8 %nsb6 to i16
  %nsb4 = shl i16 %zext, 10
  %or = or i16 %nsb4, 1
  ; With four sign bits this should be nsw.
  %shl = add i16 %nsb4, 4096
  ret i16 %shl
}

to show that the add will get a nsw. But it far from verifies that ComputeNumSignBits actually is returning the correct value. Isn't it much easier to verify the correct behavior via the unit test? Maybe I should add a lit test like the one above as a complement, to show that it is possible to trigger this via a lit test as well (even though it doesn't validate that the calculation in ComputeNumSignBits is correct).

You could use the:

  // (A & 2^C1) + A => A & (2^C1 - 1) iff bit C1 in A is a sign bit
  if (match(&I, m_c_Add(m_And(m_Value(A), m_APInt(C1)), m_Deferred(A))) &&
      C1->isPowerOf2() && (ComputeNumSignBits(A) > C1->countl_zero())) {
    Constant *NewMask = ConstantInt::get(RHS->getType(), *C1 - 1);
    return BinaryOperator::CreateAnd(A, NewMask);
  }

fold and be precise (although typicaly I would do this with nsw flag on shl, but that doesn't really work here).

[goldsteinn]

Probably should have a test (here or in lit) where ShAmt < 24

[bjope]

Replaced unit test case with a lit based test (thanks to @goldsteinn for great tip about the and+add fold!).

[nikic]

LGTM

[dtcxzyw]

LGTM. Thank you!