[MLIR][LLVM] Extend DIScopeForLLVMFuncOp to handle cross-file operatio…

[MikaOvO]

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:

# 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:

@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))

!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

[llvmbot]

@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir

Author: Zichen Lu (MikaOvO)

Changes

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

Scenario:

# a.py
@<!-- -->cute.kernel
def kernel_a(tensor):
    cute.print("a: {}", tensor)  # a.py:3
    jit_func_b(tensor)           # Calls b.py code

# b.py
@<!-- -->cute.jit
def jit_func_b(tensor):
    cute.print("b: {}", tensor)  # b.py:7

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

cuda.kernel @<!-- -->kernel_a(...) {
  cute.print("a: {}", %arg0) loc(#loc_a)  // a.py:3
  cute.print("b: {}", %arg0) loc(#loc_b)  // b.py:7 &lt;- FileLineColLoc, not CallSiteLoc
} loc(#loc_kernel)  // a.py:1

#loc1 = loc("a.py":3:.)
#loc2 = loc("b.py":7:.)
#loc_a = loc("cute.printf"(#loc1))
#loc_b = loc("cute.printf"(#loc2))

!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

---

Full diff: https://github.com/llvm/llvm-project/pull/167844.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/LLVMIR/Transforms/DIScopeForLLVMFuncOp.cpp (+37)

diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/DIScopeForLLVMFuncOp.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/DIScopeForLLVMFuncOp.cpp
index 67573c4ee6061..8862cf8e70366 100644
--- a/mlir/lib/Dialect/LLVMIR/Transforms/DIScopeForLLVMFuncOp.cpp
+++ b/mlir/lib/Dialect/LLVMIR/Transforms/DIScopeForLLVMFuncOp.cpp
@@ -122,7 +122,44 @@ 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();
+
+  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));
+
+    auto lexicalBlockFileAttr = LLVM::DILexicalBlockFileAttr::get(
+        context, scopeAttr, opFileAttr, 0);
+
+    Location newLoc = FusedLoc::get(context, {opLoc}, lexicalBlockFileAttr);
+    op->setLoc(newLoc);
+  }  
 }
 
 namespace {

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[MikaOvO]

@joker-eph @grypp Could you help review it? Thanks!

[grypp]

LLVM codebase doesn't use auto when the type is not obvious.

[MikaOvO]

Updated it, thanks!

[grypp]

how can we test this code?

[MikaOvO]

If this change is accepted, I plan to add a simple MLIR file generated from merging multiple source files and check whether the location information in the results includes multiple files.

[github-actions]

🐧 Linux x64 Test Results

• 7101 tests passed
• 594 tests skipped

[joker-eph]

Can you add some high-level documentation before this new code block? I would document the logic at the function level I think.

[MikaOvO]

Updated, thanks!