-
Notifications
You must be signed in to change notification settings - Fork 15.4k
[MLIR][LLVM] Fix blockaddress mapping to LLVM blocks #139814
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[MLIR][LLVM] Fix blockaddress mapping to LLVM blocks #139814
Conversation
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.
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-llvm Author: Bruno Cardoso Lopes (bcardosolopes) ChangesAfter each function is translated, both value and block maps are erased, which makes the current mapping of blockaddresses to llvm blocks broken - the patching happens only after all functions are translated. Simplify the overall mapping, update comments, variable names and fix the bug. Full diff: https://github.com/llvm/llvm-project/pull/139814.diff 4 Files Affected:
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
+}
|
gysit
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM
Dinistro
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM, thanks for the fix.
Co-authored-by: Christian Ulmann <[email protected]>
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
After each function is translated, both value and block maps are erased, which makes the current mapping of blockaddresses to llvm blocks broken - the patching happens only after all functions are translated.
Simplify the overall mapping, update comments, variable names and fix the bug.