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

Function run_tensor_add

dnn/test/cuda/elemwise.cpp:36–68  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

34}
35
36void run_tensor_add(
37 Handle* handle_cuda, const TensorND& a, const TensorND& b, const TensorND& c) {
38#if 1
39 cudnnHandle_t cudnn_handle;
40 cudnn_check(cudnnCreate(&cudnn_handle));
41 cuda_check(cudaDeviceSynchronize());
42 cuda_check(cudaMemcpy(
43 c.raw_ptr(), a.raw_ptr(), a.layout.span().dist_byte(),
44 cudaMemcpyDeviceToDevice));
45
46 auto bdesc = make_cudnn_tensor_desc(b.layout),
47 cdesc = make_cudnn_tensor_desc(c.layout);
48
49 float alpha = 1, beta = 1;
50 cudaProfilerStart();
51 cudnn_check(cudnnAddTensor(
52 cudnn_handle, &alpha, bdesc, b.raw_ptr(), &beta, cdesc, c.raw_ptr()));
53 cudaProfilerStop();
54
55 cudnn_check(cudnnDestroyTensorDescriptor(cdesc));
56 cudnn_check(cudnnDestroyTensorDescriptor(bdesc));
57 cudnn_check(cudnnDestroy(cudnn_handle));
58
59 cuda_check(cudaMemset(c.raw_ptr(), 0, c.layout.span().dist_byte()));
60 cuda_check(cudaDeviceSynchronize());
61#endif
62
63 auto opr = handle_cuda->create_operator<ElemwiseForward>();
64 opr->param().mode = ElemwiseForward::Mode::ADD;
65 cudaProfilerStart();
66 opr->exec({a, b}, c);
67 cudaProfilerStop();
68}
69
70} // anonymous namespace
71

Callers 1

TEST_FFunction · 0.70

Calls 6

make_cudnn_tensor_descFunction · 0.85
dist_byteMethod · 0.80
spanMethod · 0.80
raw_ptrMethod · 0.45
paramMethod · 0.45
execMethod · 0.45

Tested by

no test coverage detected