|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | + |
| 3 | +""" |
| 4 | +SASS code extraction from cubin files. |
| 5 | +
|
| 6 | +Provides utilities to disassemble NVIDIA cubin files and extract |
| 7 | +SASS assembly code using nvdisasm. This module can be used standalone |
| 8 | +via CLI or programmatically by other modules (e.g., deadlock analysis). |
| 9 | +
|
| 10 | +Usage (CLI): |
| 11 | + cutracer query sass kernel.cubin |
| 12 | + cutracer query sass kernel.cubin -o kernel.sass |
| 13 | +
|
| 14 | +Usage (programmatic): |
| 15 | + from cutracer.query.sass import dump_sass, dump_sass_to_file |
| 16 | +
|
| 17 | + # Get SASS text in memory |
| 18 | + output = dump_sass(Path("kernel.cubin")) |
| 19 | + if output: |
| 20 | + print(output.raw_text) |
| 21 | +
|
| 22 | + # Save SASS to file |
| 23 | + sass_path = dump_sass_to_file(Path("kernel.cubin")) |
| 24 | +""" |
| 25 | + |
| 26 | +from __future__ import annotations |
| 27 | + |
| 28 | +import logging |
| 29 | +import subprocess |
| 30 | +from dataclasses import dataclass, field |
| 31 | +from pathlib import Path |
| 32 | + |
| 33 | +logger: logging.Logger = logging.getLogger(__name__) |
| 34 | + |
| 35 | + |
| 36 | +@dataclass |
| 37 | +class SassOutput: |
| 38 | + """Container for nvdisasm output. |
| 39 | +
|
| 40 | + Attributes: |
| 41 | + raw_text: The disassembled SASS text. |
| 42 | + cubin_path: Path to the source cubin file. |
| 43 | + flags_used: List of nvdisasm flags that were used. |
| 44 | + """ |
| 45 | + |
| 46 | + raw_text: str |
| 47 | + cubin_path: Path |
| 48 | + flags_used: list[str] = field(default_factory=list) |
| 49 | + |
| 50 | + def save(self, output_path: Path) -> None: |
| 51 | + """Save SASS text to file. |
| 52 | +
|
| 53 | + Args: |
| 54 | + output_path: Path where SASS text will be written. |
| 55 | + """ |
| 56 | + output_path.parent.mkdir(parents=True, exist_ok=True) |
| 57 | + output_path.write_text(self.raw_text) |
| 58 | + |
| 59 | + def __len__(self) -> int: |
| 60 | + """Return length of SASS text.""" |
| 61 | + return len(self.raw_text) |
| 62 | + |
| 63 | + @property |
| 64 | + def line_count(self) -> int: |
| 65 | + """Return number of lines in SASS text.""" |
| 66 | + if not self.raw_text: |
| 67 | + return 0 |
| 68 | + return len(self.raw_text.rstrip("\n").split("\n")) |
| 69 | + |
| 70 | + |
| 71 | +def dump_sass( |
| 72 | + cubin_path: Path, |
| 73 | + *, |
| 74 | + include_source_info: bool = True, |
| 75 | + include_line_info: bool = True, |
| 76 | + timeout: int = 60, |
| 77 | +) -> SassOutput | None: |
| 78 | + """ |
| 79 | + Dump SASS assembly from cubin file using nvdisasm. |
| 80 | +
|
| 81 | + This is the core function for extracting SASS code. It runs: |
| 82 | + nvdisasm -g -c <cubin_path> |
| 83 | +
|
| 84 | + Flags: |
| 85 | + -g: Include source-level debug info (file/line annotations) |
| 86 | + -c: Output assembly with //## File comments for source mapping |
| 87 | +
|
| 88 | + The combination of -g and -c provides: |
| 89 | + - Full instruction listing with PC offsets |
| 90 | + - Source file/line annotations via //## comments |
| 91 | + - Function boundaries and labels |
| 92 | +
|
| 93 | + Args: |
| 94 | + cubin_path: Absolute path to cubin file. |
| 95 | + include_source_info: If True, add -g flag for debug info. |
| 96 | + include_line_info: If True, add -c flag for //## comments. |
| 97 | + timeout: Timeout in seconds (default: 60). |
| 98 | +
|
| 99 | + Returns: |
| 100 | + SassOutput containing raw text and metadata, or None if failed. |
| 101 | +
|
| 102 | + Example: |
| 103 | + >>> output = dump_sass(Path("/path/to/kernel.cubin")) |
| 104 | + >>> if output: |
| 105 | + ... print(f"Got {output.line_count} lines of SASS") |
| 106 | + ... output.save(Path("kernel.sass")) |
| 107 | + """ |
| 108 | + if not cubin_path.exists(): |
| 109 | + logger.warning("cubin file not found: %s", cubin_path) |
| 110 | + return None |
| 111 | + |
| 112 | + flags: list[str] = [] |
| 113 | + if include_source_info: |
| 114 | + flags.append("-g") |
| 115 | + if include_line_info: |
| 116 | + flags.append("-c") |
| 117 | + |
| 118 | + cmd = ["nvdisasm", *flags, str(cubin_path)] |
| 119 | + |
| 120 | + try: |
| 121 | + result = subprocess.run( |
| 122 | + cmd, |
| 123 | + capture_output=True, |
| 124 | + text=True, |
| 125 | + timeout=timeout, |
| 126 | + ) |
| 127 | + except FileNotFoundError: |
| 128 | + logger.warning("nvdisasm not found in PATH") |
| 129 | + return None |
| 130 | + except subprocess.TimeoutExpired: |
| 131 | + logger.warning("nvdisasm timed out after %ds", timeout) |
| 132 | + return None |
| 133 | + |
| 134 | + if result.returncode != 0: |
| 135 | + logger.warning( |
| 136 | + "nvdisasm failed (rc=%d): %s", |
| 137 | + result.returncode, |
| 138 | + result.stderr.strip(), |
| 139 | + ) |
| 140 | + return None |
| 141 | + |
| 142 | + return SassOutput( |
| 143 | + raw_text=result.stdout, |
| 144 | + cubin_path=cubin_path, |
| 145 | + flags_used=flags, |
| 146 | + ) |
| 147 | + |
| 148 | + |
| 149 | +def dump_sass_to_file( |
| 150 | + cubin_path: Path, |
| 151 | + output_path: Path | None = None, |
| 152 | + *, |
| 153 | + include_source_info: bool = True, |
| 154 | + include_line_info: bool = True, |
| 155 | + timeout: int = 60, |
| 156 | +) -> Path | None: |
| 157 | + """ |
| 158 | + Dump SASS to file. Convenience wrapper around dump_sass(). |
| 159 | +
|
| 160 | + Args: |
| 161 | + cubin_path: Path to cubin file. |
| 162 | + output_path: Output .sass file path. If None, derives from cubin path |
| 163 | + by replacing extension with .sass. |
| 164 | + include_source_info: If True, add -g flag for debug info. |
| 165 | + include_line_info: If True, add -c flag for //## comments. |
| 166 | + timeout: Timeout in seconds (default: 60). |
| 167 | +
|
| 168 | + Returns: |
| 169 | + Path to output file, or None if dump failed. |
| 170 | +
|
| 171 | + Example: |
| 172 | + >>> # Creates kernel.sass next to kernel.cubin |
| 173 | + >>> sass_path = dump_sass_to_file(Path("kernel.cubin")) |
| 174 | + >>> # Or specify explicit output path |
| 175 | + >>> sass_path = dump_sass_to_file(Path("kernel.cubin"), Path("output/kernel.sass")) |
| 176 | + """ |
| 177 | + sass_output = dump_sass( |
| 178 | + cubin_path, |
| 179 | + include_source_info=include_source_info, |
| 180 | + include_line_info=include_line_info, |
| 181 | + timeout=timeout, |
| 182 | + ) |
| 183 | + if sass_output is None: |
| 184 | + return None |
| 185 | + |
| 186 | + if output_path is None: |
| 187 | + output_path = cubin_path.with_suffix(".sass") |
| 188 | + |
| 189 | + sass_output.save(output_path) |
| 190 | + return output_path |
![meta-codesync[bot]](https://avatars.githubusercontent.com/in/1373969?v=4&size=40){kind=avatar}
0 commit comments