| 91 | } |
| 92 | |
| 93 | void SoftmaxForwardImpl::exec( |
| 94 | _megdnn_tensor_in src, _megdnn_tensor_out dst, _megdnn_workspace workspace) { |
| 95 | dt_float32 alpha = 1.0f, beta = 0.0f; |
| 96 | TensorDesc src_desc, dst_desc; |
| 97 | |
| 98 | cudnnSoftmaxMode_t mode; |
| 99 | std::vector<int> tensor_dims = init_mode(src, mode); |
| 100 | const int dimA[] = { |
| 101 | sc(tensor_dims[0]), sc(tensor_dims[1]), sc(tensor_dims[2]), |
| 102 | sc(tensor_dims[3])}; |
| 103 | const int strideA[] = { |
| 104 | sc(tensor_dims[1] * tensor_dims[2] * tensor_dims[3]), |
| 105 | sc(tensor_dims[2] * tensor_dims[3]), sc(tensor_dims[3]), 1}; |
| 106 | |
| 107 | cudnn_check(cudnnSetTensorNdDescriptor( |
| 108 | src_desc.desc, to_cudnn_dtype(src.layout.dtype), 4, dimA, strideA)); |
| 109 | cudnn_check(cudnnSetTensorNdDescriptor( |
| 110 | dst_desc.desc, to_cudnn_dtype(dst.layout.dtype), 4, dimA, strideA)); |
| 111 | |
| 112 | cudnn_check(cudnnSoftmaxForward( |
| 113 | cudnn_handle(this->handle()), CUDNN_SOFTMAX_ACCURATE, mode, &alpha, |
| 114 | src_desc.desc, src.raw_ptr(), &beta, dst_desc.desc, dst.raw_ptr())); |
| 115 | } |
| 116 | |
| 117 | //================================Softmax Backward============================ |
| 118 |
nothing calls this directly
no test coverage detected