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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Zichen Lu (MikaOvO)

<details>
<summary>Changes</summary>

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:**
```python
# 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:

```mlir
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))
```
```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
```

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


1 Files Affected:

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


``````````diff
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 {

``````````

</details>


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