[MLIR][LLVM] Extend DIScopeForLLVMFuncOp to handle cross-file operatio… (#167844)

The current `DIScopeForLLVMFuncOp` pass handles debug information for
inlined code by processing `CallSiteLoc` attributes. However, some
compilation scenarios compose code from multiple source files directly
into a single function without generating `CallSiteLoc`.

**Scenario:**
```python
# a.py
def kernel_a(tensor):
    print("a: {}", tensor)  # a.py:3
    jit_func_b(tensor)           # Calls b.py code

# b.py
def func_b(tensor):
    print("b: {}", tensor)  # b.py:7
```

The scenario executes Python at compile-time and directly inserts
operations from `b.py` into the kernel function, resulting in MLIR like:

```mlir
@kernel_a(...) {
  print("a: {}", %arg0) loc(#loc_a)  // a.py:3
  print("b: {}", %arg0) loc(#loc_b)  // b.py:7 <- FileLineColLoc, not CallSiteLoc
} loc(#loc_kernel)  // a.py:1

#loc1 = loc("a.py":3:.)
#loc2 = loc("b.py":7:.)
#loc_a = loc("print"(#loc1))
#loc_b = loc("print"(#loc2))
```
```llvm
!6 = !DIFile(filename: "a.py", directory: "...")
!9 = distinct !DISubprogram(name: "...", linkageName: "...", scope: !6, file: !6, line: 13, ...)
!10 = !DILocation(line: 7, column: ., scope: !9)  // Points to kernel's DISubprogram, not correct
```

Author: MikaOvO

mlir/lib/Dialect/LLVMIR/Transforms/DIScopeForLLVMFuncOp.cpp

@@ -109,8 +109,12 @@ static Location getNestedLoc(Operation *op, LLVM::DIScopeAttr scopeAttr,
   return FusedLoc::get(context, {loc}, lexicalBlockFileAttr);
 }
 
+/// Adds DILexicalBlockFileAttr for operations with CallSiteLoc and operations
+/// from different files than their containing function.
 static void setLexicalBlockFileAttr(Operation *op) {
-  if (auto callSiteLoc = dyn_cast<CallSiteLoc>(op->getLoc())) {
+  Location opLoc = op->getLoc();
+
+  if (auto callSiteLoc = dyn_cast<CallSiteLoc>(opLoc)) {
     auto callerLoc = callSiteLoc.getCaller();
     auto calleeLoc = callSiteLoc.getCallee();
     LLVM::DIScopeAttr scopeAttr;
@@ -122,6 +126,45 @@ static void setLexicalBlockFileAttr(Operation *op) {
       op->setLoc(
           CallSiteLoc::get(getNestedLoc(op, scopeAttr, calleeLoc), callerLoc));
     }
+
+    return;
+  }
+
+  auto funcOp = op->getParentOfType<LLVM::LLVMFuncOp>();
+  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();
+
+  // Handle cross-file operations: add DILexicalBlockFileAttr when the
+  // operation's source file differs from its containing function.
+  if (opFile != funcFile) {
+    auto funcOpLoc = llvm::dyn_cast_if_present<FusedLoc>(funcOp.getLoc());
+    if (!funcOpLoc)
+      return;
+    auto scopeAttr = dyn_cast<LLVM::DISubprogramAttr>(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));
+
+    LLVM::DILexicalBlockFileAttr lexicalBlockFileAttr =
+        LLVM::DILexicalBlockFileAttr::get(context, scopeAttr, opFileAttr, 0);
+
+    Location newLoc = FusedLoc::get(context, {opLoc}, lexicalBlockFileAttr);
+    op->setLoc(newLoc);
   }
 }
 

mlir/test/Dialect/LLVMIR/add-debuginfo-func-scope.mlir

@@ -141,3 +141,22 @@ module {
   llvm.func @func_callsiteloc() loc(callsite("foo" at "mysource.cc":10:8))
 } loc(unknown)
 
+// -----
+
+// CHECK-LABEL: llvm.func @func_cross_file_op()
+// CHECK: #di_file = #llvm.di_file<"<unknown>" in "">
+// CHECK: #di_file1 = #llvm.di_file<"caller.py" in "">
+// CHECK: #di_file2 = #llvm.di_file<"callee.py" in "">
+// CHECK: #di_subroutine_type = #llvm.di_subroutine_type<callingConvention = DW_CC_normal>
+// CHECK: #di_subprogram = #llvm.di_subprogram<id = distinct[1]<>, compileUnit = #di_compile_unit, scope = #di_file1, name = "func_cross_file_op", linkageName = "func_cross_file_op", file = #di_file1, line = 5, scopeLine = 5, subprogramFlags = "Definition|Optimized", type = #di_subroutine_type>
+// CHECK: #di_lexical_block_file = #llvm.di_lexical_block_file<scope = #di_subprogram, file = #di_file2, discriminator = 0>
+
+#loc = loc("caller.py":5:1)
+#loc1 = loc("callee.py":10:5)
+
+module {
+  llvm.func @func_cross_file_op() {
+    llvm.return loc(#loc1)
+  } loc(#loc)
+} loc(unknown)
+