| 91 | } |
| 92 | |
| 93 | std::optional<DnnGraph> build_conv_graph( |
| 94 | cu::CommandEncoder& encoder, |
| 95 | ConvBackendType backend_type, |
| 96 | Dtype dtype, |
| 97 | array& x, |
| 98 | array& w, |
| 99 | array& y, |
| 100 | const std::vector<int64_t>& stride, |
| 101 | const std::vector<int64_t>& padding_lo, |
| 102 | const std::vector<int64_t>& padding_hi, |
| 103 | const std::vector<int64_t>& dilation) { |
| 104 | auto compute_dtype = |
| 105 | (dtype == float16 || dtype == bfloat16) ? float32 : dtype; |
| 106 | DnnGraph graph(get_cudnn_handle(encoder.device()), dtype, compute_dtype); |
| 107 | auto x_ = graph.tensor_nchw("X", 'x', x); |
| 108 | auto w_ = graph.tensor_nchw("W", 'w', w); |
| 109 | |
| 110 | auto set_options = [&](auto& options) { |
| 111 | options.set_compute_data_type(dtype_to_cudnn_type(compute_dtype)) |
| 112 | .set_convolution_mode(fe::ConvolutionMode_t::CROSS_CORRELATION) |
| 113 | .set_stride(stride) |
| 114 | .set_pre_padding(padding_lo) |
| 115 | .set_post_padding(padding_hi) |
| 116 | .set_dilation(dilation); |
| 117 | }; |
| 118 | |
| 119 | std::shared_ptr<fe::graph::Tensor_attributes> y_; |
| 120 | if (backend_type == CONV_FORWARD) { |
| 121 | auto options = fe::graph::Conv_fprop_attributes(); |
| 122 | set_options(options); |
| 123 | y_ = graph.conv_fprop(x_, w_, options); |
| 124 | } else if (backend_type == CONV_BACKWARD_INPUT) { |
| 125 | auto options = fe::graph::Conv_dgrad_attributes(); |
| 126 | set_options(options); |
| 127 | y_ = graph.conv_dgrad(x_, w_, options); |
| 128 | } else if (backend_type == CONV_BACKWARD_WEIGHT) { |
| 129 | auto options = fe::graph::Conv_wgrad_attributes(); |
| 130 | set_options(options); |
| 131 | y_ = graph.conv_wgrad(w_, x_, options); |
| 132 | } |
| 133 | graph.tensor_nchw(y_, 'y', y)->set_output(true); |
| 134 | |
| 135 | if (graph.prepare().is_bad()) { |
| 136 | return std::nullopt; |
| 137 | } |
| 138 | graph.deselect_numeric_notes({fe::NumericalNote_t::DOWN_CONVERT_INPUTS}); |
| 139 | if (dtype == float32 && !env::enable_tf32()) { |
| 140 | graph.deselect_numeric_notes({fe::NumericalNote_t::TENSOR_CORE}); |
| 141 | } |
| 142 | CHECK_CUDNN_ERROR(graph.build()); |
| 143 | return graph; |
| 144 | } |
| 145 | |
| 146 | // Transpose from (C_out, H, W, C_in / groups) to (C_in, H, W, C_out / groups). |
| 147 | array group_transpose( |
no test coverage detected