Created
June 9, 2026 05:11
-
-
Save galv/5b64487b4e95b20cd03a288777a28d25 to your computer and use it in GitHub Desktop.
Stream capture works on null stream
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #!/usr/bin/env python3 | |
| """Create a green context, make it current, and capture work on its stream.""" | |
| from __future__ import annotations | |
| from cuda.bindings import driver | |
| CAPTURE_BYTES = 4096 | |
| def _decode(value: object) -> str: | |
| if isinstance(value, bytes): | |
| return value.decode("utf-8", errors="replace") | |
| return str(value) | |
| def _driver_error(err: driver.CUresult) -> str: | |
| name_result, name = driver.cuGetErrorName(err) | |
| string_result, message = driver.cuGetErrorString(err) | |
| name_text = err.name | |
| if name_result == driver.CUresult.CUDA_SUCCESS: | |
| name_text = _decode(name) | |
| message_text = "unknown" | |
| if string_result == driver.CUresult.CUDA_SUCCESS: | |
| message_text = _decode(message) | |
| return f"{name_text} ({int(err)}): {message_text}" | |
| def check(result: tuple[driver.CUresult, ...]): | |
| err = result[0] | |
| if err != driver.CUresult.CUDA_SUCCESS: | |
| raise RuntimeError(f"CUDA driver error: {_driver_error(err)}") | |
| if len(result) == 1: | |
| return None | |
| if len(result) == 2: | |
| return result[1] | |
| return result[1:] | |
| def cleanup(label: str, fn, *args) -> None: | |
| if args[0] is None: | |
| return | |
| err = fn(*args)[0] | |
| if err != driver.CUresult.CUDA_SUCCESS: | |
| print(f"warning: failed to destroy {label}: {_driver_error(err)}") | |
| def main() -> None: | |
| check(driver.cuInit(0)) | |
| driver_version = check(driver.cuDriverGetVersion()) | |
| device = check(driver.cuDeviceGet(0)) | |
| device_name = _decode(check(driver.cuDeviceGetName(256, device))).split("\x00", 1)[0] | |
| print(f"CUDA driver version: {driver_version}") | |
| print(f"Using device 0: {device_name}") | |
| sm_resource = check( | |
| driver.cuDeviceGetDevResource( | |
| device, | |
| driver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM, | |
| ) | |
| ) | |
| resource_desc = check(driver.cuDevResourceGenerateDesc([sm_resource], 1)) | |
| green_ctx = None | |
| context = None | |
| stream = None | |
| allocation = None | |
| graph = None | |
| graph_exec = None | |
| try: | |
| green_ctx = check( | |
| driver.cuGreenCtxCreate( | |
| resource_desc, | |
| device, | |
| int(driver.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM), | |
| ) | |
| ) | |
| green_id = check(driver.cuGreenCtxGetId(green_ctx)) | |
| print(f"Created green context id: {green_id}") | |
| context = check(driver.cuCtxFromGreenCtx(green_ctx)) | |
| check(driver.cuCtxSetCurrent(context)) | |
| print("Converted the green context to a CUcontext and made it current.") | |
| # this is the null stream!!! | |
| stream = 0 | |
| allocation = check(driver.cuMemAlloc(CAPTURE_BYTES)) | |
| check(driver.cuMemsetD8Async(allocation, 0, CAPTURE_BYTES, stream)) | |
| check(driver.cuStreamSynchronize(stream)) | |
| print("Beginning capture on the green-context stream.") | |
| check( | |
| driver.cuStreamBeginCapture( | |
| stream, | |
| driver.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_GLOBAL, | |
| ) | |
| ) | |
| check(driver.cuMemsetD8Async(allocation, 0x7F, CAPTURE_BYTES, stream)) | |
| capture_status = check(driver.cuStreamIsCapturing(stream)) | |
| print(f"Capture status: {capture_status.name}") | |
| graph = check(driver.cuStreamEndCapture(stream)) | |
| print("Ended stream capture and produced a CUDA graph.") | |
| graph_exec = check(driver.cuGraphInstantiate(graph, 0)) | |
| check(driver.cuGraphLaunch(graph_exec, stream)) | |
| check(driver.cuStreamSynchronize(stream)) | |
| print("Instantiated and launched the captured graph successfully.") | |
| finally: | |
| cleanup("graph exec", driver.cuGraphExecDestroy, graph_exec) | |
| cleanup("graph", driver.cuGraphDestroy, graph) | |
| cleanup("allocation", driver.cuMemFree, allocation) | |
| if context is not None: | |
| err = driver.cuCtxSetCurrent(None)[0] | |
| if err != driver.CUresult.CUDA_SUCCESS: | |
| print(f"warning: failed to clear current context: {_driver_error(err)}") | |
| cleanup("green context", driver.cuGreenCtxDestroy, green_ctx) | |
| if __name__ == "__main__": | |
| main() |
Author
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Result from execution ^
Used cuda-bindings 13.3.1