| 30 | } |
| 31 | |
| 32 | void TopK::exec( |
| 33 | int k, _megdnn_tensor_in data, _megdnn_tensor_out values, |
| 34 | _megdnn_tensor_out indices, _megdnn_workspace workspace) { |
| 35 | TensorLayout oval, oidx; |
| 36 | deduce_layout(k, data.layout, oval, oidx); |
| 37 | megdnn_assert_eq_layout(oval, values.layout); |
| 38 | int32_t* iptr = nullptr; |
| 39 | if (param().mode == Param::Mode::KTH_ONLY) { |
| 40 | megdnn_assert_eq_shape(indices.layout, TensorShape{}); |
| 41 | } else { |
| 42 | iptr = indices.ptr<int32_t>(); |
| 43 | megdnn_assert_eq_layout(oidx, indices.layout); |
| 44 | } |
| 45 | megdnn_assert( |
| 46 | workspace.size >= |
| 47 | get_workspace_in_bytes(k, data.layout, values.layout, indices.layout)); |
| 48 | if (static_cast<size_t>(std::abs(k)) > data.layout[1]) { |
| 49 | if (k > 0) { |
| 50 | k = data.layout[1]; |
| 51 | } else { |
| 52 | k = -static_cast<int>(data.layout[1]); |
| 53 | } |
| 54 | } |
| 55 | do_exec(k, data, values, iptr, workspace); |
| 56 | } |
| 57 | |
| 58 | // vim: syntax=cpp.doxygen |
nothing calls this directly
no test coverage detected