[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 (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 <- 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

[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]

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:

git-clang-format --diff origin/main HEAD --extensions cpp -- mlir/lib/Dialect/LLVMIR/Transforms/DIScopeForLLVMFuncOp.cpp --diff_from_common_commit

⚠️
The reproduction instructions above might return results for more than one PR
in a stack if you are using a stacked PR workflow. You can limit the results by
changing origin/main to the base branch/commit you want to compare against.
⚠️

View the diff from clang-format here.

diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/DIScopeForLLVMFuncOp.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/DIScopeForLLVMFuncOp.cpp
index 8862cf8e7..3ba289fbd 100644
--- a/mlir/lib/Dialect/LLVMIR/Transforms/DIScopeForLLVMFuncOp.cpp
+++ b/mlir/lib/Dialect/LLVMIR/Transforms/DIScopeForLLVMFuncOp.cpp
@@ -136,14 +136,14 @@ static void setLexicalBlockFileAttr(Operation *op) {
 
   FileLineColLoc funcFileLoc = extractFileLoc(funcOp.getLoc());
   if (!funcFileLoc)
-    return;  
+    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) 
+    if (!funcOpLoc)
       return;
     auto scopeAttr = dyn_cast<LLVM::DISubprogramAttr>(funcOpLoc.getMetadata());
     if (!scopeAttr)
@@ -154,12 +154,12 @@ static void setLexicalBlockFileAttr(Operation *op) {
         LLVM::DIFileAttr::get(context, llvm::sys::path::filename(opFile),
                               llvm::sys::path::parent_path(opFile));
 
-    auto lexicalBlockFileAttr = LLVM::DILexicalBlockFileAttr::get(
-        context, scopeAttr, opFileAttr, 0);
+    auto lexicalBlockFileAttr =
+        LLVM::DILexicalBlockFileAttr::get(context, scopeAttr, opFileAttr, 0);
 
     Location newLoc = FusedLoc::get(context, {opLoc}, lexicalBlockFileAttr);
     op->setLoc(newLoc);
-  }  
+  }
 }
 
 namespace {