Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
conv_bias_prelu.py 3.38 KB
Newer Older
yanjl1's avatar
Initial  
yanjl1 committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119 
import hipdnn
import torch


def build_conv_bias_prelu_graph(
    hipdnn_handle,
    torch_tensor_x,
    torch_tensor_w,
    torch_tensor_bias,
    padding,
    stride,
    dilation,
    negative_slope,
    hipdnn_data_type,
):
    # Create graph
    graph = hipdnn.pygraph(
        handle=hipdnn_handle,
        io_data_type=hipdnn_data_type,
        intermediate_data_type=hipdnn.data_type.FLOAT,
        compute_data_type=hipdnn.data_type.FLOAT,
        name="conv_bias_prelu",
    )

    # Create hipdnn tensors
    hipdnn_tensor_x = graph.tensor_like(torch_tensor_x)
    hipdnn_tensor_w = graph.tensor_like(torch_tensor_w)
    hipdnn_tensor_bias = graph.tensor_like(torch_tensor_bias)

    # Create op
    hipdnn_tensor_conv_output = graph.conv_fprop(
        image=hipdnn_tensor_x,
        weight=hipdnn_tensor_w,
        padding=padding,
        stride=stride,
        dilation=dilation,
        name="conv2d",
    )
    hipdnn_tensor_add_output = graph.add(
        a=hipdnn_tensor_conv_output, b=hipdnn_tensor_bias, name="bias"
    )
    hipdnn_tensor_y = graph.prelu(
        input=hipdnn_tensor_add_output, negative_slope=negative_slope, name="prelu"
    )
    hipdnn_tensor_y.set_output(True)

    graph.build(hipdnn_handle)

    return (graph, hipdnn_tensor_x, hipdnn_tensor_w, hipdnn_tensor_bias, hipdnn_tensor_y)


if __name__ == "__main__":
    # Input dimensions
    n = 1  # Batch size
    c = 16  # Number of input channels
    h = 16  # Height
    w = 16  # Width

    # Filter dimensions
    k = 16  # Number of output channels
    r = 3  # Filter height
    s = 3  # Filter width

    # Convolution parameters
    stride_h = 1  # Height stride
    stride_w = 1  # Width stride
    pad_h = 1  # Height padding
    pad_w = 1  # Width padding
    dil_h = 1  # Height dilation
    dil_w = 1  # Width dilation

    # activate parameters
    negative_slope = 0.01  # Negative slope

    hipdnn_data_type = hipdnn.data_type.HALF
    torch_data_type = torch.float16

    torch_tensor_x = torch.rand(n, c, h, w, dtype=torch_data_type, device="cuda").to(
        memory_format=torch.channels_last
    )
    torch_tensor_w = torch.rand(k, c, r, s, dtype=torch_data_type, device="cuda").to(
        memory_format=torch.channels_last
    )
    torch_tensor_bias = torch.rand(1, k, 1, 1, dtype=torch_data_type, device="cuda").to(
        memory_format=torch.channels_last
    )

    hipdnn_handle = hipdnn.create_handle()

    graph, hipdnn_tensor_x, hipdnn_tensor_w, hipdnn_tensor_bias, hipdnn_tensor_y = (
        build_conv_bias_prelu_graph(
            hipdnn_handle,
            torch_tensor_x,
            torch_tensor_w,
            torch_tensor_bias,
            [pad_h, pad_w],
            [stride_h, stride_w],
            [dil_h, dil_w],
            negative_slope,
            hipdnn_data_type,
        )
    )

    torch_tensor_y = torch.empty(
        hipdnn_tensor_y.get_dim(),
        dtype=torch_data_type,
        memory_format=torch.channels_last,
        device="cuda",
    )
    variant_pack = {
        hipdnn_tensor_x: torch_tensor_x.data_ptr(),
        hipdnn_tensor_w: torch_tensor_w.data_ptr(),
        hipdnn_tensor_bias: torch_tensor_bias.data_ptr(),
        hipdnn_tensor_y: torch_tensor_y.data_ptr(),
    }
    workspace = torch.empty(graph.get_workspace_size(), dtype=torch.uint8, device="cuda")

    graph.exec(variant_pack=variant_pack, workspace=workspace.data_ptr())
    print("conv_bias_prelu graph execution complete.")