| 6 | using namespace cambricon; |
| 7 | |
| 8 | void IndexingOneHotForwardImpl::exec( |
| 9 | _megdnn_tensor_in src, _megdnn_tensor_in index, _megdnn_tensor_out dst, |
| 10 | _megdnn_workspace workspace) { |
| 11 | check_exec(src.layout, index.layout, dst.layout, workspace.size); |
| 12 | auto dtype_dest = dst.layout.dtype.enumv(); |
| 13 | megdnn_assert( |
| 14 | check_dtype_float_ieee(dtype_dest) || |
| 15 | dtype_dest == megdnn::DTypeEnum::Int32, |
| 16 | "Cambricon unsupport IndexingOneHot with dtype:%d", |
| 17 | static_cast<int>(dtype_dest)); |
| 18 | auto cnnl_handler = cnnl_handle(this->handle()); |
| 19 | auto axis = param().axis; |
| 20 | TensorLayout index_mid_layout = index.layout; |
| 21 | index_mid_layout.add_axis_cont_inplace(axis); |
| 22 | CnnlTensorDescriptor src_desc, index_desc, dst_desc; |
| 23 | src_desc.set(src.layout); |
| 24 | index_desc.set(index_mid_layout); |
| 25 | dst_desc.set(dst.layout); |
| 26 | cnnl_check(cnnlGather( |
| 27 | cnnl_handler, axis, src_desc.desc(), src.raw_ptr(), index_desc.desc(), |
| 28 | index.raw_ptr(), dst_desc.desc(), dst.raw_ptr())); |
| 29 | } |
| 30 | |
| 31 | void IndexingSetOneHotForwardImpl::exec( |
| 32 | _megdnn_tensor_inout data, _megdnn_tensor_in index, _megdnn_tensor_in sub, |
nothing calls this directly
no test coverage detected