[ConstraintElim] Opimize abs based on known constraints

Author: el-ev

llvm/lib/Transforms/Scalar/ConstraintElimination.cpp

@@ -1128,20 +1128,19 @@ void State::addInfoFor(BasicBlock &BB) {
           FactOrCheck::getCheck(DT.getNode(&BB), cast<CallInst>(&I)));
       break;
     // Enqueue the intrinsics to add extra info.
+    case Intrinsic::abs:
     case Intrinsic::umin:
     case Intrinsic::umax:
     case Intrinsic::smin:
     case Intrinsic::smax:
     case Intrinsic::usub_sat:
-    // TODO: handle llvm.abs as well
       WorkList.push_back(
           FactOrCheck::getCheck(DT.getNode(&BB), cast<CallInst>(&I)));
       // TODO: Check if it is possible to instead only added the min/max facts
       // when simplifying uses of the min/max intrinsics.
       if (!isGuaranteedNotToBePoison(&I))
         break;
       [[fallthrough]];
-    case Intrinsic::abs:
     case Intrinsic::uadd_sat:
       WorkList.push_back(FactOrCheck::getInstFact(DT.getNode(&BB), &I));
       break;
@@ -1537,6 +1536,28 @@ static bool checkAndReplaceUSubSat(SaturatingInst *I, ConstraintInfo &Info,
   return false;
 }
 
+static bool checkAndReplaceAbs(IntrinsicInst *I, ConstraintInfo &Info,
+                               SmallVectorImpl<Instruction *> &ToRemove) {
+  assert(I->getIntrinsicID() == Intrinsic::abs && "Expected abs intrinsic");
+  Value *Op = I->getOperand(0);
+  if (checkCondition(ICmpInst::ICMP_SGE, Op, ConstantInt::get(Op->getType(), 0),
+                     I, Info)
+          .value_or(false)) {
+    I->replaceAllUsesWith(Op);
+    ToRemove.push_back(I);
+    return true;
+  }
+  if (checkCondition(ICmpInst::ICMP_SLT, Op, ConstantInt::get(Op->getType(), 0),
+                     I, Info)
+          .value_or(false)) {
+    IRBuilder<> Builder(I->getParent(), I->getIterator());
+    I->replaceAllUsesWith(Builder.CreateNeg(Op));
+    ToRemove.push_back(I);
+    return true;
+  }
+  return false;
+}
+
 static void
 removeEntryFromStack(const StackEntry &E, ConstraintInfo &Info,
                      Module *ReproducerModule,
@@ -1863,6 +1884,10 @@ static bool eliminateConstraints(Function &F, DominatorTree &DT, LoopInfo &LI,
           Changed |= checkAndReplaceUSubSat(SatIntr, Info, ToRemove);
         else
           llvm_unreachable("Unexpected intrinsic.");
+      } else if (auto *II = dyn_cast<IntrinsicInst>(Inst)) {
+        if (II->getIntrinsicID() == Intrinsic::abs) {
+          Changed |= checkAndReplaceAbs(II, Info, ToRemove);
+        }
       }
       continue;
     }

llvm/test/Transforms/ConstraintElimination/abs.ll

@@ -159,5 +159,48 @@ define i1 @abs_is_nonnegative_constant_arg() {
   ret i1 %cmp
 }
 
+define i64 @abs_assume_nonnegative(i64 %arg) {
+; CHECK-LABEL: define i64 @abs_assume_nonnegative(
+; CHECK-SAME: i64 [[ARG:%.*]]) {
+; CHECK-NEXT:    [[PRECOND:%.*]] = icmp sge i64 [[ARG]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[PRECOND]])
+; CHECK-NEXT:    [[ABS:%.*]] = tail call i64 @llvm.abs.i64(i64 [[ARG]], i1 false)
+; CHECK-NEXT:    ret i64 [[ABS]]
+;
+  %precond = icmp sge i64 %arg, 0
+  call void @llvm.assume(i1 %precond)
+  %abs = tail call i64 @llvm.abs.i64(i64 %arg, i1 false)
+  ret i64 %abs
+}
+
+define i64 @abs_assume_negative(i64 %arg) {
+; CHECK-LABEL: define i64 @abs_assume_negative(
+; CHECK-SAME: i64 [[ARG:%.*]]) {
+; CHECK-NEXT:    [[PRECOND:%.*]] = icmp slt i64 [[ARG]], 0
+; CHECK-NEXT:    call void @llvm.assume(i1 [[PRECOND]])
+; CHECK-NEXT:    [[ABS:%.*]] = tail call i64 @llvm.abs.i64(i64 [[ARG]], i1 false)
+; CHECK-NEXT:    ret i64 [[ABS]]
+;
+  %precond = icmp slt i64 %arg, 0
+  call void @llvm.assume(i1 %precond)
+  %abs = tail call i64 @llvm.abs.i64(i64 %arg, i1 false)
+  ret i64 %abs
+}
+
+; Negative test
+define i64 @abs_assume_unrelated(i64 %arg) {
+; CHECK-LABEL: define i64 @abs_assume_unrelated(
+; CHECK-SAME: i64 [[ARG:%.*]]) {
+; CHECK-NEXT:    [[PRECOND:%.*]] = icmp slt i64 [[ARG]], 3
+; CHECK-NEXT:    call void @llvm.assume(i1 [[PRECOND]])
+; CHECK-NEXT:    [[ABS:%.*]] = tail call i64 @llvm.abs.i64(i64 [[ARG]], i1 false)
+; CHECK-NEXT:    ret i64 [[ABS]]
+;
+  %precond = icmp slt i64 %arg, 3
+  call void @llvm.assume(i1 %precond)
+  %abs = tail call i64 @llvm.abs.i64(i64 %arg, i1 false)
+  ret i64 %abs
+}
+
 declare i32 @llvm.abs.i32(i32, i1 immarg)
 declare void @llvm.assume(i1)