nvjitlink_cuda_jit_linker_python_bindings_quickstart.py

1import torch
2from cuda import nvjitlink
3
4# This example demonstrates the basic workflow of the nvjitlink Python API
5# It initializes a linker, adds a dummy input, and retrieves the linked result.
6
7def run_quickstart():
8    # 1. Create a linker handle
9    # In a real scenario, you would pass options like '-arch=sm_xx'
10    handle = nvjitlink.nvJitLinkCreate([])
11
12    try:
13        # 2. Define a simple cubin or PTX input as bytes
14        # Note: In a real application, this would be valid CUDA binary data.
15        # Here we use a placeholder to illustrate the API call structure.
16        input_data = b"\x7fELF..." 
17
18        # 3. Add the data to the linker
19        # Input type 1 corresponds to NVJITLINK_INPUT_CUBIN
20        nvjitlink.nvJitLinkAddData(handle, 1, input_data, "my_module")
21
22        # 4. Perform the link
23        nvjitlink.nvJitLinkComplete(handle)
24
25        # 5. Retrieve the linked image size and data
26        size = nvjitlink.nvJitLinkGetLinkedImageSize(handle)
27        linked_image = nvjitlink.nvJitLinkGetLinkedImage(handle)
28
29        print(f"Linked image size: {size} bytes")
30
31    finally:
32        # 6. Always destroy the handle to free resources
33        nvjitlink.nvJitLinkDestroy(handle)
34
35if __name__ == "__main__":
36    run_quickstart()