[IR] BlockAddress doesn't use BasicBlock (#200772)

BlockAddress is the only non-terminator user of a BasicBlock, and it
occurs very rarely. To speed up predecessor iteration, change
BlockAddress to no longer use its BasicBlock.

This should also make uselistorder_bb obsolete.

llvm/docs/LangRef.rst

@@ -4325,16 +4325,11 @@ Use-list directives may appear at function scope or global scope. They are not
 instructions, and have no effect on the semantics of the IR. When they're at
 function scope, they must appear after the terminator of the final basic block.
 
-If basic blocks have their address taken via ``blockaddress()`` expressions,
-``uselistorder_bb`` can be used to reorder their use-lists from outside their
-function's scope.
-
 :Syntax:
 
 ::
 
     uselistorder <ty> <value>, { <order-indexes> }
-    uselistorder_bb @function, %block { <order-indexes> }
 
 :Examples:
 
@@ -4355,7 +4350,6 @@ function's scope.
     uselistorder ptr @global, { 1, 2, 0 }
     uselistorder i32 7, { 1, 0 }
     uselistorder i32 (i32) @bar, { 1, 0 }
-    uselistorder_bb @foo, %bb, { 5, 1, 3, 2, 0, 4 }
 
 .. _source_filename:
 

llvm/include/llvm/AsmParser/LLParser.h

@@ -706,7 +706,6 @@ namespace llvm {
 
     // Use-list order directives.
     bool parseUseListOrder(PerFunctionState *PFS = nullptr);
-    bool parseUseListOrderBB();
     bool parseUseListOrderIndexes(SmallVectorImpl<unsigned> &Indexes);
     bool sortUseListOrder(Value *V, ArrayRef<unsigned> Indexes, SMLoc Loc);
   };

llvm/include/llvm/AsmParser/LLToken.h

@@ -388,7 +388,6 @@ enum Kind {
 
   // Use-list order directives.
   kw_uselistorder,
-  kw_uselistorder_bb,
 
   // Summary index keywords
   kw_path,

llvm/include/llvm/IR/CFG.h

@@ -51,23 +51,9 @@ protected:
   using Self = PredIterator<Ptr, USE_iterator>;
   USE_iterator It;
 
-  inline void advancePastNonTerminators() {
-    // Loop to ignore non-terminator uses (for example BlockAddresses).
-    while (!It.atEnd()) {
-      if (auto *Inst = dyn_cast<Instruction>(*It)) {
-        assert(Inst->isTerminator() && "BasicBlock used in non-terminator");
-        break;
-      }
-
-      ++It;
-    }
-  }
-
 public:
   PredIterator() = default;
-  explicit inline PredIterator(Ptr *bb) : It(bb->user_begin()) {
-    advancePastNonTerminators();
-  }
+  explicit inline PredIterator(Ptr *bb) : It(bb->user_begin()) {}
   inline PredIterator(Ptr *bb, bool) : It(bb->user_end()) {}
 
   inline bool operator==(const Self& x) const { return It == x.It; }
@@ -75,13 +61,15 @@ public:
 
   inline reference operator*() const {
     assert(!It.atEnd() && "pred_iterator out of range!");
-    return cast<Instruction>(*It)->getParent();
+    auto *I = cast<Instruction>(*It);
+    assert(I->isTerminator() && "BasicBlock used in non-terminator");
+    return I->getParent();
   }
   inline pointer *operator->() const { return &operator*(); }
 
   inline Self& operator++() {   // Preincrement
     assert(!It.atEnd() && "pred_iterator out of range!");
-    ++It; advancePastNonTerminators();
+    ++It;
     return *this;
   }
 

llvm/include/llvm/IR/Constants.h

@@ -1082,7 +1082,9 @@ public:
 class BlockAddress final : public Constant {
   friend class Constant;
 
-  constexpr static IntrusiveOperandsAllocMarker AllocMarker{1};
+  constexpr static IntrusiveOperandsAllocMarker AllocMarker{0};
+
+  BasicBlock *Block;
 
   BlockAddress(Type *Ty, BasicBlock *BB);
 
@@ -1114,7 +1116,7 @@ public:
   /// Transparently provide more efficient getOperand methods.
   DECLARE_TRANSPARENT_OPERAND_ACCESSORS(Value);
 
-  BasicBlock *getBasicBlock() const { return cast<BasicBlock>(Op<0>().get()); }
+  BasicBlock *getBasicBlock() const { return Block; }
   Function *getFunction() const { return getBasicBlock()->getParent(); }
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast:
@@ -1125,7 +1127,7 @@ public:
 
 template <>
 struct OperandTraits<BlockAddress>
-    : public FixedNumOperandTraits<BlockAddress, 1> {};
+    : public FixedNumOperandTraits<BlockAddress, 0> {};
 
 DEFINE_TRANSPARENT_OPERAND_ACCESSORS(BlockAddress, Value)
 

llvm/lib/AsmParser/LLLexer.cpp

@@ -801,7 +801,6 @@ lltok::Kind LLLexer::LexIdentifier() {
 
   // Use-list order directives.
   KEYWORD(uselistorder);
-  KEYWORD(uselistorder_bb);
 
   KEYWORD(personality);
   KEYWORD(cleanup);

llvm/lib/AsmParser/LLParser.cpp

@@ -634,10 +634,6 @@ bool LLParser::parseTopLevelEntities() {
       if (parseUseListOrder())
         return true;
       break;
-    case lltok::kw_uselistorder_bb:
-      if (parseUseListOrderBB())
-        return true;
-      break;
     }
   }
 }
@@ -9472,53 +9468,6 @@ bool LLParser::parseUseListOrder(PerFunctionState *PFS) {
   return sortUseListOrder(V, Indexes, Loc);
 }
 
-/// parseUseListOrderBB
-///   ::= 'uselistorder_bb' @foo ',' %bar ',' UseListOrderIndexes
-bool LLParser::parseUseListOrderBB() {
-  assert(Lex.getKind() == lltok::kw_uselistorder_bb);
-  SMLoc Loc = Lex.getLoc();
-  Lex.Lex();
-
-  ValID Fn, Label;
-  SmallVector<unsigned, 16> Indexes;
-  if (parseValID(Fn, /*PFS=*/nullptr) ||
-      parseToken(lltok::comma, "expected comma in uselistorder_bb directive") ||
-      parseValID(Label, /*PFS=*/nullptr) ||
-      parseToken(lltok::comma, "expected comma in uselistorder_bb directive") ||
-      parseUseListOrderIndexes(Indexes))
-    return true;
-
-  // Check the function.
-  GlobalValue *GV;
-  if (Fn.Kind == ValID::t_GlobalName)
-    GV = M->getNamedValue(Fn.StrVal);
-  else if (Fn.Kind == ValID::t_GlobalID)
-    GV = NumberedVals.get(Fn.UIntVal);
-  else
-    return error(Fn.Loc, "expected function name in uselistorder_bb");
-  if (!GV)
-    return error(Fn.Loc,
-                 "invalid function forward reference in uselistorder_bb");
-  auto *F = dyn_cast<Function>(GV);
-  if (!F)
-    return error(Fn.Loc, "expected function name in uselistorder_bb");
-  if (F->isDeclaration())
-    return error(Fn.Loc, "invalid declaration in uselistorder_bb");
-
-  // Check the basic block.
-  if (Label.Kind == ValID::t_LocalID)
-    return error(Label.Loc, "invalid numeric label in uselistorder_bb");
-  if (Label.Kind != ValID::t_LocalName)
-    return error(Label.Loc, "expected basic block name in uselistorder_bb");
-  Value *V = F->getValueSymbolTable()->lookup(Label.StrVal);
-  if (!V)
-    return error(Label.Loc, "invalid basic block in uselistorder_bb");
-  if (!isa<BasicBlock>(V))
-    return error(Label.Loc, "expected basic block in uselistorder_bb");
-
-  return sortUseListOrder(V, Indexes, Loc);
-}
-
 /// ModuleEntry
 ///   ::= 'module' ':' '(' 'path' ':' STRINGCONSTANT ',' 'hash' ':' Hash ')'
 /// Hash ::= '(' UInt32 ',' UInt32 ',' UInt32 ',' UInt32 ',' UInt32 ')'

llvm/lib/CodeGen/IndirectBrExpandPass.cpp

@@ -135,23 +135,11 @@ bool runImpl(Function &F, const TargetLowering *TLI, DomTreeUpdater *DTU) {
     if (!IndirectBrSuccs.count(&BB))
       continue;
 
-    auto IsBlockAddressUse = [&](const Use &U) {
-      return isa<BlockAddress>(U.getUser());
-    };
-    auto BlockAddressUseIt = llvm::find_if(BB.uses(), IsBlockAddressUse);
-    if (BlockAddressUseIt == BB.use_end())
-      continue;
-
-    assert(std::none_of(std::next(BlockAddressUseIt), BB.use_end(),
-                        IsBlockAddressUse) &&
-           "There should only ever be a single blockaddress use because it is "
-           "a constant and should be uniqued.");
-
-    auto *BA = cast<BlockAddress>(BlockAddressUseIt->getUser());
+    auto *BA = BlockAddress::lookup(&BB);
 
     // Skip if the constant was formed but ended up not being used (due to DCE
     // or whatever).
-    if (!BA->isConstantUsed())
+    if (!BA || !BA->isConstantUsed())
       continue;
 
     // Compute the index we want to use for this basic block. We can't use zero

llvm/lib/IR/AsmWriter.cpp

@@ -5038,20 +5038,11 @@ void AssemblyWriter::writeAllAttributeGroups() {
 
 void AssemblyWriter::printUseListOrder(const Value *V,
                                        ArrayRef<unsigned> Shuffle) {
-  bool IsInFunction = Machine.getFunction();
-  if (IsInFunction)
+  if (Machine.getFunction())
     Out << "  ";
 
-  Out << "uselistorder";
-  if (const BasicBlock *BB = IsInFunction ? nullptr : dyn_cast<BasicBlock>(V)) {
-    Out << "_bb ";
-    writeOperand(BB->getParent(), false);
-    Out << ", ";
-    writeOperand(BB, false);
-  } else {
-    Out << " ";
-    writeOperand(V, true);
-  }
+  Out << "uselistorder ";
+  writeOperand(V, true);
 
   assert(Shuffle.size() >= 2 && "Shuffle too small");
   Out << ", { " << llvm::interleaved(Shuffle) << " }\n";

llvm/lib/IR/BasicBlock.cpp

@@ -175,8 +175,7 @@ BasicBlock::~BasicBlock() {
   // is no indirect branch).  Handle these cases by zapping the BlockAddress
   // nodes.  There are no other possible uses at this point.
   if (hasAddressTaken()) {
-    assert(!use_empty() && "There should be at least one blockaddress!");
-    BlockAddress *BA = cast<BlockAddress>(user_back());
+    BlockAddress *BA = BlockAddress::lookup(this);
 
     Constant *Replacement = ConstantInt::get(Type::getInt32Ty(getContext()), 1);
     BA->replaceAllUsesWith(

llvm/lib/IR/Constants.cpp

@@ -2056,7 +2056,7 @@ BlockAddress *BlockAddress::get(Function *F, BasicBlock *BB) {
 
 BlockAddress::BlockAddress(Type *Ty, BasicBlock *BB)
     : Constant(Ty, Value::BlockAddressVal, AllocMarker) {
-  setOperand(0, BB);
+  Block = BB;
   BB->setHasAddressTaken(true);
 }
 
@@ -2089,7 +2089,7 @@ Value *BlockAddress::handleOperandChangeImpl(Value *From, Value *To) {
   // erase invalidates iterators/references, hence the duplicate NewBB lookup.
   getContext().pImpl->BlockAddresses.erase(getBasicBlock());
   getContext().pImpl->BlockAddresses[NewBB] = this;
-  setOperand(0, NewBB);
+  Block = NewBB;
   getBasicBlock()->setHasAddressTaken(true);
 
   // If we just want to keep the existing value, then return null.

llvm/lib/IR/Value.cpp

@@ -542,8 +542,11 @@ void Value::doRAUW(Value *New, ReplaceMetadataUses ReplaceMetaUses) {
     U.set(New);
   }
 
-  if (BasicBlock *BB = dyn_cast<BasicBlock>(this))
+  if (BasicBlock *BB = dyn_cast<BasicBlock>(this)) {
     BB->replaceSuccessorsPhiUsesWith(cast<BasicBlock>(New));
+    if (BB->hasAddressTaken())
+      BlockAddress::lookup(BB)->handleOperandChange(this, New);
+  }
 }
 
 void Value::replaceAllUsesWith(Value *New) {

llvm/test/Assembler/invalid-uselistorder_bb-float-literal.ll

@@ -1,11 +0,0 @@
-; RUN: not llvm-as < %s -disable-output 2>&1 | FileCheck %s
-; CHECK: error: unexpected floating-point literal
-
-@ba1 = constant ptr blockaddress (@foo, %1)
-
-define void @foo() {
-  br label %1
-  unreachable
-}
-
-uselistorder_bb 1.0, %1, { 1, 0 }

llvm/test/Assembler/invalid-uselistorder_bb-missing-bb.ll

@@ -1,6 +0,0 @@
-; RUN: not llvm-as < %s -disable-output 2>&1 | FileCheck %s
-; CHECK: error: invalid basic block in uselistorder_bb
-define void @foo() {
-  unreachable
-}
-uselistorder_bb @foo, %bb, { 1, 0 }

llvm/test/Assembler/invalid-uselistorder_bb-missing-body.ll

@@ -1,4 +0,0 @@
-; RUN: not llvm-as < %s -disable-output 2>&1 | FileCheck %s
-; CHECK: error: invalid declaration in uselistorder_bb
-declare void @foo()
-uselistorder_bb @foo, %bb, { 1, 0 }

llvm/test/Assembler/invalid-uselistorder_bb-missing-func.ll

@@ -1,3 +0,0 @@
-; RUN: not llvm-as < %s -disable-output 2>&1 | FileCheck %s
-; CHECK: error: invalid function forward reference in uselistorder_bb
-uselistorder_bb @foo, %bb, { 1, 0 }

llvm/test/Assembler/invalid-uselistorder_bb-not-bb.ll

@@ -1,6 +0,0 @@
-; RUN: not llvm-as < %s -disable-output 2>&1 | FileCheck %s
-; CHECK: error: expected basic block in uselistorder_bb
-define i32 @foo(i32 %arg) {
-  ret i32 %arg
-}
-uselistorder_bb @foo, %arg, { 1, 0 }

llvm/test/Assembler/invalid-uselistorder_bb-not-func.ll

@@ -1,4 +0,0 @@
-; RUN: not llvm-as < %s -disable-output 2>&1 | FileCheck %s
-; CHECK: error: expected function name in uselistorder_bb
-@global = global i1 0
-uselistorder_bb @global, %bb, { 1, 0 }

llvm/test/Assembler/invalid-uselistorder_bb-numbered.ll

@@ -1,11 +0,0 @@
-; RUN: not llvm-as < %s -disable-output 2>&1 | FileCheck %s
-; CHECK: error: invalid numeric label in uselistorder_bb
-
-@ba1 = constant ptr blockaddress (@foo, %1)
-
-define void @foo() {
-  br label %1
-  unreachable
-}
-
-uselistorder_bb @foo, %1, { 1, 0 }

llvm/test/Assembler/uselistorder_bb.ll

@@ -1,42 +0,0 @@
-; RUN: llvm-as < %s -disable-output 2>&1 | FileCheck %s -allow-empty
-; CHECK-NOT: error
-; CHECK-NOT: warning
-; RUN: verify-uselistorder < %s
-
-@ba1 = constant ptr blockaddress (@bafunc1, %bb)
-@ba2 = constant ptr getelementptr (i8, ptr blockaddress (@bafunc2, %bb), i61 0)
-@ba3 = constant ptr getelementptr (i8, ptr blockaddress (@bafunc2, %bb), i61 0)
-
-define ptr @babefore() {
-  ret ptr getelementptr (i8, ptr blockaddress (@bafunc2, %bb), i61 0)
-bb1:
-  ret ptr blockaddress (@bafunc1, %bb)
-bb2:
-  ret ptr blockaddress (@bafunc3, %bb)
-}
-define void @bafunc1() {
-  br label %bb
-bb:
-  unreachable
-}
-define void @bafunc2() {
-  br label %bb
-bb:
-  unreachable
-}
-define void @bafunc3() {
-  br label %bb
-bb:
-  unreachable
-}
-define ptr @baafter() {
-  ret ptr blockaddress (@bafunc2, %bb)
-bb1:
-  ret ptr blockaddress (@bafunc1, %bb)
-bb2:
-  ret ptr blockaddress (@bafunc3, %bb)
-}
-
-uselistorder_bb @bafunc1, %bb, { 1, 0 }
-uselistorder_bb @bafunc2, %bb, { 1, 0 }
-uselistorder_bb @bafunc3, %bb, { 1, 0 }

llvm/utils/emacs/llvm-mode.el

@@ -74,7 +74,7 @@
    `(,(regexp-opt
        '(;; Toplevel entities
          "declare" "define" "module" "target" "source_filename" "global" "constant" "const" "alias" "ifunc" "comdat"
-         "attributes" "uselistorder" "uselistorder_bb"
+         "attributes" "uselistorder"
          ;; Linkage types
          "private" "internal" "weak" "weak_odr" "linkonce" "linkonce_odr" "available_externally" "appending" "common" "extern_weak" "external"
          "uninitialized" "implementation" "..."
@@ -120,7 +120,7 @@
    ;; Fast-math flags
    `(,(regexp-opt '("nnan" "ninf" "nsz" "arcp" "contract" "afn" "reassoc" "fast") 'symbols) . font-lock-keyword-face)
    ;; Use-list order directives
-   `(,(regexp-opt '("uselistorder" "uselistorder_bb") 'symbols) . font-lock-keyword-face))
+   `(,(regexp-opt '("uselistorder") 'symbols) . font-lock-keyword-face))
   "Syntax highlighting for LLVM.")
 
 (defun llvm-current-defun-name ()

llvm/utils/vim/syntax/llvm.vim

@@ -199,7 +199,6 @@ syn keyword llvmKeyword
       \ unnamed_addr
       \ unordered
       \ uselistorder
-      \ uselistorder_bb
       \ uwtable
       \ volatile
       \ weak

llvm/utils/vscode/llvm/syntaxes/ll.tmLanguage.yaml

@@ -292,7 +292,6 @@ patterns:
             \\bunnamed_addr\\b|\
             \\bunordered\\b|\
             \\buselistorder\\b|\
-            \\buselistorder_bb\\b|\
             \\buwtable\\b|\
             \\bvolatile\\b|\
             \\bweak\\b|\