| 48 | } |
| 49 | |
| 50 | TensorND NonZeroImpl::exec( |
| 51 | _megdnn_tensor_in src, _megdnn_workspace workspace, |
| 52 | DynOutMallocPolicyCall malloc_policy) { |
| 53 | #if !MGE_BUILD_WITHOUT_NAIVE_EXEC |
| 54 | auto idx_tmp = workspace.ptr<dt_int32>(); |
| 55 | |
| 56 | switch (src.layout.dtype.enumv()) { |
| 57 | #define cb(_dt) \ |
| 58 | case DTypeTrait<_dt>::enumv: { \ |
| 59 | using ctype = DTypeTrait<_dt>::ctype; \ |
| 60 | using namespace cond_take; \ |
| 61 | KParam param({}); \ |
| 62 | param.val = 0.0; \ |
| 63 | param.eps = 1e-6; \ |
| 64 | Pred<PEnum::NEQ, ctype> pred(param); \ |
| 65 | MEGDNN_DISPATCH_CPU_KERN_OPR(gen_index(idx_tmp, src, pred)); \ |
| 66 | break; \ |
| 67 | } |
| 68 | MEGDNN_FOREACH_COMPUTING_DTYPE(cb) |
| 69 | cb(::megdnn::dtype::Bool) |
| 70 | #undef cb |
| 71 | default : { |
| 72 | std::string data_type = src.layout.dtype.name(); |
| 73 | megdnn_throw( |
| 74 | "bad mask dtype,support_types is support types: [Float32, Float16, " |
| 75 | "BFloat16, Int32, Int16, Int8, Uint8, Bool]" + |
| 76 | std::string("but the data type is ") + data_type); |
| 77 | } |
| 78 | } |
| 79 | |
| 80 | static_cast<HandleImpl*>(handle())->megcore_dispatcher()->sync(); |
| 81 | size_t index_size_pos = src.layout.total_nr_elems(); |
| 82 | size_t index_size = idx_tmp[index_size_pos]; |
| 83 | TensorND ret; |
| 84 | size_t ndim = src.layout.ndim; |
| 85 | TensorShape dst_shape({ndim, index_size}); |
| 86 | ret = malloc_policy.alloc_output(0, dtype::Int32(), {ndim, index_size}); |
| 87 | MEGDNN_DISPATCH_CPU_KERN_OPR( |
| 88 | expansion_index(idx_tmp, index_size, &ret, src.layout.shape, ndim)); |
| 89 | |
| 90 | return ret; |
| 91 | #else |
| 92 | __builtin_trap(); |
| 93 | return {}; |
| 94 | #endif |
| 95 | } |
nothing calls this directly
no test coverage detected