[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