[llvm] d9e9276 - [ConstantRange] Improve ConstantRange::binaryXor (#80146)

[via llvm-commits]

Author: Yingwei Zheng
Date: 2024-02-08T22:34:52+08:00
New Revision: d9e92765c5f9b0fa7adafa769dd13d37b6bca038

URL: https://github.com/llvm/llvm-project/commit/d9e92765c5f9b0fa7adafa769dd13d37b6bca038
DIFF: https://github.com/llvm/llvm-project/commit/d9e92765c5f9b0fa7adafa769dd13d37b6bca038.diff

LOG: [ConstantRange] Improve ConstantRange::binaryXor (#80146)

`ConstantRange::binaryXor` gives poor results as it currently depends on
`KnownBits::operator^`.
Since `sub A, B` is canonicalized into `xor A, B` if `B` is the subset
of `A`, this patch reverts the transform in `ConstantRange::binaryXor`,
which will give better results.

Alive2: https://alive2.llvm.org/ce/z/bmTMV9
Fixes #79696.

Added: 
    llvm/test/Transforms/SCCP/pr79696.ll

Modified: 
    llvm/lib/IR/ConstantRange.cpp
    llvm/unittests/IR/ConstantRangeTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/IR/ConstantRange.cpp b/llvm/lib/IR/ConstantRange.cpp
index cbb64b299e648..3394a1ec8dc47 100644
--- a/llvm/lib/IR/ConstantRange.cpp
+++ b/llvm/lib/IR/ConstantRange.cpp
@@ -1467,7 +1467,22 @@ ConstantRange ConstantRange::binaryXor(const ConstantRange &Other) const {
   if (isSingleElement() && getSingleElement()->isAllOnes())
     return Other.binaryNot();
 
-  return fromKnownBits(toKnownBits() ^ Other.toKnownBits(), /*IsSigned*/false);
+  KnownBits LHSKnown = toKnownBits();
+  KnownBits RHSKnown = Other.toKnownBits();
+  KnownBits Known = LHSKnown ^ RHSKnown;
+  ConstantRange CR = fromKnownBits(Known, /*IsSigned*/ false);
+  // Typically the following code doesn't improve the result if BW = 1.
+  if (getBitWidth() == 1)
+    return CR;
+
+  // If LHS is known to be the subset of RHS, treat LHS ^ RHS as RHS -nuw/nsw
+  // LHS. If RHS is known to be the subset of LHS, treat LHS ^ RHS as LHS
+  // -nuw/nsw RHS.
+  if ((~LHSKnown.Zero).isSubsetOf(RHSKnown.One))
+    CR = CR.intersectWith(Other.sub(*this), PreferredRangeType::Unsigned);
+  else if ((~RHSKnown.Zero).isSubsetOf(LHSKnown.One))
+    CR = CR.intersectWith(this->sub(Other), PreferredRangeType::Unsigned);
+  return CR;
 }
 
 ConstantRange

diff  --git a/llvm/test/Transforms/SCCP/pr79696.ll b/llvm/test/Transforms/SCCP/pr79696.ll
new file mode 100644
index 0000000000000..a860112d5ef36
--- /dev/null
+++ b/llvm/test/Transforms/SCCP/pr79696.ll
@@ -0,0 +1,55 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt < %s -passes=ipsccp -S | FileCheck %s
+
+; Tests from PR79696
+
+define i1 @constant_range_xor(i64 %a) {
+; CHECK-LABEL: define i1 @constant_range_xor(
+; CHECK-SAME: i64 [[A:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i64 [[A]], 8192
+; CHECK-NEXT:    br i1 [[CMP]], label [[THEN:%.*]], label [[ELSE:%.*]]
+; CHECK:       then:
+; CHECK-NEXT:    [[CTLZ:%.*]] = call i64 @llvm.ctlz.i64(i64 [[A]], i1 true)
+; CHECK-NEXT:    [[CONV:%.*]] = xor i64 [[CTLZ]], 63
+; CHECK-NEXT:    ret i1 false
+; CHECK:       else:
+; CHECK-NEXT:    ret i1 false
+;
+entry:
+  %cmp = icmp ugt i64 %a, 8192
+  br i1 %cmp, label %then, label %else
+then:
+  %ctlz = call i64 @llvm.ctlz.i64(i64 %a, i1 true) ;[0, 50]
+  %conv = xor i64 %ctlz, 63                        ;[13, 63]
+  %cmp1 = icmp ult i64 %conv, 13
+  ret i1 %cmp1
+else:
+  ret i1 false
+}
+
+define i1 @constant_range_xor_negative(i64 %a) {
+; CHECK-LABEL: define i1 @constant_range_xor_negative(
+; CHECK-SAME: i64 [[A:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i64 [[A]], 8192
+; CHECK-NEXT:    br i1 [[CMP]], label [[THEN:%.*]], label [[ELSE:%.*]]
+; CHECK:       then:
+; CHECK-NEXT:    [[CTLZ:%.*]] = call i64 @llvm.ctlz.i64(i64 [[A]], i1 true)
+; CHECK-NEXT:    [[CONV:%.*]] = xor i64 [[CTLZ]], 62
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp ult i64 [[CONV]], 13
+; CHECK-NEXT:    ret i1 [[CMP1]]
+; CHECK:       else:
+; CHECK-NEXT:    ret i1 false
+;
+entry:
+  %cmp = icmp ugt i64 %a, 8192
+  br i1 %cmp, label %then, label %else
+then:
+  %ctlz = call i64 @llvm.ctlz.i64(i64 %a, i1 true) ;[0, 50]
+  %conv = xor i64 %ctlz, 62                        ;[12, 63]
+  %cmp1 = icmp ult i64 %conv, 13
+  ret i1 %cmp1
+else:
+  ret i1 false
+}

diff  --git a/llvm/unittests/IR/ConstantRangeTest.cpp b/llvm/unittests/IR/ConstantRangeTest.cpp
index e505af5d3275e..34a162a5514e9 100644
--- a/llvm/unittests/IR/ConstantRangeTest.cpp
+++ b/llvm/unittests/IR/ConstantRangeTest.cpp
@@ -2565,6 +2565,12 @@ TEST_F(ConstantRangeTest, binaryXor) {
   EXPECT_EQ(R16_35.binaryXor(R0_99), ConstantRange(APInt(8, 0), APInt(8, 128)));
   EXPECT_EQ(R0_99.binaryXor(R16_35), ConstantRange(APInt(8, 0), APInt(8, 128)));
 
+  // Treat xor A, B as sub nsw nuw A, B
+  ConstantRange R0_51(APInt(8, 0), APInt(8, 51));
+  ConstantRange R63(APInt(8, 63));
+  EXPECT_EQ(R0_51.binaryXor(R63), ConstantRange(APInt(8, 13), APInt(8, 64)));
+  EXPECT_EQ(R63.binaryXor(R0_51), ConstantRange(APInt(8, 13), APInt(8, 64)));
+
   TestBinaryOpExhaustive(
       [](const ConstantRange &CR1, const ConstantRange &CR2) {
         return CR1.binaryXor(CR2);