Repository navigation
Expand file tree
/
Copy pathsm_carveout.cpp
More file actions
138 lines (118 loc) · 6.3 KB
/
Copy pathsm_carveout.cpp
File metadata and controls
138 lines (118 loc) · 6.3 KB
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
/*
* SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: MIT
*/
#include <catch2/catch_test_macros.hpp>
#include "../utils/helpers.h"
#include <cudnn_frontend.h>
TEST_CASE("SGBN with SM carveout", "[batchnorm][graph][sm_carveout]") {
if (cudnnGetVersion() < 90300) {
SKIP("SM carveout on batchnorm not supported pre-cudnn-9.3.0");
}
namespace fe = cudnn_frontend;
fe::graph::Graph graph;
graph.set_io_data_type(fe::DataType_t::HALF)
.set_intermediate_data_type(fe::DataType_t::FLOAT)
.set_compute_data_type(fe::DataType_t::FLOAT)
.set_sm_count(8);
auto n = 8, c = 32, h = 16, w = 16;
auto X = graph.tensor(
fe::graph::Tensor_attributes().set_name("X").set_dim({n, c, h, w}).set_stride({c * h * w, 1, c * w, c}));
auto prev_running_mean = graph.tensor(fe::graph::Tensor_attributes()
.set_name("prev_running_mean")
.set_dim({1, c, 1, 1})
.set_stride({c, 1, c, c})
.set_data_type(fe::DataType_t::FLOAT));
auto prev_running_var = graph.tensor(fe::graph::Tensor_attributes()
.set_name("prev_running_var")
.set_dim({1, c, 1, 1})
.set_stride({c, 1, c, c})
.set_data_type(fe::DataType_t::FLOAT));
auto scale = graph.tensor(fe::graph::Tensor_attributes()
.set_name("scale")
.set_dim({1, c, 1, 1})
.set_stride({c, 1, c, c})
.set_data_type(fe::DataType_t::FLOAT));
auto bias = graph.tensor(fe::graph::Tensor_attributes()
.set_name("bias")
.set_dim({1, c, 1, 1})
.set_stride({c, 1, c, c})
.set_data_type(fe::DataType_t::FLOAT));
auto peer_stats_0 = graph.tensor(fe::graph::Tensor_attributes()
.set_dim({2, 4 * c, 1, 1})
.set_stride({4 * c, 1, 4 * c, 4 * c})
.set_data_type(fe::DataType_t::FLOAT));
auto peer_stats_1 = graph.tensor(fe::graph::Tensor_attributes()
.set_dim({2, 4 * c, 1, 1})
.set_stride({4 * c, 1, 4 * c, 4 * c})
.set_data_type(fe::DataType_t::FLOAT));
float epsilon_cpu = 1e-05f;
float momentum_cpu = 1e-01f;
auto epsilon = graph.tensor(epsilon_cpu);
auto momentum = graph.tensor(momentum_cpu);
auto batchnorm_options = fe::graph::Batchnorm_attributes()
.set_epsilon(epsilon)
.set_previous_running_stats(prev_running_mean, prev_running_var, momentum)
.set_peer_stats({peer_stats_0, peer_stats_1});
auto [Y, mean, inv_variance, next_running_mean, next_running_var] =
graph.batchnorm(X, scale, bias, batchnorm_options);
mean->set_output(true).set_data_type(fe::DataType_t::FLOAT);
inv_variance->set_output(true).set_data_type(fe::DataType_t::FLOAT);
next_running_mean->set_output(true).set_data_type(fe::DataType_t::FLOAT);
next_running_var->set_output(true).set_data_type(fe::DataType_t::FLOAT);
Y->set_output(true);
#if (CUDNN_VERSION < 8700)
SKIP("single GPU BN is not supported in cudnn versions prior to 8.7");
#endif
if (check_device_arch_newer_than("ampere") == false) {
SKIP("ConvBNFprop requires Ampere and up");
}
// Create a unique_ptr for the cuDNN handle
auto handle_ptr = create_cudnn_handle();
auto handle = *handle_ptr;
REQUIRE(graph.validate().is_good());
REQUIRE(graph.build_operation_graph(handle).is_good());
REQUIRE(graph.create_execution_plans({fe::HeurMode_t::A}).is_good());
REQUIRE(graph.check_support(handle).is_good());
REQUIRE(graph.build_plans(handle).is_good());
Surface<half> X_tensor(n * c * h * w);
Surface<float> Mean_tensor(c);
Surface<float> Var_tensor(c);
Surface<float> Previous_running_mean_tensor(c);
Surface<float> Previous_running_var_tensor(c);
Surface<float> Next_running_mean_tensor(c);
Surface<float> Next_running_var_tensor(c);
Surface<float> Scale_tensor(c);
Surface<float> Bias_tensor(c);
Surface<half> Y_tensor(n * c * h * w);
Surface<float> Peer_stats_0_tensor(2 * 4 * c);
std::vector<float> peer_stats_0_host(Peer_stats_0_tensor.size);
initHostImage(peer_stats_0_host.data(), static_cast<int64_t>(peer_stats_0_host.size()));
auto* peer_stats_0_bits = reinterpret_cast<uint32_t*>(peer_stats_0_host.data());
for (size_t i = 0; i < peer_stats_0_host.size(); i += 2) {
peer_stats_0_bits[i + 1] = 1u;
}
CUDA_CHECK(cudaMemcpy(Peer_stats_0_tensor.devPtr,
peer_stats_0_host.data(),
sizeof(peer_stats_0_host[0]) * peer_stats_0_host.size(),
cudaMemcpyHostToDevice));
CUDA_CHECK(cudaDeviceSynchronize());
Surface<float> Peer_stats_1_tensor(2 * 4 * c);
int64_t workspace_size = 0;
REQUIRE(graph.get_workspace_size(workspace_size).is_good());
Surface<int8_t> workspace(workspace_size);
std::unordered_map<std::shared_ptr<fe::graph::Tensor_attributes>, void*> variant_pack = {
{X, X_tensor.devPtr},
{mean, Mean_tensor.devPtr},
{inv_variance, Var_tensor.devPtr},
{prev_running_mean, Previous_running_mean_tensor.devPtr},
{prev_running_var, Previous_running_var_tensor.devPtr},
{next_running_mean, Next_running_mean_tensor.devPtr},
{next_running_var, Next_running_var_tensor.devPtr},
{scale, Scale_tensor.devPtr},
{bias, Bias_tensor.devPtr},
{Y, Y_tensor.devPtr},
{peer_stats_0, Peer_stats_0_tensor.devPtr},
{peer_stats_1, Peer_stats_1_tensor.devPtr}};
REQUIRE(graph.execute(handle, variant_pack, workspace.devPtr).is_good());
}