Spaces:
Running
Running
| import platform | |
| import shlex | |
| import subprocess | |
| import tarfile | |
| import tempfile | |
| import urllib.request | |
| from pathlib import Path | |
| import gradio as gr | |
| import nvidia | |
| def find_tool(name): | |
| path = next(p for d in nvidia.__path__ if (p := Path(d) / "cu13/bin" / name).exists()) | |
| path.chmod(0o755) | |
| return str(path) | |
| def fetch_legacy_nvdisasm(version): | |
| arch = "linux-sbsa" if platform.machine() == "aarch64" else "linux-x86_64" | |
| name = f"cuda_nvdisasm-{arch}-{version}-archive" | |
| path = Path(tempfile.gettempdir()) / name / "nvdisasm" | |
| if not path.exists(): | |
| url = f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvdisasm/{arch}/{name}.tar.xz" | |
| with urllib.request.urlopen(url) as r, tarfile.open(fileobj=r, mode="r|xz") as tar: | |
| member = next(m for m in tar if m.name == f"{name}/bin/nvdisasm") | |
| path.parent.mkdir(exist_ok=True) | |
| path.write_bytes(tar.extractfile(member).read()) | |
| path.chmod(0o755) | |
| return str(path) | |
| TOOLS = {name: find_tool(name) for name in ("cuobjdump", "nvdisasm")} | |
| # nvdisasm 13.x can't run flow analysis (CFGs) on old-format CUDA ELF (OSABI 0x33, pre-CUDA 13); 12.9 can. | |
| LEGACY_NVDISASM = fetch_legacy_nvdisasm("12.9.88") | |
| OUTPUTS = { | |
| "cuobjdump": ["--dump-sass", "--dump-ptx", "--dump-elf", "--dump-resource-usage", "--dump-elf-symbols", | |
| "--list-elf", "--list-ptx"], | |
| "nvdisasm": ["--print-code", "--output-control-flow-graph", "--output-control-flow-graph-with-basic-blocks"], | |
| } | |
| MODIFIERS = { | |
| "cuobjdump": [], | |
| "nvdisasm": ["--print-instruction-encoding", "--print-life-ranges", "--print-raw", "--print-line-info", | |
| "--print-line-info-inline"], | |
| } | |
| EXTRA_EXAMPLE = {"cuobjdump": "--function vecadd", "nvdisasm": "--emit-json"} | |
| GRAPHVIZ_HEAD = "".join(f'<script src="https://cdn.jsdelivr.net/npm/{lib}"></script>' for lib in ( | |
| "d3@7.9.0/dist/d3.min.js", "@hpcc-js/wasm@2.20.0/dist/graphviz.umd.js", "d3-graphviz@5.6.0/build/d3-graphviz.min.js")) | |
| GRAPHVIZ_JS = """ | |
| const graph = element.querySelector('.graph'); | |
| d3.select(graph).graphviz({useWorker: false, zoom: true}).renderDot(props.value, () => { | |
| const svg = graph.querySelector('svg'); | |
| svg.setAttribute('width', '100%'); | |
| svg.setAttribute('height', Math.min(svg.viewBox.baseVal.height * 4 / 3, 600)); | |
| }); | |
| """ | |
| EXAMPLES = sorted(str(p) for p in Path(__file__).with_name("examples").glob("*.cubin")) | |
| def version(exe): | |
| return subprocess.run([exe, "--version"], capture_output=True, text=True).stdout.split("release ")[-1].splitlines()[0] | |
| def is_old_elf(file): | |
| with open(file, "rb") as f: | |
| head = f.read(8) | |
| return head[:4] == b"\x7fELF" and head[7] == 0x33 | |
| def executable(tool, file): | |
| return LEGACY_NVDISASM if tool == "nvdisasm" and is_old_elf(file) else TOOLS[tool] | |
| def run_one(tool, args, file): | |
| try: | |
| r = subprocess.run([executable(tool, file), *args, file], capture_output=True, text=True, timeout=60) | |
| except subprocess.TimeoutExpired: | |
| return "Timed out after 60s." | |
| return (r.stdout + r.stderr or "(no output)") + (f"\n[exit code {r.returncode}]" if r.returncode else "") | |
| def run(file, tool, outputs, modifiers, extra): | |
| if file is None: | |
| return {"error": "Upload a cubin first."} | |
| common = [*modifiers, *shlex.split(extra or "")] | |
| return {" ".join(args): run_one(tool, args, file) for args in ([o, *common] for o in outputs)} or \ | |
| {" ".join(common) or tool: run_one(tool, common, file)} | |
| def select_tool(tool): | |
| return (gr.CheckboxGroup(OUTPUTS[tool], value=OUTPUTS[tool]), | |
| gr.CheckboxGroup(MODIFIERS[tool], value=[], visible=bool(MODIFIERS[tool])), | |
| gr.Textbox(placeholder=EXTRA_EXAMPLE[tool])) | |
| with gr.Blocks(title="cuobjdump / nvdisasm") as demo: | |
| gr.Markdown("# cuobjdump / nvdisasm\n" | |
| "Disassemble CUDA cubins, fatbins and ELF objects with NVIDIA's tools (" | |
| + ", ".join(f"`{t}` {version(exe)}" for t, exe in TOOLS.items()) | |
| + f"; old-format CUDA ELF goes to `nvdisasm` {version(LEGACY_NVDISASM)}).") | |
| with gr.Row(equal_height=True): | |
| with gr.Column(scale=1): | |
| file = gr.File(label="cubin / fatbin / CUDA ELF", type="filepath") | |
| with gr.Column(scale=2): | |
| tool = gr.Radio(list(TOOLS), value="cuobjdump", label="Tool") | |
| outputs = gr.CheckboxGroup(OUTPUTS["cuobjdump"], value=OUTPUTS["cuobjdump"], label="Outputs (each runs separately)") | |
| modifiers = gr.CheckboxGroup(MODIFIERS["cuobjdump"], label="Display options (applied to every run)", visible=False) | |
| with gr.Row(equal_height=True): | |
| extra = gr.Textbox(label="Extra arguments", placeholder=EXTRA_EXAMPLE["cuobjdump"], scale=4) | |
| btn = gr.Button("Run", variant="primary", scale=1) | |
| gr.Examples(EXAMPLES, file) | |
| results = gr.State({}) | |
| def show(results): | |
| for label, output in results.items(): | |
| if output.startswith("digraph"): | |
| gr.Markdown(f"`{label}` (scroll to zoom, drag to pan)") | |
| gr.HTML(output.replace("\t", " "), head=GRAPHVIZ_HEAD, js_on_load=GRAPHVIZ_JS, | |
| html_template='<div class="graph" style="background:white;' | |
| 'border:1px solid var(--border-color-primary)"></div>') | |
| with gr.Accordion(f"{label} (DOT source)", open=False): | |
| gr.Code(output, max_lines=60) | |
| else: | |
| gr.Code(output, label=label, max_lines=60) | |
| tool.change(select_tool, tool, [outputs, modifiers, extra]) | |
| (btn.click(lambda: gr.Button("Running…", interactive=False), None, btn) | |
| .then(run, [file, tool, outputs, modifiers, extra], results) | |
| .then(lambda: gr.Button("Run", interactive=True), None, btn)) | |
| if __name__ == "__main__": | |
| demo.launch() | |