| 6410 | } |
| 6411 | |
| 6412 | bool CudnnSupport::DeriveOutputBatchDescriptor( |
| 6413 | const dnn::BatchDescriptor& batch_descriptor, |
| 6414 | const dnn::FilterDescriptor& filter_descriptor, |
| 6415 | const dnn::ConvolutionDescriptor& convolution_descriptor, |
| 6416 | dnn::BatchDescriptor* output_batch_descriptor) { |
| 6417 | CudnnTensorDescriptor input_nd(batch_descriptor, CUDNN_DATA_FLOAT); |
| 6418 | CudnnFilterDescriptor filter(filter_descriptor, CUDNN_DATA_FLOAT); |
| 6419 | CudnnConvolutionDescriptor conv(convolution_descriptor, CUDNN_DATA_FLOAT); |
| 6420 | |
| 6421 | int dn = batch_descriptor.ndims() + 2; |
| 6422 | std::vector<int> dims(dn); // in BDYX |
| 6423 | const auto status = [&] { |
| 6424 | RETURN_IF_CUDNN_ERROR(cudnnGetConvolutionNdForwardOutputDim( |
| 6425 | conv.handle(), input_nd.handle(), filter.handle(), dn, dims.data())); |
| 6426 | output_batch_descriptor->set_count(dims[0]) |
| 6427 | .set_feature_map_count(dims[1]) |
| 6428 | .set_layout(batch_descriptor.layout()); |
| 6429 | |
| 6430 | for (int i = 0; i < batch_descriptor.ndims(); i++) { |
| 6431 | output_batch_descriptor->set_spatial_dim(static_cast<dnn::DimIndex>(i), |
| 6432 | dims.rbegin()[i]); |
| 6433 | } |
| 6434 | return port::Status::OK(); |
| 6435 | }(); |
| 6436 | return IsStatusOk(status, /*report_error=*/true); |
| 6437 | } |
| 6438 | |
| 6439 | } // namespace gpu |
| 6440 | |