[llvm] feat: add flag to disable known bits analysis (PR #228000)

[Luisa Cicolini via llvm-commits]

https://github.com/luisacicolini created https://github.com/llvm/llvm-project/pull/228000

None

>From 80f16248fbee35a7418e5355e4a1ce9685de282a Mon Sep 17 00:00:00 2001
From: luisacicolini <lc985 at cam.ac.uk>
Date: Thu, 1 Oct 2026 10:03:28 +0100
Subject: [PATCH] remove kba

---
 llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp     | 8 ++++++++
 llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp | 8 ++++++++
 llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp         | 7 +++++++
 llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp        | 8 ++++----
 4 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp b/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp
index 598893ef3c0d7dc..b4eb464d528930c 100644
--- a/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/GISelValueTracking.cpp
@@ -35,6 +35,7 @@
 #include "llvm/IR/FMF.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/MC/TargetRegistry.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/KnownFPClass.h"
 #include "llvm/Target/TargetMachine.h"
@@ -44,6 +45,10 @@
 using namespace llvm;
 using namespace MIPatternMatch;
 
+static cl::opt<bool> DisableGISelKnownBits(
+    "gisel-disable-known-bits", cl::Hidden, cl::init(false),
+    cl::desc("Disable GlobalISel known bits analysis (always return unknown)"));
+
 char llvm::GISelValueTrackingAnalysisLegacy::ID = 0;
 
 INITIALIZE_PASS(GISelValueTrackingAnalysisLegacy, DEBUG_TYPE,
@@ -278,6 +283,9 @@ void GISelValueTracking::computeKnownBitsImpl(Register R, KnownBits &Known,
   unsigned BitWidth = DstTy.getScalarSizeInBits();
   Known = KnownBits(BitWidth); // Don't know anything
 
+  if (DisableGISelKnownBits)
+    return;
+
   // Depth may get bigger than max depth if it gets passed to a different
   // GISelValueTracking object.
   // This may happen when say a generic part uses a GISelValueTracking object
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
index 5b6329c2560eb15..fc636fd91acb9e9 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
@@ -3438,6 +3438,14 @@ ExpandShiftWithKnownAmountBit(SDNode *N, SDValue &Lo, SDValue &Hi) {
   APInt HighBitMask = APInt::getHighBitsSet(ShBits, ShBits - Log2_32(NVTBits));
   KnownBits Known = DAG.computeKnownBits(Amt);
 
+  // Recognize a masked shift amount structurally, so that termination of
+  // ExpandIntRes_ShiftThroughStack does not depend on computeKnownBits (which
+  // may be disabled via -dag-disable-known-bits).
+  if (Amt.getOpcode() == ISD::AND)
+    if (auto *C = dyn_cast<ConstantSDNode>(Amt.getOperand(1)))
+      if (!C->getAPIntValue().intersects(HighBitMask))
+        Known.Zero |= HighBitMask;
+
   // If we don't know anything about the high bits, exit.
   if (((Known.Zero | Known.One) & HighBitMask) == 0)
     return false;
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index e22907c88f65437..2c96affdd6920a0 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -120,6 +120,10 @@ static cl::opt<int> VScaleUnrollLimit(
     cl::desc("Maximum vscale for which vector unrolling is allowed."),
     cl::Hidden, cl::init(64));
 
+static cl::opt<bool> DisableDAGKnownBits(
+    "dag-disable-known-bits", cl::Hidden, cl::init(false),
+    cl::desc("Disable SelectionDAG known bits analysis (always return unknown)"));
+
 static void NewSDValueDbgMsg(SDValue V, StringRef Msg, SelectionDAG *G) {
   LLVM_DEBUG(dbgs() << Msg; V.getNode()->dump(G););
 }
@@ -3352,6 +3356,9 @@ KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
 
   KnownBits Known(BitWidth);   // Don't know anything.
 
+  if (DisableDAGKnownBits)
+    return Known;
+
   if (auto OptAPInt = Op->bitcastToAPInt()) {
     // We know all of the bits for a constant!
     return KnownBits::makeConstant(*std::move(OptAPInt));
diff --git a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
index 48c0e37151b9c4d..fcca01562fc106f 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
@@ -3511,10 +3511,10 @@ static bool isBitfieldPositioningOpFromAnd(SelectionDAG *CurDAG, SDValue Op,
   //   1) (AndImm & (1 << POS) == 0)
   //   2) the result of AND is not zero at POS bit (according to NonZeroBits)
   //
-  // 1) and 2) don't agree so something must be wrong (e.g., in
-  // 'SelectionDAG::computeKnownBits')
-  assert((~AndImm & NonZeroBits) == 0 &&
-         "Something must be wrong (e.g., in SelectionDAG::computeKnownBits)");
+  // 1) and 2) don't agree, which can happen if known bits analysis is
+  // imprecise (e.g., disabled via -dag-disable-known-bits). Bail out.
+  if ((~AndImm & NonZeroBits) != 0)
+    return false;
 
   SDValue AndOp0 = Op.getOperand(0);