[triton][beta] Backport '[Triton] Fix LLVMDILocalVariable pass crash with LLVM_EXTRACT_DI_LOCAL_VARIABLES (#10042)' (#1390)

Summary:

This is a backport of an upstream PR: triton-lang/triton#10042

Upstream commit message:
```
> [Triton] Fix LLVMDILocalVariable pass crash with LLVM_EXTRACT_DI_LOCAL_VARIABLES (#10042)

> Summary:
> Fix multiple bugs in the LLVMDILocalVariable pass that caused crashes
> when running with LLVM_EXTRACT_DI_LOCAL_VARIABLES=1.

> Root cause: The MLIR DIRecursiveTypeAttrInterface uses isRecSelf=true as
> a placeholder and isRecSelf=false as its resolver. During MLIR-to-LLVM
> IR translation, the resolver recursiveNodeMap entry is temporary (popped
> after translation). Any isRecSelf=true references encountered outside
> the resolver translation scope trigger an assertion.

> The add_di_scope pass creates a DISubprogramAttr with isRecSelf=true and
> embeds it in function locations and DILexicalBlockFileAttr scopes in
> operation locations. The LLVMDILocalVariable pass was supposed to
> resolve these but had several issues.

> Fixes:
> 1. Added null/scope guard for diSubprogramAttr to skip module-level ops
> that appear before or between functions in the walk.
> 2. Fixed getDISubprogramAttr(Operation op) to take Operation *op.
> 3. Restructured fuseFuncArgVariables to create two sets of
> DILocalVariableAttr: one with isRecSelf=true scope for retainedNodes
> (correct recursive pattern), and separate ones with the resolved
> isRecSelf=false scope for DbgValueOps (avoids unresolvable refs).
> 4. Added fixLexicalBlockScopes to walk operations and replace
> DILexicalBlockFileAttr scopes that reference the old isRecSelf=true
> subprogram with the resolved version.
> 5. Value-initialized diSubprogramAttr to avoid undefined behavior.

> Test Plan:
> TRITON_ALWAYS_COMPILE=1 LLVM_EXTRACT_DI_LOCAL_VARIABLES=1 \
>   python3 python/tutorials/02-fused-softmax.py

> <!---
> The core Triton is a small number of people, and we receive many PRs
> (thank
> you!).  To help us review your code more quickly, **if you are a new
> contributor (less than 3 PRs merged) we ask that you complete the
> following
> tasks and include the filled-out checklist in your PR description.**

> Complete the following tasks before sending your PR, and replace `[ ]`
> with
> `[x]` to indicate you have done them.
> -->

> # New contributor declaration
> - [x] I am not making a trivial change, such as fixing a typo in a
> comment.

> - [x] I have written a PR description following these
>   [rules](https://cbea.ms/git-commit/#why-not-how).

> - [x] I have run `pre-commit run --from-ref origin/main --to-ref HEAD`.

> - Select one of the following.
>   - [X] I have added tests.
>     - `/test` for `lit` tests
>     - `/unittest` for C++ tests
>     - `/python/test` for end-to-end tests
>   - [ ] This PR does not need a test because `FILL THIS IN`.

> - Select one of the following.
>   - [ ] I have not added any `lit` tests.
> - [x] The `lit` tests I have added follow these [best
> practices](https://mlir.llvm.org/getting_started/TestingGuide/#filecheck-best-practices),
> including the "tests should be minimal" section. (Usually running Python
> code
>     and using the instructions it generates is not minimal.)

> Co-authored-by: Bill Yoshimi <byoshimi@fb.com>
> Co-authored-by: Corbin Robeck <robeck@meta.com>
```

***Do not remove the following line from this commit***
Reactor Backport Revision: 5b3965b27eb331094854963b751709bcb148c3d4
---

This diff was generated by running:
```
buck run fbcode//triton/tools/reactor:reactor -- backport --pr 10042
```

Differential Revision: D102859687

Author: byoshimi-github

lib/Target/LLVMIR/LLVMDILocalVariable.cpp

@@ -34,6 +34,12 @@ struct LLVMDILocalVariablePass
       return;
     }
 
+    // Skip ops outside of a function (e.g. module-level globals) where
+    // diSubprogramAttr may not be valid for this op's scope.
+    if (!diSubprogramAttr || !op->getParentOfType<LLVM::LLVMFuncOp>()) {
+      return;
+    }
+
     MLIRContext *context = op->getContext();
     OpBuilder builder(context);
     Location loc = op->getLoc();
@@ -201,7 +207,7 @@ struct LLVMDILocalVariablePass
 
   // construct a subprogram of an operation by using its parent function's
   // DISubprogramAttr construction
-  LLVM::DISubprogramAttr getDISubprogramAttr(Operation op) {
+  LLVM::DISubprogramAttr getDISubprogramAttr(Operation &op) {
     auto funcOp = op.getParentOfType<LLVM::LLVMFuncOp>();
     return getDISubprogramAttr(funcOp);
   }
@@ -213,20 +219,31 @@ struct LLVMDILocalVariablePass
     MLIRContext *context = &getContext();
     OpBuilder builder(context);
     builder.setInsertionPointToStart(&funcOp.getBody().front());
-    llvm::SmallVector<mlir::LLVM::DINodeAttr> retainedNodes;
 
     LLVM::DIFileAttr fileAttr = subprogramAttr.getFile();
     LLVM::DISubroutineTypeAttr subroutineTypeAttr = subprogramAttr.getType();
     int64_t line = subprogramAttr.getLine();
-    auto localScopeAttr = dyn_cast<LLVM::DILocalScopeAttr>(subprogramAttr);
+    // The input subprogramAttr has isRecSelf=true. Use it as the scope for
+    // retainedNodes variables (proper recursive self-reference pattern).
+    auto selfRefScopeAttr = dyn_cast<LLVM::DILocalScopeAttr>(subprogramAttr);
     auto diFlag = LLVM::DIFlags::Zero;
 
-    // Extract function arguments and add them to retainedNodes:
-    // 0. Extract function argument types from subroutineTypeAttr
-    // 1. Create DILocalVariable and DebugValueOp for each arg
-    // 2. Add each arg as DILocalVariableAttr to retainedNodes
+    // Collect argument info and create retainedNodes variables scoped to the
+    // isRecSelf=true placeholder (these are nested inside the definition and
+    // resolved during translation).
     auto argTypeAttrs = subroutineTypeAttr.getTypes();
     unsigned resNum = funcOp.getNumResults() ? funcOp.getNumResults() : 1;
+
+    struct ArgInfo {
+      unsigned argIdx;
+      BlockArgument arg;
+      Location childLoc;
+      StringAttr nameAttr;
+      LLVM::DITypeAttr typeAttr;
+    };
+    llvm::SmallVector<ArgInfo> argInfos;
+    llvm::SmallVector<mlir::LLVM::DINodeAttr> retainedNodes;
+
     for (unsigned idx = resNum; idx < argTypeAttrs.size(); idx++) {
       LLVM::DITypeAttr argTypeAttr = argTypeAttrs[idx];
       unsigned argIdx = idx - resNum;
@@ -236,20 +253,18 @@ struct LLVMDILocalVariablePass
       auto nameLoc = dyn_cast<NameLoc>(argLoc);
       if (!nameLoc)
         continue;
-      Location childLoc = nameLoc.getChildLoc();
-      StringAttr nameAttr = nameLoc.getName();
 
+      // Create variable with isRecSelf=true scope for retainedNodes
       auto argVarAttr = LLVM::DILocalVariableAttr::get(
-          context, localScopeAttr, nameAttr, fileAttr, line, argIdx + 1, 0,
-          argTypeAttr, diFlag);
-
-      auto exprAttr = LLVM::DIExpressionAttr::get(context);
-      (void)LLVM::DbgValueOp::create(builder, childLoc, arg, argVarAttr,
-                                     exprAttr);
-
+          context, selfRefScopeAttr, nameLoc.getName(), fileAttr, line,
+          argIdx + 1, 0, argTypeAttr, diFlag);
       retainedNodes.push_back(argVarAttr);
+
+      argInfos.push_back(
+          {argIdx, arg, nameLoc.getChildLoc(), nameLoc.getName(), argTypeAttr});
     }
 
+    // Create the resolved subprogram (isRecSelf=false) with retainedNodes.
     mlir::DistinctAttr recId = subprogramAttr.getRecId();
     mlir::DistinctAttr id =
         mlir::DistinctAttr::create(mlir::UnitAttr::get(context));
@@ -262,23 +277,112 @@ struct LLVMDILocalVariablePass
         funcNameAttr, funcNameAttr, fileAttr, line, line, subprogramFlags,
         subroutineTypeAttr, retainedNodes, /*annotations=*/{});
 
+    // Now create DbgValueOps with variables scoped to the RESOLVED subprogram.
+    // The isRecSelf=true scope must NOT appear in DbgValueOps because the
+    // translator's recursiveNodeMap entry is only valid during translation of
+    // the isRecSelf=false definition and is popped afterwards.
+    auto resolvedScopeAttr = dyn_cast<LLVM::DILocalScopeAttr>(subprogramAttr);
+    for (auto &info : argInfos) {
+      auto argVarAttr = LLVM::DILocalVariableAttr::get(
+          context, resolvedScopeAttr, info.nameAttr, fileAttr, line,
+          info.argIdx + 1, 0, info.typeAttr, diFlag);
+      auto exprAttr = LLVM::DIExpressionAttr::get(context);
+      (void)LLVM::DbgValueOp::create(builder, info.childLoc, info.arg,
+                                     argVarAttr, exprAttr);
+    }
+
     Location loc = funcOp.getLoc();
-    // Reset the subprogramAttr with retainedNodes to the funcOp
+    // Unwrap the old FusedLoc to avoid nesting a stale isRecSelf=true
+    // subprogram in the location chain.
+    if (auto fusedLoc = dyn_cast<FusedLoc>(loc)) {
+      if (fusedLoc.getMetadata() &&
+          isa<LLVM::DISubprogramAttr>(fusedLoc.getMetadata())) {
+        loc = FusedLoc::get(context, fusedLoc.getLocations(), Attribute());
+      }
+    }
     funcOp->setLoc(mlir::FusedLoc::get(context, {loc}, subprogramAttr));
     return subprogramAttr;
   }
 
+  // After creating the resolved subprogram, fix any DILexicalBlockFileAttr
+  // locations in the function that still reference the old isRecSelf=true
+  // subprogram. These were created by add_di_scope before this pass ran.
+  void fixLexicalBlockScopes(LLVM::LLVMFuncOp funcOp,
+                             LLVM::DISubprogramAttr oldSubprogram,
+                             LLVM::DISubprogramAttr newSubprogram) {
+    MLIRContext *context = &getContext();
+    funcOp.walk([&](Operation *op) {
+      Location loc = op->getLoc();
+      Location newLoc =
+          replaceLexicalBlockScope(context, loc, oldSubprogram, newSubprogram);
+      if (newLoc != loc)
+        op->setLoc(newLoc);
+    });
+  }
+
+  Location replaceLexicalBlockScope(MLIRContext *context, Location loc,
+                                    LLVM::DISubprogramAttr oldSP,
+                                    LLVM::DISubprogramAttr newSP) {
+    if (auto fusedLoc = dyn_cast<FusedLoc>(loc)) {
+      // Recursively fix inner locations
+      SmallVector<Location> newLocs;
+      bool changed = false;
+      for (Location inner : fusedLoc.getLocations()) {
+        Location fixed = replaceLexicalBlockScope(context, inner, oldSP, newSP);
+        newLocs.push_back(fixed);
+        if (fixed != inner)
+          changed = true;
+      }
+
+      auto metadata = fusedLoc.getMetadata();
+      if (auto lexBlock =
+              dyn_cast_or_null<LLVM::DILexicalBlockFileAttr>(metadata)) {
+        if (auto scope =
+                dyn_cast_or_null<LLVM::DISubprogramAttr>(lexBlock.getScope())) {
+          if (scope == oldSP) {
+            auto newBlock = LLVM::DILexicalBlockFileAttr::get(
+                context, newSP, lexBlock.getFile(),
+                lexBlock.getDiscriminator());
+            return FusedLoc::get(context, newLocs, newBlock);
+          }
+        }
+      }
+
+      if (changed)
+        return FusedLoc::get(context, newLocs, metadata);
+    } else if (auto callSiteLoc = dyn_cast<CallSiteLoc>(loc)) {
+      Location newCallee = replaceLexicalBlockScope(
+          context, callSiteLoc.getCallee(), oldSP, newSP);
+      Location newCaller = replaceLexicalBlockScope(
+          context, callSiteLoc.getCaller(), oldSP, newSP);
+      if (newCallee != callSiteLoc.getCallee() ||
+          newCaller != callSiteLoc.getCaller())
+        return CallSiteLoc::get(newCallee, newCaller);
+    }
+    return loc;
+  }
+
   // set it while traversing into a function
-  LLVM::DISubprogramAttr diSubprogramAttr;
+  LLVM::DISubprogramAttr diSubprogramAttr = {};
 
   void runOnOperation() override {
     Operation *op = getOperation();
 
     getOperation()->walk<WalkOrder::PreOrder>([&](Operation *op) -> void {
       if (isa<LLVM::LLVMFuncOp>(op)) {
         auto funcOp = cast<LLVM::LLVMFuncOp>(op);
-        diSubprogramAttr = getDISubprogramAttr(funcOp);
-        diSubprogramAttr = fuseFuncArgVariables(funcOp, diSubprogramAttr);
+        auto oldSubprogram = getDISubprogramAttr(funcOp);
+        // External declarations (e.g., runtime builtins like vprintf) have no
+        // body, so we cannot insert debug value intrinsics into them.
+        if (!funcOp.isExternal()) {
+          diSubprogramAttr = fuseFuncArgVariables(funcOp, oldSubprogram);
+          // Fix DILexicalBlockFileAttr locations that still reference the old
+          // isRecSelf=true subprogram from add_di_scope.
+          if (oldSubprogram.getIsRecSelf())
+            fixLexicalBlockScopes(funcOp, oldSubprogram, diSubprogramAttr);
+        } else {
+          diSubprogramAttr = oldSubprogram;
+        }
       } else {
         fuseDILocalVariable(op);
       }

test/LLVMIR/di-local-variable-with-callsite-loc.mlir

@@ -0,0 +1,41 @@
+// RUN: LLVM_EXTRACT_DI_LOCAL_VARIABLES=1 triton-opt %s -o - --mlir-print-debuginfo --mlir-use-nameloc-as-prefix --enable-line-info --extract-variable-info | \
+// RUN: mlir-translate --mlir-to-llvmir | FileCheck %s
+
+// Regression test for a crash when CallSiteLoc operations coexist with
+// NameLoc operations in the same function. When enable-line-info runs with
+// LLVM_EXTRACT_DI_LOCAL_VARIABLES=1, it creates a DISubprogramAttr with
+// isRecSelf=true and wraps CallSiteLoc ops in DILexicalBlockFileAttr that
+// references the isRecSelf=true subprogram. The extract-variable-info pass
+// must fix these lexical block scopes to reference the resolved
+// (isRecSelf=false) subprogram via fixLexicalBlockScopes, otherwise
+// mlir-translate crashes with an assertion in DebugTranslation.
+
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+  llvm.func @kernel_with_callsite(%arg0: !llvm.ptr<1> loc(#loc_arg0), %arg1: i32 loc(#loc_arg1), %arg2: !llvm.ptr<1>) {
+    // CHECK-DAG: distinct !DISubprogram({{.*}}, retainedNodes:
+    // CHECK-DAG: !DILocalVariable(name: "input_ptr", arg: 1, scope:
+    // CHECK-DAG: !DILocalVariable(name: "n_elements", arg: 2, scope:
+
+    %c = llvm.mlir.constant(42 : i32) : i32
+
+    // This op has a CallSiteLoc, triggering DILexicalBlockFileAttr creation
+    // in add_di_scope. Without fixLexicalBlockScopes, the stale isRecSelf=true
+    // reference would cause an assertion failure in mlir-translate.
+    %sum = llvm.add %arg1, %c : i32 loc(#loc_callsite1)
+
+    // This op has a NameLoc, triggering DILocalVariable creation in
+    // fuseDILocalVariable.
+    // CHECK-DAG: !DILocalVariable(name: "result", scope:
+    %result = llvm.mul %sum, %c : i32 loc(#loc_result)
+
+    llvm.return
+  }
+}
+#loc_base = loc("kernel.py":10:0)
+#loc_arg0 = loc("input_ptr"(#loc_base))
+#loc_arg1 = loc("n_elements"(#loc_base))
+#loc_callee = loc("helper.py":5:3)
+#loc_caller = loc("kernel.py":20:5)
+#loc_callsite1 = loc(callsite(#loc_callee at #loc_caller))
+#loc_result_base = loc("kernel.py":25:10)
+#loc_result = loc("result"(#loc_result_base))