MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / make_cudnn_tensor_desc

Function make_cudnn_tensor_desc

dnn/test/cuda/elemwise.cpp:18–34  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

16#define cudnn_check(e) megdnn_assert((e) == CUDNN_STATUS_SUCCESS)
17namespace {
18__attribute__((unused)) cudnnTensorDescriptor_t make_cudnn_tensor_desc(
19 const TensorLayout& ly) {
20 megdnn_assert(ly.ndim && ly.ndim <= 4 && ly.is_contiguous());
21 int dim[4] = {1, 1, 1, 1}, stride[4] = {1, 1, 1, 1};
22 for (size_t i = 0; i < ly.ndim; ++i) {
23 dim[i] = ly.shape[i];
24 stride[i] = ly.stride[i];
25 }
26 cudnnTensorDescriptor_t ret;
27 cudnn_check(cudnnCreateTensorDescriptor(&ret));
28 // cudnn requires tensors to be at-least 4D
29 cudnn_check(cudnnSetTensor4dDescriptorEx(
30 ret, CUDNN_DATA_FLOAT, dim[0], dim[1], dim[2], dim[3], stride[0], stride[1],
31 stride[2], stride[3]));
32
33 return ret;
34}
35
36void run_tensor_add(
37 Handle* handle_cuda, const TensorND& a, const TensorND& b, const TensorND& c) {

Callers 1

run_tensor_addFunction · 0.85

Calls 1

is_contiguousMethod · 0.45

Tested by

no test coverage detected