diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/DIScopeForLLVMFuncOp.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/DIScopeForLLVMFuncOp.cpp index 67573c4ee6061..d40046589e16c 100644 --- a/mlir/lib/Dialect/LLVMIR/Transforms/DIScopeForLLVMFuncOp.cpp +++ b/mlir/lib/Dialect/LLVMIR/Transforms/DIScopeForLLVMFuncOp.cpp @@ -110,7 +110,9 @@ static Location getNestedLoc(Operation *op, LLVM::DIScopeAttr scopeAttr, } static void setLexicalBlockFileAttr(Operation *op) { - if (auto callSiteLoc = dyn_cast(op->getLoc())) { + Location opLoc = op->getLoc(); + + if (auto callSiteLoc = dyn_cast(opLoc)) { auto callerLoc = callSiteLoc.getCaller(); auto calleeLoc = callSiteLoc.getCallee(); LLVM::DIScopeAttr scopeAttr; @@ -122,6 +124,43 @@ static void setLexicalBlockFileAttr(Operation *op) { op->setLoc( CallSiteLoc::get(getNestedLoc(op, scopeAttr, calleeLoc), callerLoc)); } + + return; + } + + auto funcOp = op->getParentOfType(); + if (!funcOp) + return; + + FileLineColLoc opFileLoc = extractFileLoc(opLoc); + if (!opFileLoc) + return; + + FileLineColLoc funcFileLoc = extractFileLoc(funcOp.getLoc()); + if (!funcFileLoc) + return; + + StringRef opFile = opFileLoc.getFilename().getValue(); + StringRef funcFile = funcFileLoc.getFilename().getValue(); + + if (opFile != funcFile) { + auto funcOpLoc = llvm::dyn_cast_if_present(funcOp.getLoc()); + if (!funcOpLoc) + return; + auto scopeAttr = dyn_cast(funcOpLoc.getMetadata()); + if (!scopeAttr) + return; + + auto *context = op->getContext(); + LLVM::DIFileAttr opFileAttr = + LLVM::DIFileAttr::get(context, llvm::sys::path::filename(opFile), + llvm::sys::path::parent_path(opFile)); + + auto lexicalBlockFileAttr = + LLVM::DILexicalBlockFileAttr::get(context, scopeAttr, opFileAttr, 0); + + Location newLoc = FusedLoc::get(context, {opLoc}, lexicalBlockFileAttr); + op->setLoc(newLoc); } }