[Mlir-commits] [mlir] [MLIR][LLVM] Fix blockaddress mapping to LLVM blocks (PR #139814)
Bruno Cardoso Lopes
llvmlistbot at llvm.org
Tue May 13 18:08:54 PDT 2025
https://github.com/bcardosolopes updated https://github.com/llvm/llvm-project/pull/139814
>From 67489aa5128b3bc6679479a58170e70196003e8f Mon Sep 17 00:00:00 2001
From: Bruno Cardoso Lopes <bruno.cardoso at gmail.com>
Date: Tue, 13 May 2025 16:32:28 -0700
Subject: [PATCH 1/2] [MLIR][LLVM] Transation: fix bug in blockaddress mappings
of blocks
Stop using mapped llvm blocks because they have a per function lifetime and
conversion for block addresses is potentially delayed until all functions are
emitted. Keep a separate map for block tags to llvm blocks.
---
.../mlir/Target/LLVMIR/ModuleTranslation.h | 24 +++++++--------
.../LLVMIR/LLVMToLLVMIRTranslation.cpp | 15 ++++------
mlir/lib/Target/LLVMIR/ModuleTranslation.cpp | 6 +---
mlir/test/Target/LLVMIR/blockaddress.mlir | 29 +++++++++++++++++++
4 files changed, 45 insertions(+), 29 deletions(-)
diff --git a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
index 30c190e50a4f7..f0ad1dedb8760 100644
--- a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
+++ b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
@@ -145,18 +145,15 @@ class ModuleTranslation {
"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 which is fine if the given
- // labels are verified to be unique.
- blockTagMapping[attr] = blockTag;
+ /// Maps a BlockAddressAttr to its corresponding LLVM basic block.
+ void mapBlockAddress(BlockAddressAttr attr, llvm::BasicBlock *block) {
+ assert(!blockAddressToLLVMMapping.count(blockAddressToLLVMMapping));
+ blockAddressToLLVMMapping[attr] = block;
}
- /// Finds an MLIR block that corresponds to the given MLIR call
- /// operation.
- BlockTagOp lookupBlockTag(BlockAddressAttr attr) const {
- return blockTagMapping.lookup(attr);
+ /// Finds the LLVM basic block that corresponds to the given BlockAddressAttr.
+ llvm::BasicBlock *lookupBlockAddress(BlockAddressAttr attr) const {
+ return blockAddressToLLVMMapping.lookup(attr);
}
/// Removes the mapping for blocks contained in the region and values defined
@@ -463,10 +460,9 @@ class ModuleTranslation {
/// 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;
+ /// Mapping from a BlockAddressAttr attribute to it's matching LLVM basic
+ /// block.
+ DenseMap<BlockAddressAttr, llvm::BasicBlock *> blockAddressToLLVMMapping;
/// Stack of user-specified state elements, useful when translating operations
/// with regions.
diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
index 4ea313019f34d..9470b54c9f3aa 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
@@ -690,19 +690,13 @@ convertOperationImpl(Operation &opInst, llvm::IRBuilderBase &builder,
// 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::BasicBlock *llvmBlock =
+ moduleTranslation.lookupBlockAddress(blockAddressAttr);
llvm::Value *llvmValue = nullptr;
StringRef fnName = blockAddressAttr.getFunction().getValue();
- if (llvm::BasicBlock *llvmBlock =
- moduleTranslation.lookupBlock(blockTagOp->getBlock())) {
+ if (llvmBlock) {
llvm::Function *llvmFn = moduleTranslation.lookupFunction(fnName);
llvmValue = llvm::BlockAddress::get(llvmFn, llvmBlock);
} else {
@@ -736,7 +730,8 @@ convertOperationImpl(Operation &opInst, llvm::IRBuilderBase &builder,
FlatSymbolRefAttr::get(&moduleTranslation.getContext(),
funcOp.getName()),
blockTagOp.getTag());
- moduleTranslation.mapBlockTag(blockAddressAttr, blockTagOp);
+ moduleTranslation.mapBlockAddress(blockAddressAttr,
+ builder.GetInsertBlock());
return success();
}
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index 1168b9f339904..95b8ee0331c55 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -1843,17 +1843,13 @@ LogicalResult ModuleTranslation::convertComdats() {
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());
+ llvm::BasicBlock *llvmBlock = lookupBlockAddress(blockAddressAttr);
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();
}
diff --git a/mlir/test/Target/LLVMIR/blockaddress.mlir b/mlir/test/Target/LLVMIR/blockaddress.mlir
index fb3d853531122..4473f91c4bdb5 100644
--- a/mlir/test/Target/LLVMIR/blockaddress.mlir
+++ b/mlir/test/Target/LLVMIR/blockaddress.mlir
@@ -34,3 +34,32 @@ llvm.func @blockaddr0() -> !llvm.ptr {
// CHECK: [[RET]]:
// CHECK: ret ptr blockaddress(@blockaddr0, %1)
// CHECK: }
+
+// -----
+
+llvm.mlir.global private @h() {addr_space = 0 : i32, dso_local} : !llvm.ptr {
+ %0 = llvm.blockaddress <function = @h3, tag = <id = 0>> : !llvm.ptr
+ llvm.return %0 : !llvm.ptr
+}
+
+// CHECK: @h = private global ptr blockaddress(@h3, %[[BB_ADDR:.*]])
+
+// CHECK: define void @h3() {
+// CHECK: br label %[[BB_ADDR]]
+
+// CHECK: [[BB_ADDR]]:
+// CHECK: ret void
+// CHECK: }
+
+// CHECK: define void @h0()
+
+llvm.func @h3() {
+ llvm.br ^bb1
+^bb1:
+ llvm.blocktag <id = 0>
+ llvm.return
+}
+
+llvm.func @h0() {
+ llvm.return
+}
>From 23da2adda966440d4951743f59f7ff3ebf83c5cb Mon Sep 17 00:00:00 2001
From: Bruno Cardoso Lopes <bruno.cardoso at gmail.com>
Date: Tue, 13 May 2025 18:08:32 -0700
Subject: [PATCH 2/2] fix assert
---
mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h | 9 ++++++---
1 file changed, 6 insertions(+), 3 deletions(-)
diff --git a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
index f0ad1dedb8760..9e4dad5cb9953 100644
--- a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
+++ b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
@@ -142,13 +142,16 @@ class ModuleTranslation {
auto result = unresolvedBlockAddressMapping.try_emplace(op, cst);
(void)result;
assert(result.second &&
- "attempting to map a blockaddress that is already mapped");
+ "attempting to map a blockaddress operation that is already mapped");
}
/// Maps a BlockAddressAttr to its corresponding LLVM basic block.
void mapBlockAddress(BlockAddressAttr attr, llvm::BasicBlock *block) {
- assert(!blockAddressToLLVMMapping.count(blockAddressToLLVMMapping));
- blockAddressToLLVMMapping[attr] = block;
+ auto result = blockAddressToLLVMMapping.try_emplace(attr, block);
+ (void)result;
+ assert(
+ result.second &&
+ "attempting to map a blockaddress attritbute that is already mapped");
}
/// Finds the LLVM basic block that corresponds to the given BlockAddressAttr.
More information about the Mlir-commits
mailing list