| 31 | } |
| 32 | |
| 33 | TensorND NonZeroImpl::exec( |
| 34 | _megdnn_tensor_in data, _megdnn_workspace workspace, |
| 35 | DynOutMallocPolicyCall malloc_policy) { |
| 36 | size_t size = data.layout.total_nr_elems(); |
| 37 | if (size == 0) { |
| 38 | TensorShape target_shape({data.layout.ndim, 0}); |
| 39 | TensorND rst = malloc_policy.alloc_output(0, dtype::Int32(), target_shape); |
| 40 | return rst; |
| 41 | } |
| 42 | auto wk_bundle = make_bundle(data.layout); |
| 43 | wk_bundle.set(workspace.raw_ptr); |
| 44 | |
| 45 | auto idx_tmp = static_cast<megdnn::cuda::cond_take::IdxType*>(wk_bundle.get(0)); |
| 46 | |
| 47 | CondTake::Param param; |
| 48 | param.mode = CondTake::Param::Mode::NEQ; |
| 49 | param.val = 0; |
| 50 | param.eps = 1e-06; |
| 51 | megdnn::cond_take::KParam kparam(param); |
| 52 | |
| 53 | auto stream = cuda_stream(handle()); |
| 54 | size_t out_size; |
| 55 | switch (data.layout.dtype.enumv()) { |
| 56 | #define cb(_dt) \ |
| 57 | case DTypeTrait<_dt>::enumv: { \ |
| 58 | using ctype = DTypeTrait<_dt>::ctype; \ |
| 59 | out_size = megdnn::cuda::cond_take::gen_idx( \ |
| 60 | wk_bundle.get(1), wk_bundle.get_size(1), idx_tmp, data.ptr<ctype>(), \ |
| 61 | size, static_cast<uint32_t>(param.mode), kparam, stream); \ |
| 62 | break; \ |
| 63 | } |
| 64 | MEGDNN_FOREACH_COMPUTING_DTYPE(cb) |
| 65 | cb(::megdnn::dtype::Bool) |
| 66 | #undef cb |
| 67 | default : { |
| 68 | std::string data_type = data.layout.dtype.name(); |
| 69 | megdnn_throw( |
| 70 | "bad mask dtype,support_types is support types: [Float32, Float16, " |
| 71 | "BFloat16, Int32, Int16, Int8, Uint8, Bool]" + |
| 72 | std::string("but the data type is ") + data_type); |
| 73 | } |
| 74 | } |
| 75 | |
| 76 | TensorShape dst_shape({data.layout.ndim, out_size}); |
| 77 | TensorND out_idx = malloc_policy.alloc_output(0, dtype::Int32(), dst_shape); |
| 78 | dt_int32* out_idx_ptr = out_idx.ptr<dt_int32>(); |
| 79 | |
| 80 | switch (data.layout.dtype.enumv()) { |
| 81 | #define cb(_dt) \ |
| 82 | case DTypeTrait<_dt>::enumv: { \ |
| 83 | for (size_t idx = 0; idx < data.layout.ndim; idx++) { \ |
| 84 | dt_int32* copy_idx_ptr = out_idx_ptr + idx * out_size; \ |
| 85 | copy_idx(copy_idx_ptr, idx_tmp, size, stream); \ |
| 86 | } \ |
| 87 | break; \ |
| 88 | } |
| 89 | MEGDNN_FOREACH_COMPUTING_DTYPE(cb) |
| 90 | cb(::megdnn::dtype::Bool) |
nothing calls this directly
no test coverage detected