cudnn_python_frontend_convolution_graph_with_pytorch_tensors.py

1import cudnn
2import torch
3
4# Create a cuDNN graph
5graph = cudnn.pygraph(intermediate_data_type=cudnn.data_type.FLOAT, compute_precision=cudnn.data_type.FLOAT)
6
7# Define input, filter, and output dimensions (NCHW format)
8X = graph.tensor(name="X", name_format="NCHW", data_type=cudnn.data_type.FLOAT, shape=[1, 32, 16, 16])
9W = graph.tensor(name="W", name_format="KCHW", data_type=cudnn.data_type.FLOAT, shape=[64, 32, 3, 3])
10
11# Define the convolution operation
12Y = graph.conv_fprop(image=X, weight=W, padding=[1, 1], stride=[1, 1], dilation=[1, 1])
13Y.set_output(True).set_data_type(cudnn.data_type.FLOAT)
14
15# Build the graph for a specific backend (CUDA)
16graph.build([cudnn.backend_behavior.QUERY, cudnn.backend_behavior.COMPILED])
17
18# Prepare data using PyTorch
19x_gpu = torch.randn(1, 32, 16, 16, device="cuda", dtype=torch.float32)
20w_gpu = torch.randn(64, 32, 3, 3, device="cuda", dtype=torch.float32)
21y_gpu = torch.empty(1, 64, 16, 16, device="cuda", dtype=torch.float32)
22
23# Execute the graph
24workspace = torch.empty(graph.get_workspace_size(), device="cuda", dtype=torch.uint8)
25graph.execute(inputs={"X": x_gpu, "W": w_gpu}, outputs={"Y": y_gpu}, workspace=workspace)
26
27print("Convolution execution complete. Output shape:", y_gpu.shape)