| 4 | namespace megdnn { |
| 5 | namespace cuda { |
| 6 | void MaskedFillImpl::exec( |
| 7 | _megdnn_tensor_in origin, _megdnn_tensor_in index, _megdnn_tensor_out dest) { |
| 8 | check_exec(origin.layout, index.layout, dest.layout); |
| 9 | |
| 10 | megdnn_assert(index.layout.is_contiguous()); |
| 11 | uint32_t mask_stride = TensorLayout(origin.layout, origin.layout.dtype) |
| 12 | .stride[index.layout.ndim - 1]; |
| 13 | ElemwiseOpParamN<1> ele_param; |
| 14 | ele_param[0] = origin; |
| 15 | ele_param.init_from_given_tensor(); |
| 16 | auto stream = cuda_stream(handle()); |
| 17 | |
| 18 | #define cb(DType) \ |
| 19 | if (origin.layout.dtype == DType()) { \ |
| 20 | using T = typename DTypeTrait<DType>::ctype; \ |
| 21 | auto value = static_cast<T>(param().value); \ |
| 22 | run_elemwise<MaskedFillScalarKernOp<T>, T, 1>( \ |
| 23 | ele_param, stream, {dest, index, value, mask_stride}); \ |
| 24 | return; \ |
| 25 | } |
| 26 | MEGDNN_FOREACH_COMPUTING_DTYPE(cb) |
| 27 | cb(::megdnn::dtype::Bool) |
| 28 | #undef cb |
| 29 | } |
| 30 | } // namespace cuda |
| 31 | } // namespace megdnn |
nothing calls this directly
no test coverage detected