| 7 | namespace cuda { |
| 8 | |
| 9 | void DotForwardImpl::exec( |
| 10 | _megdnn_tensor_in A, _megdnn_tensor_in B, _megdnn_tensor_out C, |
| 11 | _megdnn_workspace workspace) { |
| 12 | check_exec(A.layout, B.layout, C.layout, workspace.size); |
| 13 | megdnn_assert(A.layout.dtype.category() == DTypeCategory::FLOAT); |
| 14 | auto handle = cublas_handle(this->handle()); |
| 15 | if (A.layout.dtype == dtype::Float32()) { |
| 16 | cublas_check(cublasSdot( |
| 17 | handle, A.layout.total_nr_elems(), A.ptr<dt_float32>(), |
| 18 | A.layout.stride[0], B.ptr<dt_float32>(), B.layout.stride[0], |
| 19 | C.ptr<dt_float32>())); |
| 20 | } else { |
| 21 | megdnn_assert_internal(A.layout.dtype == dtype::Float16()); |
| 22 | dot::run<dt_float16>( |
| 23 | A.ptr<dt_float16>(), B.ptr<dt_float16>(), C.ptr<dt_float16>(), |
| 24 | workspace.ptr<dt_float32>(), A.layout.total_nr_elems(), |
| 25 | A.layout.stride[0], B.layout.stride[0], cuda_stream(this->handle())); |
| 26 | } |
| 27 | } |
| 28 | |
| 29 | } // namespace cuda |
| 30 | } // namespace megdnn |
nothing calls this directly
no test coverage detected