Method set

dnn/src/cuda/cudnn_wrapper.cpp:109–150

}

void TensorDesc::set(
        const TensorLayout& layout, const param::Convolution::Format format) {
    // Layout can be not contiguous; group conv needs it.
    // megdnn_assert_contiguous(layout);
    if (format == param::Convolution::Format::NCHW4 ||
        format == param::Convolution::Format::NCHW32)
        megdnn_assert_eq_size_t(layout.ndim, 5_z);
    else
        megdnn_assert_eq_size_t(layout.ndim, 4_z);

    size_t c_pos, spatial_pos;
    if (format == param::Convolution::Format::NCHW ||
        format == param::Convolution::Format::NCHW4 ||
        format == param::Convolution::Format::NCHW32) {
        c_pos = 1;
        spatial_pos = 2;
    } else {
        megdnn_assert(format == param::Convolution::Format::NHWC);
        c_pos = 3;
        spatial_pos = 1;
    }
    if (format == param::Convolution::Format::NCHW4) {
        megdnn_assert(layout.is_physical_contiguous());
        cudnn_check(cudnnSetTensor4dDescriptor(
                desc, to_cudnn_format(format), to_cudnn_dtype(layout.dtype, format),
                layout.shape[0], layout.shape[c_pos] * 4, layout.shape[spatial_pos + 0],
                layout.shape[spatial_pos + 1]));
    } else if (format == param::Convolution::Format::NCHW32) {
        megdnn_assert(layout.is_physical_contiguous());
        cudnn_check(cudnnSetTensor4dDescriptor(
                desc, to_cudnn_format(format), to_cudnn_dtype(layout.dtype, format),
                layout.shape[0], layout.shape[c_pos] * 32,
                layout.shape[spatial_pos + 0], layout.shape[spatial_pos + 1]));

    } else {
        cudnn_check(cudnnSetTensor4dDescriptorEx(
                desc, to_cudnn_dtype(layout.dtype), layout.shape[0],
                layout.shape[c_pos], layout.shape[spatial_pos + 0],
                layout.shape[spatial_pos + 1], layout.stride[0], layout.stride[c_pos],
                layout.stride[spatial_pos + 0], layout.stride[spatial_pos + 1]));
    }
}

std::string TensorDesc::to_string() {
    cudnnDataType_t data_type;