[Mlir-commits] [mlir] [MLIR][LLVM] Block address support (PR #134335)

Bruno Cardoso Lopes llvmlistbot at llvm.org
Mon Apr 7 14:46:37 PDT 2025


https://github.com/bcardosolopes updated https://github.com/llvm/llvm-project/pull/134335

>From 6616cdfbb448e54be2f69c9f100e0bba67d2dc6a Mon Sep 17 00:00:00 2001
From: Bruno Cardoso Lopes <bruno.cardoso at gmail.com>
Date: Tue, 1 Apr 2025 16:15:48 -0700
Subject: [PATCH 1/6] [MLIR][LLVM] Block address support

Add support for import and translate.

MLIR does not support using basic block references outside a function (like LLVM does),
This PR does not consider changes to MLIR to that respect. It instead introduces two
new ops: `llvm.blockaddress` and `llvm.blocktag`. Here's an example:

```
llvm.func @ba() -> !llvm.ptr {
  %0 = llvm.blockaddress <function = @ba, tag = <id = 1>> : !llvm.ptr
  llvm.br ^bb1
^bb1:  // pred: ^bb0
  llvm.blocktag <id = 1>
  llvm.return %0 : !llvm.ptr
}
```

Value `%0` hold the address of block tagged as `id = 1` in function `@ba`. Block
tags need to be unique within a function and use of `llvm.blockaddress` requires
a matching tag in a `llvm.blocktag`.
---
 .../mlir/Dialect/LLVMIR/LLVMAttrDefs.td       | 19 +++++
 mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td   | 78 +++++++++++++++++++
 .../mlir/Target/LLVMIR/ModuleTranslation.h    | 35 +++++++++
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp    | 74 ++++++++++++++++++
 .../LLVMIR/LLVMToLLVMIRTranslation.cpp        | 53 +++++++++++++
 mlir/lib/Target/LLVMIR/ModuleImport.cpp       | 28 ++++++-
 mlir/lib/Target/LLVMIR/ModuleTranslation.cpp  | 26 +++++++
 .../LLVMIR/blockaddress-canonicalize.mlir     | 11 +++
 .../test/Dialect/LLVMIR/constant-folding.mlir | 15 ++++
 mlir/test/Dialect/LLVMIR/invalid.mlir         | 22 ++++++
 mlir/test/Dialect/LLVMIR/roundtrip.mlir       | 21 +++++
 .../test/Target/LLVMIR/Import/blockaddress.ll | 32 ++++++++
 .../Target/LLVMIR/Import/import-failure.ll    | 26 -------
 mlir/test/Target/LLVMIR/blockaddress.mlir     | 36 +++++++++
 14 files changed, 447 insertions(+), 29 deletions(-)
 create mode 100644 mlir/test/Dialect/LLVMIR/blockaddress-canonicalize.mlir
 create mode 100644 mlir/test/Target/LLVMIR/Import/blockaddress.ll
 create mode 100644 mlir/test/Target/LLVMIR/blockaddress.mlir

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
index 549a37de2e412..690243525ede4 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
@@ -1224,6 +1224,25 @@ def LLVM_DSOLocalEquivalentAttr : LLVM_Attr<"DSOLocalEquivalent",
   let assemblyFormat = "$sym";
 }
 
+//===----------------------------------------------------------------------===//
+// BlockAddressAttr
+//===----------------------------------------------------------------------===//
+
+def LLVM_BlockTagAttr : LLVM_Attr<"BlockTag", "blocktag"> {
+  let parameters = (ins "uint32_t":$id);
+  let assemblyFormat = "`<` struct(params) `>`";
+}
+
+/// Folded into from LLVM_BlockAddressAttr.
+def LLVM_BlockAddressAttr : LLVM_Attr<"BlockAddress", "blockaddress"> {
+  let description = [{
+    Describes a block address identified by a pair of `$function` and `$tag`.
+  }];
+  let parameters = (ins "FlatSymbolRefAttr":$function,
+                        "BlockTagAttr":$tag);
+  let assemblyFormat = "`<` struct(params) `>`";
+}
+
 //===----------------------------------------------------------------------===//
 // VecTypeHintAttr
 //===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 423cf948b03e1..b107b64e55b46 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -1625,6 +1625,84 @@ def LLVM_DSOLocalEquivalentOp : LLVM_Op<"dso_local_equivalent",
   let hasFolder = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// BlockAddressOp & BlockTagOp
+//===----------------------------------------------------------------------===//
+
+def LLVM_BlockAddressOp : LLVM_Op<"blockaddress",
+    [Pure, ConstantLike, DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
+  let arguments = (ins LLVM_BlockAddressAttr:$block_addr);
+  let results = (outs LLVM_AnyPointer:$res);
+
+  let summary = "Creates a LLVM blockaddress ptr";
+
+  let description = [{
+    Creates an SSA value containing a pointer to a basic block. The block
+    address information (function and block) is given by the `BlockAddressAttr`
+    attribute. This operation assumes an existing `llvm.blocktag` operation
+    identifying an existing MLIR block within a function. Example:
+
+    ```mlir
+    llvm.mlir.global private @g() : !llvm.ptr {
+      %0 = llvm.blockaddress <function = @fn, tag = <id = 0>> : !llvm.ptr
+      llvm.return %0 : !llvm.ptr
+    }
+
+    llvm.func @fn() {
+      llvm.br ^bb1
+    ^bb1:  // pred: ^bb0
+      llvm.blocktag <id = 0>
+      llvm.return
+    }
+    ```
+  }];
+
+  let assemblyFormat = [{
+    $block_addr
+    attr-dict `:` qualified(type($res))
+  }];
+
+  let extraClassDeclaration = [{
+    /// Return the llvm.func operation that is referenced here.
+    LLVMFuncOp getFunction(SymbolTableCollection &symbolTable);
+
+    /// Search for the matching `llvm.blocktag` operation. This is performed
+    /// by walking the function in `block_addr`.
+    BlockTagOp getBlockTagOp();
+  }];
+
+  let hasVerifier = 1;
+  let hasFolder = 1;
+}
+
+def LLVM_BlockTagOp : LLVM_Op<"blocktag"> {
+  let description = [{
+    This operation uses a `tag` to uniquely identify an MLIR block in a
+    function. The same tag is used by `llvm.blockaddress` in order to compute
+    the target address.
+
+    A given function should have at most one `llvm.blocktag` operation with a
+    given `tag`. This operation cannot be used as a terminator but can be
+    placed everywhere else in a block.
+
+    Example:
+
+    ```mlir
+    llvm.func @f() -> !llvm.ptr {
+      %addr = llvm.blockaddress <function = @f, tag = <id = 1>> : !llvm.ptr
+      llvm.br ^bb1
+    ^bb1:
+      llvm.blocktag <id = 1>
+      llvm.return %addr : !llvm.ptr
+    }
+    ```
+  }];
+  let arguments = (ins LLVM_BlockTagAttr:$tag);
+  let assemblyFormat = [{ $tag attr-dict }];
+  // Covered as part of LLVMFuncOp verifier.
+  let hasVerifier = 0;
+}
+
 def LLVM_ComdatSelectorOp : LLVM_Op<"comdat_selector", [Symbol]> {
   let arguments = (ins
     SymbolNameAttr:$sym_name,
diff --git a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
index 99b1b65aeb6a5..a61537463f55f 100644
--- a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
+++ b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
@@ -136,6 +136,29 @@ class ModuleTranslation {
     return callMapping.lookup(op);
   }
 
+  /// Maps a blockaddress operation to its corresponding placeholder LLVM
+  /// value.
+  void mapUnresolvedBlockAddress(BlockAddressOp op, llvm::Value *cst) {
+    auto result = unresolvedBlockAddressMapping.try_emplace(op, cst);
+    (void)result;
+    assert(result.second &&
+           "attempting to map a blockaddress that is already mapped");
+  }
+
+  /// Maps a blockaddress operation to its corresponding placeholder LLVM
+  /// value.
+  void mapBlockTag(BlockAddressAttr attr, BlockTagOp blockTag) {
+    // Attempts to map already mapped block labels are fine given labels are
+    // verified to be unique.
+    blockTagMapping[attr] = blockTag;
+  }
+
+  /// Finds an MLIR block that corresponds to the given MLIR call
+  /// operation.
+  BlockTagOp lookupBlockTag(BlockAddressAttr attr) const {
+    return blockTagMapping.lookup(attr);
+  }
+
   /// Removes the mapping for blocks contained in the region and values defined
   /// in these blocks.
   void forgetMapping(Region &region);
@@ -338,6 +361,8 @@ class ModuleTranslation {
   LogicalResult convertFunctions();
   LogicalResult convertComdats();
 
+  LogicalResult convertUnresolvedBlockAddress();
+
   /// Handle conversion for both globals and global aliases.
   ///
   /// - Create named global variables that correspond to llvm.mlir.global
@@ -433,6 +458,16 @@ class ModuleTranslation {
   /// This map is populated on module entry.
   DenseMap<ComdatSelectorOp, llvm::Comdat *> comdatMapping;
 
+  /// Mapping from llvm.blockaddress operations to their corresponding LLVM
+  /// constant placeholders. After all basic blocks are translated, this
+  /// mapping is used to replace the placeholders with the LLVM block addresses.
+  DenseMap<BlockAddressOp, llvm::Value *> unresolvedBlockAddressMapping;
+
+  /// Mapping from a BlockAddressAttr attribute to a matching BlockTagOp. This
+  /// is used to cache BlockTagOp locations instead of walking a LLVMFuncOp in
+  /// search for those.
+  DenseMap<BlockAddressAttr, BlockTagOp> blockTagMapping;
+
   /// Stack of user-specified state elements, useful when translating operations
   /// with regions.
   SmallVector<std::unique_ptr<StackFrame>> stack;
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 252bdd1425d5e..0deb2900fdfe5 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -2305,6 +2305,28 @@ static LogicalResult verifyComdat(Operation *op,
   return success();
 }
 
+static LogicalResult verifyBlockTags(LLVMFuncOp funcOp) {
+  llvm::DenseSet<BlockTagAttr> blockTags;
+  BlockTagOp badBlockTagOp;
+  if (funcOp
+          .walk([&](BlockTagOp blockTagOp) {
+            if (blockTags.count(blockTagOp.getTag())) {
+              badBlockTagOp = blockTagOp;
+              return WalkResult::interrupt();
+            }
+            blockTags.insert(blockTagOp.getTag());
+            return WalkResult::advance();
+          })
+          .wasInterrupted()) {
+    badBlockTagOp.emitError()
+        << "duplicate block tag '" << badBlockTagOp.getTag().getId()
+        << "' in the same function: ";
+    return failure();
+  }
+
+  return success();
+}
+
 /// Parse common attributes that might show up in the same order in both
 /// GlobalOp and AliasOp.
 template <typename OpType>
@@ -3060,6 +3082,9 @@ LogicalResult LLVMFuncOp::verify() {
     return emitError(diagnosticMessage);
   }
 
+  if (failed(verifyBlockTags(*this)))
+    return failure();
+
   return success();
 }
 
@@ -3815,6 +3840,55 @@ void InlineAsmOp::getEffects(
   }
 }
 
+//===----------------------------------------------------------------------===//
+// BlockAddressOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+BlockAddressOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  Operation *symbol = symbolTable.lookupSymbolIn(parentLLVMModule(*this),
+                                                 getBlockAddr().getFunction());
+  auto function = dyn_cast_or_null<LLVMFuncOp>(symbol);
+
+  if (!function)
+    return emitOpError("must reference a function defined by 'llvm.func'");
+
+  return success();
+}
+
+LLVMFuncOp BlockAddressOp::getFunction(SymbolTableCollection &symbolTable) {
+  return dyn_cast_or_null<LLVMFuncOp>(symbolTable.lookupSymbolIn(
+      parentLLVMModule(*this), getBlockAddr().getFunction()));
+}
+
+BlockTagOp BlockAddressOp::getBlockTagOp() {
+  auto m = (*this)->getParentOfType<ModuleOp>();
+  auto funcOp = cast<LLVMFuncOp>(mlir::SymbolTable::lookupNearestSymbolFrom(
+      m, getBlockAddr().getFunction()));
+
+  BlockTagOp blockTagOp = nullptr;
+  funcOp.walk([&](LLVM::BlockTagOp labelOp) {
+    if (labelOp.getTag() == getBlockAddr().getTag()) {
+      blockTagOp = labelOp;
+      return WalkResult::interrupt();
+    }
+    return WalkResult::advance();
+  });
+  return blockTagOp;
+}
+
+LogicalResult BlockAddressOp::verify() {
+  if (!getBlockTagOp())
+    return emitOpError(
+        "expects an existing block label target in the referenced function");
+
+  return success();
+}
+
+/// Fold a blockaddress operation to a dedicated blockaddress
+/// attribute.
+OpFoldResult BlockAddressOp::fold(FoldAdaptor) { return getBlockAddr(); }
+
 //===----------------------------------------------------------------------===//
 // AssumeOp (intrinsic)
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
index 10b68a333bcbd..738f036bb376a 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
@@ -555,6 +555,59 @@ convertOperationImpl(Operation &opInst, llvm::IRBuilderBase &builder,
     return success();
   }
 
+  // Emit blockaddress. We first need to find the LLVM block referenced by this
+  // operation and then create a LLVM block address for it.
+  if (auto blockAddressOp = dyn_cast<LLVM::BlockAddressOp>(opInst)) {
+    // getBlockTagOp() walks a function to search for block labels. Check
+    // whether it's in cache first.
+    BlockAddressAttr blockAddressAttr = blockAddressOp.getBlockAddr();
+    BlockTagOp blockTagOp = moduleTranslation.lookupBlockTag(blockAddressAttr);
+    if (!blockTagOp) {
+      blockTagOp = blockAddressOp.getBlockTagOp();
+      moduleTranslation.mapBlockTag(blockAddressAttr, blockTagOp);
+    }
+
+    llvm::Value *llvmValue = nullptr;
+    StringRef fnName = blockAddressAttr.getFunction().getValue();
+    if (llvm::BasicBlock *llvmBlock =
+            moduleTranslation.lookupBlock(blockTagOp->getBlock())) {
+      llvm::Function *llvmFn = moduleTranslation.lookupFunction(fnName);
+      llvmValue = llvm::BlockAddress::get(llvmFn, llvmBlock);
+    } else {
+      // The matching LLVM block is not yet emitted, a placeholder is created
+      // in its place. When the LLVM block is emitted later in translation,
+      // the llvmValue is replaced with the actual llvm::BlockAddress.
+      // A GlobalVariable is chosen as placeholder because in general LLVM
+      // constants are uniqued and are not proper for RAUW, since that could
+      // harm unrelated uses of the constant.
+      llvmValue = new llvm::GlobalVariable(
+          *moduleTranslation.getLLVMModule(),
+          llvm::PointerType::getUnqual(moduleTranslation.getLLVMContext()),
+          /*isConstant=*/true, llvm::GlobalValue::LinkageTypes::ExternalLinkage,
+          /*Initializer=*/nullptr,
+          Twine("__mlir_block_address_")
+              .concat(Twine(fnName))
+              .concat(Twine((uint64_t)blockAddressOp.getOperation())));
+      moduleTranslation.mapUnresolvedBlockAddress(blockAddressOp, llvmValue);
+    }
+
+    moduleTranslation.mapValue(blockAddressOp.getResult(), llvmValue);
+    return success();
+  }
+
+  // Emit block label. If this label is seen before BlockAddressOp is
+  // translated, go ahead and already map it.
+  if (auto blockTagOp = dyn_cast<LLVM::BlockTagOp>(opInst)) {
+    auto funcOp = blockTagOp->getParentOfType<LLVMFuncOp>();
+    BlockAddressAttr blockAddressAttr = BlockAddressAttr::get(
+        &moduleTranslation.getContext(),
+        FlatSymbolRefAttr::get(&moduleTranslation.getContext(),
+                               funcOp.getName()),
+        blockTagOp.getTag());
+    moduleTranslation.mapBlockTag(blockAddressAttr, blockTagOp);
+    return success();
+  }
+
   return failure();
 }
 
diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index 5f047a59a9828..b71330445d047 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -1381,9 +1381,19 @@ FailureOr<Value> ModuleImport::convertConstant(llvm::Constant *constant) {
     return builder.create<LLVM::ZeroOp>(loc, targetExtType).getRes();
   }
 
+  if (auto *blockAddr = dyn_cast<llvm::BlockAddress>(constant))
+    return builder
+        .create<BlockAddressOp>(
+            loc, convertType(blockAddr->getType()),
+            BlockAddressAttr::get(
+                context,
+                FlatSymbolRefAttr::get(context,
+                                       blockAddr->getFunction()->getName()),
+                BlockTagAttr::get(context,
+                                  blockAddr->getBasicBlock()->getNumber())))
+        .getRes();
+
   StringRef error = "";
-  if (isa<llvm::BlockAddress>(constant))
-    error = " since blockaddress(...) is unsupported";
 
   if (isa<llvm::ConstantPtrAuth>(constant))
     error = " since ptrauth(...) is unsupported";
@@ -2448,8 +2458,13 @@ LogicalResult ModuleImport::processFunction(llvm::Function *func) {
   SmallVector<llvm::BasicBlock *> reachableBasicBlocks;
   for (llvm::BasicBlock &basicBlock : *func) {
     // Skip unreachable blocks.
-    if (!reachable.contains(&basicBlock))
+    if (!reachable.contains(&basicBlock)) {
+      if (basicBlock.hasAddressTaken())
+        emitWarning(funcOp.getLoc())
+            << "unreachable block '" << basicBlock.getName()
+            << "' with address taken";
       continue;
+    }
     Region &body = funcOp.getBody();
     Block *block = builder.createBlock(&body, body.end());
     mapBlock(&basicBlock, block);
@@ -2606,6 +2621,13 @@ LogicalResult ModuleImport::processBasicBlock(llvm::BasicBlock *bb,
       }
     }
   }
+
+  if (bb->hasAddressTaken()) {
+    OpBuilder::InsertionGuard guard(builder);
+    builder.setInsertionPointToStart(block);
+    builder.create<BlockTagOp>(block->getParentOp()->getLoc(),
+                               BlockTagAttr::get(context, bb->getNumber()));
+  }
   return success();
 }
 
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index d30cb8a7d7974..ee7dc3a5231f4 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -1824,6 +1824,27 @@ LogicalResult ModuleTranslation::convertComdats() {
   return success();
 }
 
+LogicalResult ModuleTranslation::convertUnresolvedBlockAddress() {
+  for (auto &[blockAddressOp, llvmCst] : unresolvedBlockAddressMapping) {
+    BlockAddressAttr blockAddressAttr = blockAddressOp.getBlockAddr();
+    BlockTagOp blockTagOp = lookupBlockTag(blockAddressAttr);
+    assert(blockTagOp && "expected all block tags to be already seen");
+
+    llvm::BasicBlock *llvmBlock = lookupBlock(blockTagOp->getBlock());
+    assert(llvmBlock && "expected LLVM blocks to be already translated");
+
+    // Update mapping with new block address constant.
+    auto *llvmBlockAddr = llvm::BlockAddress::get(
+        lookupFunction(blockAddressAttr.getFunction().getValue()), llvmBlock);
+    llvmCst->replaceAllUsesWith(llvmBlockAddr);
+    mapValue(blockAddressOp.getResult(), llvmBlockAddr);
+    assert(llvmCst->use_empty() && "expected all uses to be replaced");
+    cast<llvm::GlobalVariable>(llvmCst)->eraseFromParent();
+  }
+  unresolvedBlockAddressMapping.clear();
+  return success();
+}
+
 void ModuleTranslation::setAccessGroupsMetadata(AccessGroupOpInterface op,
                                                 llvm::Instruction *inst) {
   if (llvm::MDNode *node = loopAnnotationTranslation->getAccessGroups(op))
@@ -2236,6 +2257,11 @@ mlir::translateModuleToLLVMIR(Operation *module, llvm::LLVMContext &llvmContext,
   if (failed(translator.convertFunctions()))
     return nullptr;
 
+  // Now that all MLIR blocks are resolved into LLVM ones, patch block address
+  // constants to point to the correct blocks.
+  if (failed(translator.convertUnresolvedBlockAddress()))
+    return nullptr;
+
   // Once we've finished constructing elements in the module, we should convert
   // it to use the debug info format desired by LLVM.
   // See https://llvm.org/docs/RemoveDIsDebugInfo.html
diff --git a/mlir/test/Dialect/LLVMIR/blockaddress-canonicalize.mlir b/mlir/test/Dialect/LLVMIR/blockaddress-canonicalize.mlir
new file mode 100644
index 0000000000000..0f18cc6d8cd3e
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/blockaddress-canonicalize.mlir
@@ -0,0 +1,11 @@
+// RUN: mlir-opt %s -canonicalize | FileCheck %s
+
+// CHECK-LABEL: llvm.func @ba()
+llvm.func @ba() -> !llvm.ptr {
+  %0 = llvm.blockaddress <function = @ba, tag = <id = 1>> : !llvm.ptr
+  llvm.br ^bb1
+^bb1:
+  // CHECK: llvm.blocktag <id = 1>
+  llvm.blocktag <id = 1>
+  llvm.return %0 : !llvm.ptr
+}
diff --git a/mlir/test/Dialect/LLVMIR/constant-folding.mlir b/mlir/test/Dialect/LLVMIR/constant-folding.mlir
index 99f657f0aefec..0616f19b8fddb 100644
--- a/mlir/test/Dialect/LLVMIR/constant-folding.mlir
+++ b/mlir/test/Dialect/LLVMIR/constant-folding.mlir
@@ -196,3 +196,18 @@ llvm.func @dso_local_equivalent_select(%arg: i1) -> !llvm.ptr {
 }
 
 llvm.func @yay()
+
+// -----
+
+// CHECK-LABEL: llvm.func @blockaddress_select
+llvm.func @blockaddress_select(%arg: i1) -> !llvm.ptr {
+  // CHECK-NEXT: %[[ADDR:.+]] = llvm.blockaddress <function = @blockaddress_select, tag = <id = 1>>
+  %0 = llvm.blockaddress <function = @blockaddress_select, tag = <id = 1>> : !llvm.ptr
+  %1 = llvm.blockaddress <function = @blockaddress_select, tag = <id = 1>> : !llvm.ptr
+  %2 = arith.select %arg, %0, %1 : !llvm.ptr
+  // CHECK-NEXT: llvm.br ^bb1
+  llvm.br ^bb1
+^bb1:
+  llvm.blocktag <id = 1>
+  llvm.return %1 : !llvm.ptr
+}
diff --git a/mlir/test/Dialect/LLVMIR/invalid.mlir b/mlir/test/Dialect/LLVMIR/invalid.mlir
index fb9631d99b91a..e70e3185af236 100644
--- a/mlir/test/Dialect/LLVMIR/invalid.mlir
+++ b/mlir/test/Dialect/LLVMIR/invalid.mlir
@@ -1780,3 +1780,25 @@ module {
   // expected-error at +1 {{failed to parse ModuleFlagAttr parameter 'value' which is to be a `uint32_t`}}
   llvm.module_flags [#llvm.mlir.module_flag<error, "wchar_size", "yolo">]
 }
+
+// -----
+
+llvm.func @t0() -> !llvm.ptr {
+  %0 = llvm.blockaddress <function = @t0, tag = <id = 1>> : !llvm.ptr
+  llvm.blocktag <id = 1>
+  llvm.br ^bb1
+^bb1:
+  // expected-error at +1 {{duplicate block tag '1' in the same function}}
+  llvm.blocktag <id = 1>
+  llvm.return %0 : !llvm.ptr
+}
+
+// -----
+
+llvm.func @t1() -> !llvm.ptr {
+  // expected-error at +1 {{expects an existing block label target in the referenced function}}
+  %0 = llvm.blockaddress <function = @t1, tag = <id = 1>> : !llvm.ptr
+  llvm.br ^bb1
+^bb1:
+  llvm.return %0 : !llvm.ptr
+}
diff --git a/mlir/test/Dialect/LLVMIR/roundtrip.mlir b/mlir/test/Dialect/LLVMIR/roundtrip.mlir
index d0aa65d14a176..88460fe374d87 100644
--- a/mlir/test/Dialect/LLVMIR/roundtrip.mlir
+++ b/mlir/test/Dialect/LLVMIR/roundtrip.mlir
@@ -1002,3 +1002,24 @@ llvm.func @intrinsic_call_arg_attrs_bundles(%arg0: i32) -> i32 {
   %0 = llvm.call_intrinsic "llvm.riscv.sha256sig0"(%arg0) ["adazdazd"()] {constant} : (i32 {llvm.signext}) -> (i32)
   llvm.return %0 : i32
 }
+
+llvm.mlir.global private @blockaddr_global() {addr_space = 0 : i32, dso_local} : !llvm.ptr {
+  %0 = llvm.blockaddress <function = @blockaddr_fn, tag = <id = 0>> : !llvm.ptr
+  llvm.return %0 : !llvm.ptr
+}
+
+// CHECK: llvm.mlir.global private @blockaddr_global() {{.*}}
+// CHECK-NEXT:   %{{.*}} = llvm.blockaddress <function = @blockaddr_fn, tag = <id = 0>> : !llvm.ptr
+// CHECK-NEXT:   llvm.return %{{.*}} : !llvm.ptr
+
+llvm.func @blockaddr_fn() {
+  llvm.br ^bb1
+^bb1:
+  llvm.blocktag <id = 0>
+  llvm.return
+}
+
+// CHECK-LABEL: llvm.func @blockaddr_fn
+// CHECK-NEXT:  llvm.br ^bb1
+// CHECK-NEXT:^bb1:
+// CHECK-NEXT:  llvm.blocktag <id = 0>
diff --git a/mlir/test/Target/LLVMIR/Import/blockaddress.ll b/mlir/test/Target/LLVMIR/Import/blockaddress.ll
new file mode 100644
index 0000000000000..fb6ef1b6c3a2b
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/Import/blockaddress.ll
@@ -0,0 +1,32 @@
+; RUN: mlir-translate -import-llvm -split-input-file %s | FileCheck %s
+
+ at g = private global ptr blockaddress(@fn, %bb1)
+define void @fn() {
+  br label %bb1
+bb1:
+  ret void
+}
+
+; CHECK:  llvm.mlir.global private @g()
+; CHECK:     %[[ADDR:.*]] = llvm.blockaddress <function = @fn, tag = <id = [[ID:.*]]>> : !llvm.ptr
+; CHECK:     llvm.return %[[ADDR]] : !llvm.ptr
+
+; CHECK:   llvm.func @fn() {
+; CHECK:     llvm.br ^[[RET_BB:.*]]
+; CHECK:   ^[[RET_BB]]:
+; CHECK:     llvm.blocktag <id = [[ID]]>
+; CHECK:     llvm.return
+; CHECK:   }
+
+; // -----
+
+; CHECK-LABEL: blockaddr0
+define ptr @blockaddr0() {
+  br label %bb1
+  ; CHECK: %[[BLOCKADDR:.*]] = llvm.blockaddress <function = @blockaddr0, tag = <id = [[BLOCK_ID:.*]]>> : !llvm.ptr
+bb1:
+  ; CHECK: [[BLOCKADDR]]:
+  ; CHECK: llvm.blocktag <id = [[BLOCK_ID]]>
+  ; CHECK-NEXT: llvm.return %[[BLOCKADDR]] : !llvm.ptr
+  ret ptr blockaddress(@blockaddr0, %bb1)
+}
diff --git a/mlir/test/Target/LLVMIR/Import/import-failure.ll b/mlir/test/Target/LLVMIR/Import/import-failure.ll
index d3ea3a510d7f8..b7ca2ad36c46c 100644
--- a/mlir/test/Target/LLVMIR/Import/import-failure.ll
+++ b/mlir/test/Target/LLVMIR/Import/import-failure.ll
@@ -12,32 +12,6 @@ bb2:
 
 ; // -----
 
-; CHECK:      <unknown>
-; CHECK-SAME: unhandled constant: ptr blockaddress(@unhandled_constant, %bb1) since blockaddress(...) is unsupported
-; CHECK:      <unknown>
-; CHECK-SAME: error: unhandled instruction: ret ptr blockaddress(@unhandled_constant, %bb1)
-define ptr @unhandled_constant() {
-  br label %bb1
-bb1:
-  ret ptr blockaddress(@unhandled_constant, %bb1)
-}
-
-; // -----
-
-; CHECK:      <unknown>
-; CHECK-SAME: unhandled constant: ptr blockaddress(@unhandled_global, %bb1) since blockaddress(...) is unsupported
-; CHECK:      <unknown>
-; CHECK-SAME: error: unhandled global variable: @private = private global ptr blockaddress(@unhandled_global, %bb1)
- at private = private global ptr blockaddress(@unhandled_global, %bb1)
-
-define void @unhandled_global() {
-  br label %bb1
-bb1:
-  ret void
-}
-
-; // -----
-
 ; Check that debug intrinsics with an unsupported argument are dropped.
 
 declare void @llvm.dbg.value(metadata, metadata, metadata)
diff --git a/mlir/test/Target/LLVMIR/blockaddress.mlir b/mlir/test/Target/LLVMIR/blockaddress.mlir
new file mode 100644
index 0000000000000..bf7dfef637b63
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/blockaddress.mlir
@@ -0,0 +1,36 @@
+// RUN: mlir-translate -mlir-to-llvmir %s -split-input-file | FileCheck %s
+
+llvm.mlir.global private @g() {addr_space = 0 : i32, dso_local} : !llvm.ptr {
+  %0 = llvm.blockaddress <function = @fn, tag = <id = 0>> : !llvm.ptr
+  llvm.return %0 : !llvm.ptr
+}
+
+llvm.func @fn() {
+  llvm.br ^bb1
+^bb1:
+  llvm.blocktag <id = 0>
+  llvm.return
+}
+
+// CHECK: @g = private global ptr blockaddress(@fn, %1)
+// CHECK: define void @fn() {
+// CHECK:   br label %1
+// CHECK: 1:
+// CHECK:   ret void
+// CHECK: }
+
+// -----
+
+llvm.func @blockaddr0() -> !llvm.ptr {
+  %0 = llvm.blockaddress <function = @blockaddr0, tag = <id = 1>> : !llvm.ptr
+  llvm.br ^bb1
+^bb1:
+  llvm.blocktag <id = 1>
+  llvm.return %0 : !llvm.ptr
+}
+
+// CHECK: define ptr @blockaddr0() {
+// CHECK:   br label %1
+// CHECK: 1:
+// CHECK:   ret ptr blockaddress(@blockaddr0, %1)
+// CHECK: }

>From a90e0e65a4f9975b88668517eb506ab17d2cd4d1 Mon Sep 17 00:00:00 2001
From: Bruno Cardoso Lopes <bruno.cardoso at gmail.com>
Date: Fri, 4 Apr 2025 10:54:19 -0700
Subject: [PATCH 2/6] address some of the review points

---
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp |  5 ++---
 mlir/lib/Target/LLVMIR/ModuleImport.cpp    | 17 ++++++++---------
 2 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 0deb2900fdfe5..e8353425c4213 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -2310,7 +2310,7 @@ static LogicalResult verifyBlockTags(LLVMFuncOp funcOp) {
   BlockTagOp badBlockTagOp;
   if (funcOp
           .walk([&](BlockTagOp blockTagOp) {
-            if (blockTags.count(blockTagOp.getTag())) {
+            if (blockTags.contains(blockTagOp.getTag())) {
               badBlockTagOp = blockTagOp;
               return WalkResult::interrupt();
             }
@@ -3862,9 +3862,8 @@ LLVMFuncOp BlockAddressOp::getFunction(SymbolTableCollection &symbolTable) {
 }
 
 BlockTagOp BlockAddressOp::getBlockTagOp() {
-  auto m = (*this)->getParentOfType<ModuleOp>();
   auto funcOp = cast<LLVMFuncOp>(mlir::SymbolTable::lookupNearestSymbolFrom(
-      m, getBlockAddr().getFunction()));
+      parentLLVMModule(*this), getBlockAddr().getFunction()));
 
   BlockTagOp blockTagOp = nullptr;
   funcOp.walk([&](LLVM::BlockTagOp labelOp) {
diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index b71330445d047..0ea0b1fd197c1 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -1381,17 +1381,16 @@ FailureOr<Value> ModuleImport::convertConstant(llvm::Constant *constant) {
     return builder.create<LLVM::ZeroOp>(loc, targetExtType).getRes();
   }
 
-  if (auto *blockAddr = dyn_cast<llvm::BlockAddress>(constant))
+  if (auto *blockAddr = dyn_cast<llvm::BlockAddress>(constant)) {
+    auto fnSym =
+        FlatSymbolRefAttr::get(context, blockAddr->getFunction()->getName());
+    auto blockTag =
+        BlockTagAttr::get(context, blockAddr->getBasicBlock()->getNumber());
     return builder
-        .create<BlockAddressOp>(
-            loc, convertType(blockAddr->getType()),
-            BlockAddressAttr::get(
-                context,
-                FlatSymbolRefAttr::get(context,
-                                       blockAddr->getFunction()->getName()),
-                BlockTagAttr::get(context,
-                                  blockAddr->getBasicBlock()->getNumber())))
+        .create<BlockAddressOp>(loc, convertType(blockAddr->getType()),
+                                BlockAddressAttr::get(context, fnSym, blockTag))
         .getRes();
+  }
 
   StringRef error = "";
 

>From 85d2918593eb9b9a46edf05f421b048ff152b403 Mon Sep 17 00:00:00 2001
From: Bruno Cardoso Lopes <bruno.cardoso at gmail.com>
Date: Fri, 4 Apr 2025 11:03:31 -0700
Subject: [PATCH 3/6] Emit error when taking the address of unrecheable blocks

---
 mlir/lib/Target/LLVMIR/ModuleImport.cpp          |  2 +-
 mlir/test/Target/LLVMIR/Import/import-failure.ll | 10 ++++++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index 0ea0b1fd197c1..8c4f6b6eb9040 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -2459,7 +2459,7 @@ LogicalResult ModuleImport::processFunction(llvm::Function *func) {
     // Skip unreachable blocks.
     if (!reachable.contains(&basicBlock)) {
       if (basicBlock.hasAddressTaken())
-        emitWarning(funcOp.getLoc())
+        emitError(funcOp.getLoc())
             << "unreachable block '" << basicBlock.getName()
             << "' with address taken";
       continue;
diff --git a/mlir/test/Target/LLVMIR/Import/import-failure.ll b/mlir/test/Target/LLVMIR/Import/import-failure.ll
index b7ca2ad36c46c..4fbf187659a7b 100644
--- a/mlir/test/Target/LLVMIR/Import/import-failure.ll
+++ b/mlir/test/Target/LLVMIR/Import/import-failure.ll
@@ -350,3 +350,13 @@ bb2:
 declare i32 @g()
 
 declare i32 @__gxx_personality_v0(...)
+
+; // -----
+
+ at g = private global ptr blockaddress(@fn, %bb1)
+define void @fn() {
+  ret void
+; CHECK: unreachable block 'bb1' with address taken
+bb1:
+  ret void
+}

>From 73529928e64a98a8bbc2c830a17097516eb97262 Mon Sep 17 00:00:00 2001
From: Bruno Cardoso Lopes <bruno.cardoso at gmail.com>
Date: Fri, 4 Apr 2025 11:22:00 -0700
Subject: [PATCH 4/6] Do not inline functions containing llvm.blocktag

---
 .../LLVMIR/Transforms/InlinerInterfaceImpl.cpp   |  5 +++--
 mlir/test/Dialect/LLVMIR/inlining.mlir           | 16 ++++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
index 1edf7fd070b27..bcc0fd4aac8f8 100644
--- a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
@@ -731,8 +731,9 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
   }
 
   bool isLegalToInline(Operation *op, Region *, bool, IRMapping &) const final {
-    // The inliner cannot handle variadic function arguments.
-    return !isa<LLVM::VaStartOp>(op);
+    // The inliner cannot handle variadic function arguments nor blocktag
+    // operations.
+    return !(isa<LLVM::VaStartOp>(op) || isa<LLVM::BlockTagOp>(op));
   }
 
   /// Handle the given inlined return by replacing it with a branch. This
diff --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir
index 136d0f85d509a..4e37b0d4be6a0 100644
--- a/mlir/test/Dialect/LLVMIR/inlining.mlir
+++ b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -692,3 +692,19 @@ llvm.func @caller(%x : i32) -> i32 {
   %z = llvm.call @unreachable_func(%x) : (i32) -> (i32)
   llvm.return %z : i32
 }
+
+// -----
+// Check that @func is not inlined because of llvm.blocktag
+
+func.func @func(%arg0 : i32) -> i32  {
+  llvm.blocktag <id = 1>
+  llvm.return %arg0 : i32
+}
+// CHECK-LABEL: @llvm_ret
+// CHECK-NOT: llvm.blocktag
+// CHECK: %[[R:.*]] = call
+// CHECK: return %[[R]]
+func.func @llvm_ret(%arg0 : i32) -> i32 {
+  %res = call @func(%arg0) : (i32) -> (i32)
+  return %res : i32
+}

>From f85e6d529f87733c434d04a299d246ddeeb92bd0 Mon Sep 17 00:00:00 2001
From: Bruno Cardoso Lopes <bruno.cardoso at gmail.com>
Date: Fri, 4 Apr 2025 12:35:23 -0700
Subject: [PATCH 5/6] Add a region-simplify test

---
 .../LLVMIR/blockaddress-canonicalize.mlir     | 30 ++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

diff --git a/mlir/test/Dialect/LLVMIR/blockaddress-canonicalize.mlir b/mlir/test/Dialect/LLVMIR/blockaddress-canonicalize.mlir
index 0f18cc6d8cd3e..9b2c35ca86954 100644
--- a/mlir/test/Dialect/LLVMIR/blockaddress-canonicalize.mlir
+++ b/mlir/test/Dialect/LLVMIR/blockaddress-canonicalize.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -canonicalize | FileCheck %s
+// RUN: mlir-opt %s -pass-pipeline='builtin.module(llvm.func(canonicalize{region-simplify=aggressive}))' -split-input-file | FileCheck %s
 
 // CHECK-LABEL: llvm.func @ba()
 llvm.func @ba() -> !llvm.ptr {
@@ -9,3 +9,31 @@ llvm.func @ba() -> !llvm.ptr {
   llvm.blocktag <id = 1>
   llvm.return %0 : !llvm.ptr
 }
+
+// -----
+
+
+llvm.mlir.global private @g() {addr_space = 0 : i32, dso_local} : !llvm.ptr {
+  %0 = llvm.blockaddress <function = @fn, tag = <id = 0>> : !llvm.ptr
+  llvm.return %0 : !llvm.ptr
+}
+
+llvm.mlir.global private @h() {addr_space = 0 : i32, dso_local} : !llvm.ptr {
+  %0 = llvm.blockaddress <function = @fn, tag = <id = 1>> : !llvm.ptr
+  llvm.return %0 : !llvm.ptr
+}
+
+// CHECK-LABEL: llvm.func @fn
+llvm.func @fn(%cond : i1, %arg0 : i32, %arg1 : i32) -> i32 {
+  llvm.cond_br %cond, ^bb1, ^bb2
+^bb1:
+  // CHECK: llvm.blocktag <id = 0>
+  // CHECK: llvm.return
+  llvm.blocktag <id = 0>
+  llvm.return %arg0 : i32
+^bb2:
+  // CHECK: llvm.blocktag <id = 1>
+  // CHECK: llvm.return
+  llvm.blocktag <id = 1>
+  llvm.return %arg1 : i32
+}

>From a523b7d2cbf436a7fe432bf1d6515338800ea217 Mon Sep 17 00:00:00 2001
From: Bruno Cardoso Lopes <bruno.cardoso at gmail.com>
Date: Mon, 7 Apr 2025 14:42:10 -0700
Subject: [PATCH 6/6] Address last round of comments

---
 mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h      | 4 ++--
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp               | 4 +++-
 .../Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp   | 5 +++--
 mlir/lib/Target/LLVMIR/ModuleImport.cpp                  | 6 +++---
 mlir/test/Dialect/LLVMIR/blockaddress-canonicalize.mlir  | 9 +++++++++
 mlir/test/Dialect/LLVMIR/inlining.mlir                   | 7 ++++---
 mlir/test/Target/LLVMIR/blockaddress.mlir                | 8 ++++----
 7 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
index a61537463f55f..30c190e50a4f7 100644
--- a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
+++ b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
@@ -148,8 +148,8 @@ class ModuleTranslation {
   /// Maps a blockaddress operation to its corresponding placeholder LLVM
   /// value.
   void mapBlockTag(BlockAddressAttr attr, BlockTagOp blockTag) {
-    // Attempts to map already mapped block labels are fine given labels are
-    // verified to be unique.
+    // Attempts to map already mapped block labels which is fine if the given
+    // labels are verified to be unique.
     blockTagMapping[attr] = blockTag;
   }
 
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index e8353425c4213..42a3839e8c638 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -3862,8 +3862,10 @@ LLVMFuncOp BlockAddressOp::getFunction(SymbolTableCollection &symbolTable) {
 }
 
 BlockTagOp BlockAddressOp::getBlockTagOp() {
-  auto funcOp = cast<LLVMFuncOp>(mlir::SymbolTable::lookupNearestSymbolFrom(
+  auto funcOp = dyn_cast<LLVMFuncOp>(mlir::SymbolTable::lookupNearestSymbolFrom(
       parentLLVMModule(*this), getBlockAddr().getFunction()));
+  if (!funcOp)
+    return nullptr;
 
   BlockTagOp blockTagOp = nullptr;
   funcOp.walk([&](LLVM::BlockTagOp labelOp) {
diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
index bcc0fd4aac8f8..7f3afffc9645e 100644
--- a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
@@ -731,8 +731,9 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
   }
 
   bool isLegalToInline(Operation *op, Region *, bool, IRMapping &) const final {
-    // The inliner cannot handle variadic function arguments nor blocktag
-    // operations.
+    // The inliner cannot handle variadic function arguments and blocktag
+    // operations prevent inlining since they the blockaddress operations
+    // reference them via the callee symbol.
     return !(isa<LLVM::VaStartOp>(op) || isa<LLVM::BlockTagOp>(op));
   }
 
diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index 8c4f6b6eb9040..2859abdb41772 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -2459,9 +2459,9 @@ LogicalResult ModuleImport::processFunction(llvm::Function *func) {
     // Skip unreachable blocks.
     if (!reachable.contains(&basicBlock)) {
       if (basicBlock.hasAddressTaken())
-        emitError(funcOp.getLoc())
-            << "unreachable block '" << basicBlock.getName()
-            << "' with address taken";
+        return emitError(funcOp.getLoc())
+               << "unreachable block '" << basicBlock.getName()
+               << "' with address taken";
       continue;
     }
     Region &body = funcOp.getBody();
diff --git a/mlir/test/Dialect/LLVMIR/blockaddress-canonicalize.mlir b/mlir/test/Dialect/LLVMIR/blockaddress-canonicalize.mlir
index 9b2c35ca86954..11dd6f0d97f78 100644
--- a/mlir/test/Dialect/LLVMIR/blockaddress-canonicalize.mlir
+++ b/mlir/test/Dialect/LLVMIR/blockaddress-canonicalize.mlir
@@ -1,5 +1,10 @@
 // RUN: mlir-opt %s -pass-pipeline='builtin.module(llvm.func(canonicalize{region-simplify=aggressive}))' -split-input-file | FileCheck %s
 
+llvm.mlir.global private @x() {addr_space = 0 : i32, dso_local} : !llvm.ptr {
+  %0 = llvm.blockaddress <function = @ba, tag = <id = 2>> : !llvm.ptr
+  llvm.return %0 : !llvm.ptr
+}
+
 // CHECK-LABEL: llvm.func @ba()
 llvm.func @ba() -> !llvm.ptr {
   %0 = llvm.blockaddress <function = @ba, tag = <id = 1>> : !llvm.ptr
@@ -7,6 +12,10 @@ llvm.func @ba() -> !llvm.ptr {
 ^bb1:
   // CHECK: llvm.blocktag <id = 1>
   llvm.blocktag <id = 1>
+  llvm.br ^bb2
+^bb2:
+  // CHECK: llvm.blocktag <id = 2>
+  llvm.blocktag <id = 2>
   llvm.return %0 : !llvm.ptr
 }
 
diff --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir
index 4e37b0d4be6a0..551e0c97912d0 100644
--- a/mlir/test/Dialect/LLVMIR/inlining.mlir
+++ b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -700,11 +700,12 @@ func.func @func(%arg0 : i32) -> i32  {
   llvm.blocktag <id = 1>
   llvm.return %arg0 : i32
 }
+
 // CHECK-LABEL: @llvm_ret
-// CHECK-NOT: llvm.blocktag
-// CHECK: %[[R:.*]] = call
-// CHECK: return %[[R]]
 func.func @llvm_ret(%arg0 : i32) -> i32 {
+  // CHECK-NOT: llvm.blocktag
+  // CHECK: %[[R:.*]] = call
   %res = call @func(%arg0) : (i32) -> (i32)
+  // CHECK: return %[[R]]
   return %res : i32
 }
diff --git a/mlir/test/Target/LLVMIR/blockaddress.mlir b/mlir/test/Target/LLVMIR/blockaddress.mlir
index bf7dfef637b63..fb3d853531122 100644
--- a/mlir/test/Target/LLVMIR/blockaddress.mlir
+++ b/mlir/test/Target/LLVMIR/blockaddress.mlir
@@ -14,8 +14,8 @@ llvm.func @fn() {
 
 // CHECK: @g = private global ptr blockaddress(@fn, %1)
 // CHECK: define void @fn() {
-// CHECK:   br label %1
-// CHECK: 1:
+// CHECK:   br label %[[RET:.*]]
+// CHECK: [[RET]]:
 // CHECK:   ret void
 // CHECK: }
 
@@ -30,7 +30,7 @@ llvm.func @blockaddr0() -> !llvm.ptr {
 }
 
 // CHECK: define ptr @blockaddr0() {
-// CHECK:   br label %1
-// CHECK: 1:
+// CHECK:   br label %[[RET:.*]]
+// CHECK: [[RET]]:
 // CHECK:   ret ptr blockaddress(@blockaddr0, %1)
 // CHECK: }



More information about the Mlir-commits mailing list