Repository navigation
Expand file tree
/
Copy pathfp8_fprop.cpp
More file actions
118 lines (93 loc) · 4.9 KB
/
Copy pathfp8_fprop.cpp
File metadata and controls
118 lines (93 loc) · 4.9 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
/*
* SPDX-FileCopyrightText: Copyright (c) 2024 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("Convolution fp8 precision", "[conv][graph]") {
if (cudnnGetCudartVersion() < 12000) {
SKIP("Test requires cuda toolkit 12.0 or above");
}
if (cudnnGetVersion() < 8600) {
SKIP("TEST REQUIRES minimum cudnn version 8.6.0");
}
if (check_device_arch_newer_than("hopper") == false) {
SKIP("TEST REQUIRES device hopper arch or newer");
}
namespace fe = cudnn_frontend;
// conv problem size
int64_t n = 16, c = 128, h = 64, w = 64, k = 256, r = 1, s = 1;
// Initialize input tensors with int8_t as proxy for fp8
auto graph = std::make_shared<fe::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);
auto X = graph->tensor(fe::graph::Tensor_attributes()
.set_name("image")
.set_dim({n, c, h, w})
.set_stride({c * h * w, 1, c * w, c})
.set_data_type(fe::DataType_t::FP8_E4M3));
auto W = graph->tensor(fe::graph::Tensor_attributes()
.set_name("filter")
.set_dim({k, c, r, s})
.set_stride({c * r * s, 1, c * s, c})
.set_data_type(fe::DataType_t::FP8_E4M3));
auto conv_options =
fe::graph::Conv_fprop_attributes().set_padding({0, 0}).set_stride({1, 1}).set_dilation({1, 1}).set_name("conv");
auto conv_output_fp8 = graph->conv_fprop(X, W, conv_options);
auto descale_x = graph->tensor(fe::graph::Tensor_attributes()
.set_name("descale_x")
.set_dim({1, 1, 1, 1})
.set_stride({1, 1, 1, 1})
.set_data_type(fe::DataType_t::FLOAT));
auto descale_w = graph->tensor(fe::graph::Tensor_attributes()
.set_name("descale_w")
.set_dim({1, 1, 1, 1})
.set_stride({1, 1, 1, 1})
.set_data_type(fe::DataType_t::FLOAT));
auto scale_y = graph->tensor(fe::graph::Tensor_attributes()
.set_name("scale_y")
.set_dim({1, 1, 1, 1})
.set_stride({1, 1, 1, 1})
.set_data_type(fe::DataType_t::FLOAT));
auto scale_options = fe::graph::Pointwise_attributes().set_mode(fe::PointwiseMode_t::MUL);
auto after_descale_x = graph->pointwise(conv_output_fp8, descale_x, scale_options);
auto after_descale_w = graph->pointwise(after_descale_x, descale_w, scale_options);
auto Y = graph->pointwise(after_descale_w, scale_y, scale_options);
Y->set_output(true).set_data_type(fe::DataType_t::FP8_E4M3);
auto amax = graph->reduction(after_descale_w,
fe::graph::Reduction_attributes()
.set_mode(fe::ReductionMode_t::AMAX)
.set_compute_data_type(fe::DataType_t::FLOAT));
amax->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({1, 1, 1, 1});
REQUIRE(graph->validate().is_good());
// Create a unique_ptr for the cuDNN handle
auto handle_ptr = create_cudnn_handle();
auto handle = *handle_ptr;
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, fe::BuildPlanPolicy_t::HEURISTICS_CHOICE).is_good());
// Use int8_t as proxy for fp8
Surface<int8_t> X_gpu(n * c * h * w);
Surface<int8_t> W_gpu(k * c * r * s);
Surface<int8_t> Y_gpu(n * k * h * w);
Surface<float> X_descale_gpu(1);
Surface<float> W_descale_gpu(1);
Surface<float> Y_scale_gpu(1);
Surface<float> amax_gpu(1);
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_gpu.devPtr},
{W, W_gpu.devPtr},
{Y, Y_gpu.devPtr},
{descale_x, X_descale_gpu.devPtr},
{descale_w, W_descale_gpu.devPtr},
{scale_y, Y_scale_gpu.devPtr},
{amax, amax_gpu.devPtr}};
std::cout << graph->print() << std::endl;
REQUIRE(graph->execute(handle, variant_pack, workspace.devPtr).is_good());
}