use std::path::Path;
use anyhow::{Context, Result, bail};
use rusqlite::{Connection, params};
use crate::kernel::normalize_kernel_name;
const MAX_IMPORT_BYTES: u64 = 1 * 1024 * 1024 * 1024;
const MAX_NSYS_ROWS: usize = 1_000_000;
const MAX_NSYS_ROWS_PLUS_ONE: usize = MAX_NSYS_ROWS + 1;
pub fn import_nsys_rep(dest: &Connection, nsys_path: &Path, layer_id: i64) -> Result<()> {
let size = std::fs::metadata(nsys_path)
.with_context(|| format!("stat {}", nsys_path.display()))?
.len();
if size > MAX_IMPORT_BYTES {
bail!(
"NSYS capture {} is too large ({} bytes; maximum is {} bytes)",
nsys_path.display(),
size,
MAX_IMPORT_BYTES
);
}
let src = Connection::open_with_flags(
nsys_path,
rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY | rusqlite::OpenFlags::SQLITE_OPEN_NO_MUTEX,
)
.with_context(|| format!("cannot open {}", nsys_path.display()))?;
let recognized = [
"CUPTI_ACTIVITY_KIND_KERNEL",
"CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL",
"CUPTI_ACTIVITY_KIND_MEMCPY",
"CUDA_GPU_MEMORY_USAGE_EVENTS",
"NVTX_EVENTS",
"CUPTI_ACTIVITY_KIND_RUNTIME",
]
.iter()
.any(|table| find_table(&src, &[*table]).is_ok());
if !recognized {
bail!("NSYS report contains no recognized activity tables");
}
let has_kernels = import_kernels(dest, &src, layer_id)?;
import_transfers(dest, &src, layer_id)?;
import_allocations(dest, &src, layer_id)?;
import_nvtx_regions(dest, &src, layer_id)?;
import_device_info(dest, &src)?;
let runtime_api_count = if !has_kernels {
import_runtime_api(dest, &src, layer_id)?
} else {
0
};
import_wall_time(dest, layer_id)?;
let imported: i64 = dest.query_row(
"SELECT
(SELECT COUNT(*) FROM launches WHERE layer_id=?1) +
(SELECT COUNT(*) FROM transfers WHERE layer_id=?1) +
(SELECT COUNT(*) FROM allocations WHERE layer_id=?1) +
(SELECT COUNT(*) FROM regions WHERE layer_id=?1)",
params![layer_id],
|row| row.get(0),
)?;
if imported == 0 && runtime_api_count == 0 {
bail!("NSYS report contains no recognized activity rows");
}
Ok(())
}
fn import_kernels(dest: &Connection, src: &Connection, layer_id: i64) -> Result<bool> {
let table = match find_table(
src,
&[
"CUPTI_ACTIVITY_KIND_KERNEL",
"CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL",
],
) {
Ok(t) => t,
Err(_) => return Ok(false),
};
let has_string_ids = find_table(src, &["StringIds"]).is_ok();
let name_is_integer = is_column_integer(src, &table, "demangledName");
let use_join = has_string_ids && name_is_integer;
let sql = if use_join {
format!(
"SELECT s.value, k.start, k.end,
k.gridX, k.gridY, k.gridZ,
k.blockX, k.blockY, k.blockZ,
k.streamId, k.correlationId
FROM {table} k
JOIN StringIds s ON s.id = k.demangledName
ORDER BY k.start LIMIT {}",
MAX_NSYS_ROWS_PLUS_ONE
)
} else {
format!(
"SELECT demangledName, start, end,
gridX, gridY, gridZ,
blockX, blockY, blockZ,
streamId, correlationId
FROM {table}
ORDER BY start LIMIT {}",
MAX_NSYS_ROWS_PLUS_ONE
)
};
let mut read = src.prepare(&sql)?;
let mut write = dest.prepare(
"INSERT INTO launches
(kernel_name, duration_us, grid_x, grid_y, grid_z,
block_x, block_y, block_z, stream_id, start_us,
correlation_id, layer_id)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)",
)?;
let rows = read.query_map([], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, i64>(1)?, row.get::<_, i64>(2)?, row.get::<_, u32>(3)?, row.get::<_, u32>(4)?,
row.get::<_, u32>(5)?,
row.get::<_, u32>(6)?, row.get::<_, u32>(7)?,
row.get::<_, u32>(8)?,
row.get::<_, u32>(9)?, row.get::<_, i64>(10)?, ))
})?;
let mut row_count = 0;
for row in rows {
row_count += 1;
if row_count > MAX_NSYS_ROWS {
bail!("NSYS kernel table exceeds the maximum of {MAX_NSYS_ROWS} rows");
}
let (raw_name, start_ns, end_ns, gx, gy, gz, bx, by, bz, sid, cid) = row?;
let name = normalize_kernel_name(&raw_name);
let duration_us = duration_us(start_ns, end_ns, "kernel")?;
let start_us = start_ns as f64 / 1000.0;
write.execute(params![
name,
duration_us,
gx,
gy,
gz,
bx,
by,
bz,
sid,
start_us,
cid,
layer_id
])?;
}
Ok(row_count > 0)
}
fn import_transfers(dest: &Connection, src: &Connection, layer_id: i64) -> Result<()> {
let table = match find_table(src, &["CUPTI_ACTIVITY_KIND_MEMCPY"]) {
Ok(t) => t,
Err(_) => return Ok(()),
};
let mut read = src.prepare(&format!(
"SELECT copyKind, start, end, bytes, streamId FROM {table} ORDER BY start LIMIT {MAX_NSYS_ROWS_PLUS_ONE}"
))?;
let mut write = dest.prepare(
"INSERT INTO transfers (kind, bytes, duration_us, start_us, stream_id, layer_id)
VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
)?;
let rows = read.query_map([], |row| {
Ok((
row.get::<_, i32>(0)?,
row.get::<_, i64>(1)?, row.get::<_, i64>(2)?, row.get::<_, i64>(3)?, row.get::<_, u32>(4)?, ))
})?;
let mut row_count = 0;
for row in rows {
row_count += 1;
if row_count > MAX_NSYS_ROWS {
bail!("NSYS transfer table exceeds the maximum of {MAX_NSYS_ROWS} rows");
}
let (kind, start_ns, end_ns, bytes, sid) = row?;
let kind_str = match kind {
1 => "H2D",
2 => "D2H",
3 => "D2D",
_ => "Peer",
};
write.execute(params![
kind_str,
bytes,
duration_us(start_ns, end_ns, "memory transfer")?,
start_ns as f64 / 1000.0,
sid,
layer_id
])?;
}
Ok(())
}
fn import_allocations(dest: &Connection, src: &Connection, layer_id: i64) -> Result<()> {
let table = match find_table(src, &["CUDA_GPU_MEMORY_USAGE_EVENTS"]) {
Ok(t) => t,
Err(_) => return Ok(()),
};
let select_sid = if has_column(src, &table, "streamId") {
"streamId"
} else {
"NULL"
};
let mut read = src.prepare(&format!(
"SELECT start, memoryOperationType, address, bytes, {select_sid} FROM {table} ORDER BY start LIMIT {MAX_NSYS_ROWS_PLUS_ONE}"
))?;
let mut write = dest.prepare(
"INSERT INTO allocations (op, address, bytes, start_us, stream_id, layer_id)
VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
)?;
let rows = read.query_map([], |row| {
Ok((
row.get::<_, i64>(0)?, row.get::<_, i32>(1)?, row.get::<_, i64>(2)?, row.get::<_, i64>(3)?, row.get::<_, Option<u32>>(4)?, ))
})?;
let mut row_count = 0;
for row in rows {
row_count += 1;
if row_count > MAX_NSYS_ROWS {
bail!("NSYS allocation table exceeds the maximum of {MAX_NSYS_ROWS} rows");
}
let (start_ns, oper, addr, bytes, sid) = row?;
let op = match oper {
0 => "alloc",
1 => "free",
_ => continue,
};
write.execute(params![
op,
addr,
bytes,
start_ns as f64 / 1000.0,
sid,
layer_id
])?;
}
Ok(())
}
fn import_nvtx_regions(dest: &Connection, src: &Connection, layer_id: i64) -> Result<()> {
let table = match find_table(src, &["NVTX_EVENTS"]) {
Ok(t) => t,
Err(_) => return Ok(()),
};
let mut read = src.prepare(&format!(
"SELECT text, start, end FROM {table}
WHERE text IS NOT NULL
ORDER BY start LIMIT {MAX_NSYS_ROWS_PLUS_ONE}"
))?;
let mut write = dest.prepare(
"INSERT INTO regions (name, start_us, duration_us, layer_id)
VALUES (?1, ?2, ?3, ?4)",
)?;
let rows = read.query_map([], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, i64>(1)?,
row.get::<_, i64>(2)?,
))
})?;
let mut row_count = 0;
for row in rows {
row_count += 1;
if row_count > MAX_NSYS_ROWS {
bail!("NSYS NVTX table exceeds the maximum of {MAX_NSYS_ROWS} rows");
}
let (name, start_ns, end_ns) = row?;
write.execute(params![
name,
start_ns as f64 / 1000.0,
duration_us(start_ns, end_ns, "NVTX region")?,
layer_id
])?;
}
Ok(())
}
fn import_runtime_api(dest: &Connection, src: &Connection, _layer_id: i64) -> Result<usize> {
let table = match find_table(src, &["CUPTI_ACTIVITY_KIND_RUNTIME"]) {
Ok(t) => t,
Err(_) => return Ok(0),
};
let has_strings = find_table(src, &["StringIds"]).is_ok();
if !has_strings {
return Ok(0);
}
let sql = format!(
"SELECT s.value, r.start, r.end, r.correlationId
FROM {table} r
JOIN StringIds s ON s.id = r.nameId
WHERE s.value LIKE 'cudaLaunchKernel%'
ORDER BY r.start LIMIT {MAX_NSYS_ROWS_PLUS_ONE}"
);
let mut read = src.prepare(&sql)?;
let rows = read.query_map([], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, i64>(1)?,
row.get::<_, i64>(2)?,
row.get::<_, Option<i64>>(3)?,
))
})?;
let mut count = 0;
let mut total_runtime_us = 0.0;
let mut row_count = 0;
for row in rows {
row_count += 1;
if row_count > MAX_NSYS_ROWS {
bail!("NSYS runtime table exceeds the maximum of {MAX_NSYS_ROWS} rows");
}
let (_api_name, start_ns, end_ns, _corr_id) = row?;
let duration_us = duration_us(start_ns, end_ns, "runtime API event")?;
total_runtime_us += duration_us;
count += 1;
}
if count > 0 {
dest.execute(
"INSERT OR REPLACE INTO meta (key, value) VALUES ('nsys_warning', ?1)",
params!["GPU kernel tracing unavailable (WSL2 or missing permissions). \
Showing CPU-side cudaLaunchKernel API timing only. \
For full GPU profiling, run on native Linux with root or appropriate permissions."],
)?;
dest.execute(
"INSERT OR REPLACE INTO meta (key, value) VALUES ('nsys_runtime_api_count', ?1)",
params![count.to_string()],
)?;
dest.execute(
"INSERT OR REPLACE INTO meta (key, value) VALUES ('nsys_runtime_api_time_us', ?1)",
params![total_runtime_us.to_string()],
)?;
}
Ok(count)
}
fn import_device_info(dest: &Connection, src: &Connection) -> Result<()> {
let table = match find_table(src, &["TARGET_INFO_CUDA_GPU", "TARGET_INFO_GPU"]) {
Ok(t) => t,
Err(_) => return Ok(()),
};
let name: Option<String> = src
.query_row(&format!("SELECT name FROM {table} LIMIT 1"), [], |row| {
row.get(0)
})
.ok();
if let Some(name) = name {
dest.execute(
"INSERT OR REPLACE INTO meta (key, value) VALUES ('device', ?1)",
params![name],
)?;
}
Ok(())
}
pub(crate) fn import_wall_time(dest: &Connection, layer_id: i64) -> Result<()> {
let wall: f64 = dest.query_row(
"SELECT COALESCE(MAX(end_us) - MIN(start_us), 0) FROM (
SELECT start_us, start_us + duration_us AS end_us
FROM launches WHERE start_us IS NOT NULL AND layer_id = ?1
UNION ALL
SELECT start_us, start_us + duration_us AS end_us
FROM transfers WHERE start_us IS NOT NULL AND layer_id = ?1
)",
params![layer_id],
|row| row.get(0),
)?;
if !wall.is_finite() || wall < 0.0 {
bail!("NSYS wall-time span is invalid: {wall}");
}
dest.execute(
"INSERT OR REPLACE INTO meta (key, value) VALUES ('wall_time_us', ?1)",
params![wall.to_string()],
)?;
Ok(())
}
fn duration_us(start_ns: i64, end_ns: i64, kind: &str) -> Result<f64> {
if end_ns < start_ns {
bail!("NSYS {kind} has invalid temporal ordering: end {end_ns} precedes start {start_ns}");
}
let delta = end_ns
.checked_sub(start_ns)
.ok_or_else(|| anyhow::anyhow!("NSYS {kind} timestamp difference overflows i64"))?;
Ok(delta as f64 / 1000.0)
}
fn find_table(conn: &Connection, candidates: &[&str]) -> Result<String> {
for name in candidates {
let exists: bool = conn
.query_row(
"SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name=?1",
[name],
|row| row.get::<_, i64>(0),
)
.map(|c| c > 0)
.unwrap_or(false);
if exists {
return Ok(name.to_string());
}
}
bail!("no matching table (tried: {})", candidates.join(", "))
}
fn has_column(conn: &Connection, table: &str, column: &str) -> bool {
let sql = format!("PRAGMA table_info({table})");
let Ok(mut stmt) = conn.prepare(&sql) else {
return false;
};
let rows = stmt.query_map([], |row| row.get::<_, String>(1));
match rows {
Ok(rows) => rows.flatten().any(|name| name.eq_ignore_ascii_case(column)),
Err(_) => false,
}
}
fn is_column_integer(conn: &Connection, table: &str, column: &str) -> bool {
let sql = format!("PRAGMA table_info({table})");
let mut stmt = match conn.prepare(&sql) {
Ok(s) => s,
Err(_) => return false,
};
let rows = stmt.query_map([], |row| {
Ok((
row.get::<_, String>(1)?, row.get::<_, String>(2)?, ))
});
match rows {
Ok(rows) => {
for row in rows.flatten() {
if row.0.eq_ignore_ascii_case(column) {
return row.1.eq_ignore_ascii_case("INTEGER");
}
}
false
}
Err(_) => false,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::db::GpuDb;
#[test]
fn find_table_missing() {
let conn = Connection::open_in_memory().unwrap();
assert!(find_table(&conn, &["NOPE"]).is_err());
}
#[test]
fn nsys_rejects_invalid_order_and_extreme_timestamp_difference() {
let error = duration_us(20, 10, "kernel").unwrap_err().to_string();
assert!(error.contains("invalid temporal ordering"));
let error = duration_us(i64::MIN, i64::MAX, "kernel")
.unwrap_err()
.to_string();
assert!(error.contains("overflows i64"));
}
#[test]
fn select_for_allocations_omits_stream_id_when_absent() {
let src = Connection::open_in_memory().unwrap();
src.execute_batch(
"CREATE TABLE CUDA_GPU_MEMORY_USAGE_EVENTS (
start INTEGER,
memoryOperationType INTEGER,
address INTEGER,
bytes INTEGER
);
INSERT INTO CUDA_GPU_MEMORY_USAGE_EVENTS VALUES (100, 0, 1, 4096);",
)
.unwrap();
assert!(!has_column(
&src,
"CUDA_GPU_MEMORY_USAGE_EVENTS",
"streamId"
));
let select_sid = if has_column(&src, "CUDA_GPU_MEMORY_USAGE_EVENTS", "streamId") {
"streamId"
} else {
"NULL"
};
let sql = format!(
"SELECT start, memoryOperationType, address, bytes, {select_sid} \
FROM CUDA_GPU_MEMORY_USAGE_EVENTS ORDER BY start"
);
let mut stmt = src.prepare(&sql).expect("SELECT must compile");
let got: Option<u32> = stmt.query_row([], |r| r.get::<_, Option<u32>>(4)).unwrap();
assert!(got.is_none(), "stream_id must be NULL on 2023 schema");
}
#[test]
fn select_for_allocations_uses_stream_id_when_present() {
let src = Connection::open_in_memory().unwrap();
src.execute_batch(
"CREATE TABLE CUDA_GPU_MEMORY_USAGE_EVENTS (
start INTEGER,
memoryOperationType INTEGER,
address INTEGER,
bytes INTEGER,
streamId INTEGER
);
INSERT INTO CUDA_GPU_MEMORY_USAGE_EVENTS VALUES (100, 0, 1, 4096, 7);",
)
.unwrap();
assert!(has_column(&src, "CUDA_GPU_MEMORY_USAGE_EVENTS", "streamId"));
let select_sid = "streamId"; let sql = format!(
"SELECT start, memoryOperationType, address, bytes, {select_sid} \
FROM CUDA_GPU_MEMORY_USAGE_EVENTS"
);
let got: u32 = src
.prepare(&sql)
.unwrap()
.query_row([], |r| r.get(4))
.unwrap();
assert_eq!(got, 7);
}
#[test]
fn has_column_detects_presence_and_absence() {
let conn = Connection::open_in_memory().unwrap();
conn.execute_batch("CREATE TABLE t (a INTEGER, b TEXT)")
.unwrap();
assert!(has_column(&conn, "t", "a"));
assert!(has_column(&conn, "t", "B")); assert!(!has_column(&conn, "t", "c"));
}
#[test]
fn import_wall_time_empty() {
let db = GpuDb::create(&tempfile::tempdir().unwrap().keep().join("t.db")).unwrap();
let layer = db.add_layer("nsys", "empty", None, None, None).unwrap();
import_wall_time(&db.conn, layer).unwrap();
assert_eq!(db.meta("wall_time_us"), "0");
}
#[test]
fn import_wall_time_does_not_mix_capture_layers() {
let db = GpuDb::create(&tempfile::tempdir().unwrap().keep().join("layers.db")).unwrap();
let old = db.add_layer("torch", "old", None, None, None).unwrap();
let current = db.add_layer("nsys", "new", None, None, None).unwrap();
db.conn
.execute(
"INSERT INTO launches (kernel_name, start_us, duration_us, layer_id)
VALUES ('old', 10000, 500, ?1), ('new', 20, 5, ?2)",
params![old, current],
)
.unwrap();
import_wall_time(&db.conn, current).unwrap();
assert_eq!(db.meta("wall_time_us"), "5");
}
#[test]
fn empty_or_unrecognized_nsys_report_is_rejected() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("empty.sqlite");
Connection::open(&path).unwrap();
let db = GpuDb::create(&dir.path().join("dest.db")).unwrap();
let layer = db
.add_layer("nsys", &path.to_string_lossy(), None, None, None)
.unwrap();
let error = import_nsys_rep(&db.conn, &path, layer).unwrap_err();
assert!(error.to_string().contains("no recognized activity tables"));
}
#[test]
fn empty_kernel_table_does_not_suppress_runtime_fallback() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("runtime.sqlite");
let src = Connection::open(&path).unwrap();
src.execute_batch(
"CREATE TABLE CUPTI_ACTIVITY_KIND_KERNEL (
demangledName TEXT, start INTEGER, end INTEGER,
gridX INTEGER, gridY INTEGER, gridZ INTEGER,
blockX INTEGER, blockY INTEGER, blockZ INTEGER,
streamId INTEGER, correlationId INTEGER
);
CREATE TABLE StringIds (id INTEGER, value TEXT);
CREATE TABLE CUPTI_ACTIVITY_KIND_RUNTIME (
nameId INTEGER, start INTEGER, end INTEGER, correlationId INTEGER
);
INSERT INTO StringIds VALUES (1, 'cudaLaunchKernel');
INSERT INTO CUPTI_ACTIVITY_KIND_RUNTIME VALUES (1, 100, 200, 7);",
)
.unwrap();
drop(src);
let db = GpuDb::create(&dir.path().join("dest.db")).unwrap();
let layer = db
.add_layer("nsys", &path.to_string_lossy(), None, None, None)
.unwrap();
import_nsys_rep(&db.conn, &path, layer).unwrap();
assert_eq!(db.total_launch_count(), 0);
assert_eq!(db.meta("nsys_runtime_api_count"), "1");
assert_eq!(db.meta("nsys_runtime_api_time_us"), "0.1");
}
}