From 77816f6d0e8dbf65b7171d523b8ab92c9176a828 Mon Sep 17 00:00:00 2001
From: Melih Elibol <elibol@users.noreply.github.com>
Date: Sat, 10 Oct 2026 14:21:20 -0700
Subject: [PATCH] fix(cutile): attribute generated entry setup to the user
 signature

Signed-off-by: Melih Elibol <elibol@users.noreply.github.com>
---
 cutile-rs/CHANGELOG.md                        |  6 ++
 .../src/kernel_entry_generator.rs             |  7 ++-
 cutile-rs/cutile/tests/debug_info.rs          | 55 +++++++++++++++++++
 3 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/cutile-rs/CHANGELOG.md b/cutile-rs/CHANGELOG.md
index 4b7ec6eff9..9b224c06e5 100644
--- a/cutile-rs/CHANGELOG.md
+++ b/cutile-rs/CHANGELOG.md
@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
 
 ## [Unreleased]
 
+### Fixed
+
+- Generated kernel entry setup and assumptions now use the user function's
+  signature as their source location, preserving the inlined body's own
+  statement locations and helper scopes.
+
 ### Changed
 
 - cuTile Rust now lives in the [NVIDIA/cuda-rust](https://github.com/NVIDIA/cuda-rust)
diff --git a/cutile-rs/cutile-compiler/src/kernel_entry_generator.rs b/cutile-rs/cutile-compiler/src/kernel_entry_generator.rs
index 33ce9dce04..7de337b9d2 100644
--- a/cutile-rs/cutile-compiler/src/kernel_entry_generator.rs
+++ b/cutile-rs/cutile-compiler/src/kernel_entry_generator.rs
@@ -22,6 +22,7 @@ use proc_macro2::Span;
 use quote::ToTokens;
 use std::collections::HashMap;
 use syn::punctuated::Punctuated;
+use syn::spanned::Spanned;
 use syn::visit_mut::VisitMut;
 use syn::{Expr, FnArg, GenericArgument, ItemFn, ItemImpl, Lit, Stmt, Token};
 
@@ -657,8 +658,10 @@ pub fn generate_entry_point(
     statements.push(final_stmnt);
     fn_entry.block.stmts = statements;
 
-    // This source code is not visible.
-    let mut visitor = SpanSetter::new(Span::call_site());
+    // Generated entry setup has no user-written body. Attribute it, including
+    // assumptions and the call into the kernel, to the user's signature.
+    // The original kernel body keeps its own spans when it is inlined.
+    let mut visitor = SpanSetter::new(fn_item.sig.span());
     visitor.visit_item_fn_mut(&mut fn_entry);
 
     Ok((
diff --git a/cutile-rs/cutile/tests/debug_info.rs b/cutile-rs/cutile/tests/debug_info.rs
index 4b1734f85c..e0759b3a60 100644
--- a/cutile-rs/cutile/tests/debug_info.rs
+++ b/cutile-rs/cutile/tests/debug_info.rs
@@ -123,6 +123,61 @@ fn helper_instructions_keep_exact_lines_and_call_sites() {
     });
 }
 
+#[test]
+fn entry_wrapper_and_assumptions_use_the_user_signature() {
+    common::with_test_stack(|| {
+        use cutile_ir::bytecode::Opcode;
+
+        let signature_line = source_line(
+            include_str!("debug_info.rs"),
+            "fn source_kernel<const S: [i32; 1]>(out: &mut Tensor<i32, S>, x: i32) {",
+        );
+        let artifacts = compile(DebugInfoLevel::Full);
+        let mut wrapper_ops = 0;
+        let mut wrapper_assumes = 0;
+        let mut inline_ops = 0;
+        visit_ops(
+            artifacts.module(),
+            &artifacts.module().functions,
+            &mut |op| {
+                let mut location = &op.location;
+                while let Location::CallSite { caller, .. } = location {
+                    inline_ops += 1;
+                    location = caller;
+                }
+                // Constants built from synthesized types have a separate source
+                // span path. This test covers generated syntax and inline callers.
+                if op.opcode == Opcode::Constant
+                    && !matches!(op.location, Location::CallSite { .. })
+                {
+                    return;
+                }
+                if let Location::FileLineCol { filename, line, .. } = location {
+                    assert!(filename.ends_with("debug_info.rs"), "{location:?}");
+                    assert_eq!(*line, signature_line, "{:?}: {location:?}", op.opcode);
+                    if !matches!(op.location, Location::CallSite { .. }) {
+                        wrapper_ops += 1;
+                        wrapper_assumes += usize::from(op.opcode == Opcode::Assume);
+                    }
+                } else {
+                    panic!("missing wrapper origin: {location:?}");
+                }
+            },
+        );
+        assert!(wrapper_ops > 1);
+        assert!(wrapper_assumes > 0);
+        assert!(inline_ops > 0);
+
+        let dump = cutile_ir::decode_bytecode(&artifacts.bytecode().unwrap()).unwrap();
+        assert!(
+            dump.contains(&format!(
+                "line={signature_line}, name=\"source_kernel\", linkage=\"source_kernel_entry\""
+            )),
+            "{dump}"
+        );
+    });
+}
+
 #[test]
 fn explicit_debug_modes_override_build_default_and_change_cache_keys() {
     use cutile_compiler::cuda_tile_runtime_utils::TileirasOptions;