| 88 | } |
| 89 | |
| 90 | void TopKImpl::do_exec( |
| 91 | int k, _megdnn_tensor_in data, _megdnn_tensor_out values, int32_t* indices, |
| 92 | _megdnn_workspace workspace) { |
| 93 | megdnn_assert( |
| 94 | param().mode != Param::Mode::KTH_ONLY, |
| 95 | "CNNL topk do not support KTH_ONLY mode"); |
| 96 | auto _handle = cnnl_handle(this->handle()); |
| 97 | TensorLayout indices_layout = values.layout; |
| 98 | indices_layout.dtype = dtype::Int32{}; |
| 99 | auto ws_bundle = make_bundle(k, data.layout, values.layout, indices_layout); |
| 100 | ws_bundle.set(workspace.raw_ptr); |
| 101 | TopKCnnlDescs descs(data.layout, values.layout, param().mode); |
| 102 | |
| 103 | void* target_src = data.raw_ptr(); |
| 104 | if (!data.layout.is_contiguous()) { |
| 105 | TensorLayout dst; |
| 106 | dst = data.layout; |
| 107 | dst.init_contiguous_stride(); |
| 108 | CnnlTensorDescriptor dst_desc; |
| 109 | dst_desc.set(dst); |
| 110 | cnnl_check(cnnlCopy( |
| 111 | _handle, descs.data_desc.desc(), data.raw_ptr(), dst_desc.desc(), |
| 112 | ws_bundle.get(0))); |
| 113 | target_src = ws_bundle.get(0); |
| 114 | } |
| 115 | switch (data.layout.dtype.enumv()) { |
| 116 | #define cb(t) \ |
| 117 | case DTypeTrait<t>::enumv: \ |
| 118 | do { \ |
| 119 | using ct = DTypeTrait<t>::ctype; \ |
| 120 | dispatch_with_ctype<ct>( \ |
| 121 | k, static_cast<ct*>(target_src), values.ptr<ct>(), indices, \ |
| 122 | ws_bundle.get(1), ws_bundle.get_size(1), descs); \ |
| 123 | return; \ |
| 124 | } while (0); |
| 125 | cb(::megdnn::dtype::Float32); |
| 126 | DNN_INC_FLOAT16(cb(::megdnn::dtype::Float16)); |
| 127 | MEGDNN_FOREACH_COMPUTING_DTYPE_INT(cb); |
| 128 | default: |
| 129 | megdnn_throw(ssprintf( |
| 130 | "cambricon topk not support dtype=%s", data.layout.dtype.name())); |
| 131 | #undef cb |
| 132 | } |
| 133 | } |
| 134 | |
| 135 | } // namespace cambricon |
| 136 | } // namespace megdnn |
nothing calls this directly
no test coverage detected