Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
bn_finalize.py 6.94 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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186 
import hipdnn
import torch


def build_bn_finalize_graph(
    hipdnn_handle,
    torch_tensor_sum,
    torch_tensor_sq_sum,
    torch_tensor_scale,
    torch_tensor_bias,
    torch_tensor_prev_running_mean,
    torch_tensor_prev_running_variance,
    torch_tensor_momentum,
    torch_tensor_epsilon,
    torch_tensor_accum_count,
    hipdnn_data_type,
):
    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="bn_finalize",
    )

    hipdnn_tensor_sum = graph.tensor_like(torch_tensor_sum)
    hipdnn_tensor_sq_sum = graph.tensor_like(torch_tensor_sq_sum)
    hipdnn_tensor_scale = graph.tensor_like(torch_tensor_scale)
    hipdnn_tensor_bias = graph.tensor_like(torch_tensor_bias)
    hipdnn_tensor_prev_running_mean = graph.tensor_like(torch_tensor_prev_running_mean)
    hipdnn_tensor_prev_running_variance = graph.tensor_like(torch_tensor_prev_running_variance)
    hipdnn_tensor_momentum = graph.tensor_like(torch_tensor_momentum)
    hipdnn_tensor_momentum.set_value(0.001)
    hipdnn_tensor_epsilon = graph.tensor_like(torch_tensor_epsilon)
    hipdnn_tensor_epsilon.set_value(1e-5)
    hipdnn_tensor_accum_count = graph.tensor_like(torch_tensor_accum_count)
    hipdnn_tensor_accum_count.set_value(torch_tensor_accum_count.item())

    (
        hipdnn_tensor_eq_scale,
        hipdnn_tensor_eq_bias,
        hipdnn_tensor_mean,
        hipdnn_tensor_inv_variance,
        hipdnn_tensor_next_running_mean,
        hipdnn_tensor_next_running_variance,
    ) = graph.bn_finalize(
        sum=hipdnn_tensor_sum,
        sq_sum=hipdnn_tensor_sq_sum,
        scale=hipdnn_tensor_scale,
        bias=hipdnn_tensor_bias,
        epsilon=hipdnn_tensor_epsilon,
        accum_count=hipdnn_tensor_accum_count,
        prev_running_mean=hipdnn_tensor_prev_running_mean,
        prev_running_variance=hipdnn_tensor_prev_running_variance,
        momentum=hipdnn_tensor_momentum,
        name="bn_finalize_node",
    )
    hipdnn_tensor_eq_scale.set_output(True)
    hipdnn_tensor_eq_bias.set_output(True)
    hipdnn_tensor_mean.set_output(True)
    hipdnn_tensor_inv_variance.set_output(True)
    hipdnn_tensor_next_running_mean.set_output(True)
    hipdnn_tensor_next_running_variance.set_output(True)

    graph.build(hipdnn_handle)

    return (
        graph,
        hipdnn_tensor_sum,
        hipdnn_tensor_sq_sum,
        hipdnn_tensor_scale,
        hipdnn_tensor_bias,
        hipdnn_tensor_prev_running_mean,
        hipdnn_tensor_prev_running_variance,
        hipdnn_tensor_momentum,
        hipdnn_tensor_epsilon,
        hipdnn_tensor_accum_count,
        hipdnn_tensor_eq_scale,
        hipdnn_tensor_eq_bias,
        hipdnn_tensor_mean,
        hipdnn_tensor_inv_variance,
        hipdnn_tensor_next_running_mean,
        hipdnn_tensor_next_running_variance,
    )


if __name__ == "__main__":
    n = 1
    c = 32
    h = 1
    w = 1

    hipdnn_data_type = hipdnn.data_type.FLOAT
    torch_data_type = torch.float32

    torch_tensor_sum = torch.rand(n, c, h, w, dtype=torch_data_type, device="cuda")
    torch_tensor_sq_sum = torch.rand(n, c, h, w, dtype=torch_data_type, device="cuda")
    torch_tensor_scale = torch.rand(n, c, h, w, dtype=torch_data_type, device="cuda")
    torch_tensor_bias = torch.rand(n, c, h, w, dtype=torch_data_type, device="cuda")
    torch_tensor_prev_running_mean = torch.rand(n, c, h, w, dtype=torch_data_type, device="cuda")
    torch_tensor_prev_running_variance = torch.rand(
        n, c, h, w, dtype=torch_data_type, device="cuda"
    )
    torch_tensor_momentum = torch.full(
        (1, 1, 1, 1), 0.001, dtype=torch.float32, requires_grad=False, device="cuda"
    )
    torch_tensor_epsilon = torch.full(
        (1, 1, 1, 1), 1e-5, dtype=torch.float32, requires_grad=False, device="cuda"
    )
    torch_tensor_accum_count = torch.full(
        (1, 1, 1, 1), n * h * w, dtype=torch.int32, requires_grad=False, device="cuda"
    )

    hipdnn_handle = hipdnn.create_handle()

    (
        graph,
        hipdnn_tensor_sum,
        hipdnn_tensor_sq_sum,
        hipdnn_tensor_scale,
        hipdnn_tensor_bias,
        hipdnn_tensor_prev_running_mean,
        hipdnn_tensor_prev_running_variance,
        hipdnn_tensor_momentum,
        hipdnn_tensor_epsilon,
        hipdnn_tensor_accum_count,
        hipdnn_tensor_eq_scale,
        hipdnn_tensor_eq_bias,
        hipdnn_tensor_mean,
        hipdnn_tensor_inv_variance,
        hipdnn_tensor_next_running_mean,
        hipdnn_tensor_next_running_variance,
    ) = build_bn_finalize_graph(
        hipdnn_handle,
        torch_tensor_sum,
        torch_tensor_sq_sum,
        torch_tensor_scale,
        torch_tensor_bias,
        torch_tensor_prev_running_mean,
        torch_tensor_prev_running_variance,
        torch_tensor_momentum,
        torch_tensor_epsilon,
        torch_tensor_accum_count,
        hipdnn_data_type,
    )

    torch_tensor_eq_scale = torch.empty(
        hipdnn_tensor_eq_scale.get_dim(), dtype=torch_data_type, device="cuda"
    )
    torch_tensor_eq_bias = torch.empty(
        hipdnn_tensor_eq_bias.get_dim(), dtype=torch_data_type, device="cuda"
    )
    torch_tensor_mean = torch.empty(
        hipdnn_tensor_mean.get_dim(), dtype=torch_data_type, device="cuda"
    )
    torch_tensor_inv_variance = torch.empty(
        hipdnn_tensor_inv_variance.get_dim(), dtype=torch_data_type, device="cuda"
    )
    torch_tensor_next_running_mean = torch.empty(
        hipdnn_tensor_next_running_mean.get_dim(), dtype=torch_data_type, device="cuda"
    )
    torch_tensor_next_running_variance = torch.empty(
        hipdnn_tensor_next_running_variance.get_dim(), dtype=torch_data_type, device="cuda"
    )

    variant_pack = {
        hipdnn_tensor_sum: torch_tensor_sum.data_ptr(),
        hipdnn_tensor_sq_sum: torch_tensor_sq_sum.data_ptr(),
        hipdnn_tensor_scale: torch_tensor_scale.data_ptr(),
        hipdnn_tensor_bias: torch_tensor_bias.data_ptr(),
        hipdnn_tensor_prev_running_mean: torch_tensor_prev_running_mean.data_ptr(),
        hipdnn_tensor_prev_running_variance: torch_tensor_prev_running_variance.data_ptr(),
        hipdnn_tensor_momentum: torch_tensor_momentum.data_ptr(),
        hipdnn_tensor_epsilon: torch_tensor_epsilon.data_ptr(),
        hipdnn_tensor_accum_count: torch_tensor_accum_count.data_ptr(),
        hipdnn_tensor_eq_scale: torch_tensor_eq_scale.data_ptr(),
        hipdnn_tensor_eq_bias: torch_tensor_eq_bias.data_ptr(),
        hipdnn_tensor_mean: torch_tensor_mean.data_ptr(),
        hipdnn_tensor_inv_variance: torch_tensor_inv_variance.data_ptr(),
        hipdnn_tensor_next_running_mean: torch_tensor_next_running_mean.data_ptr(),
        hipdnn_tensor_next_running_variance: torch_tensor_next_running_variance.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("Batch normalization finalize graph execution complete.")