| 46 | |
| 47 | template <typename Dtype> |
| 48 | void test::run_topk_test(Handle* handle, bool test_kth_only) { |
| 49 | Checker<TopK> checker{handle}; |
| 50 | using Mode = TopK::Param::Mode; |
| 51 | |
| 52 | bool tie_breaking_mode = false; |
| 53 | Mode cur_mode; |
| 54 | auto output_canonizer = [&](const CheckerHelper::TensorValueArray& arr) { |
| 55 | if (cur_mode == Mode::KTH_ONLY) { |
| 56 | return; |
| 57 | } |
| 58 | auto pinp = arr[0].ptr<typename DTypeTrait<Dtype>::ctype>(); |
| 59 | auto pval = arr[1].ptr<typename DTypeTrait<Dtype>::ctype>(); |
| 60 | auto pidx = arr.at(2).ptr<int>(); |
| 61 | size_t m = arr[1].layout[0], n = arr[1].layout[1]; |
| 62 | using idx_val = std::pair<int, typename DTypeTrait<Dtype>::ctype>; |
| 63 | std::vector<idx_val> data(n); |
| 64 | auto compare = [](const idx_val& it1, const idx_val& it2) { |
| 65 | return (it1.second > it2.second); |
| 66 | }; |
| 67 | for (size_t i = 0; i < m; ++i) { |
| 68 | if (cur_mode == Mode::VALUE_IDX_NOSORT) { |
| 69 | // sort output pairs to canonize |
| 70 | for (size_t j = 0; j < n; ++j) { |
| 71 | data[j].first = pidx[i * n + j]; |
| 72 | data[j].second = pval[i * n + j]; |
| 73 | } |
| 74 | std::sort(data.begin(), data.end(), compare); |
| 75 | for (size_t j = 0; j < n; ++j) { |
| 76 | pidx[i * n + j] = data[j].first; |
| 77 | pval[i * n + j] = data[j].second; |
| 78 | } |
| 79 | } |
| 80 | if (tie_breaking_mode) { |
| 81 | // check if indices are correct and mark all indices to be zero |
| 82 | for (size_t j = 0; j < n; ++j) { |
| 83 | auto idx = pidx[i * n + j]; |
| 84 | auto val = pval[i * n + j]; |
| 85 | // + 0 can change the type, such as changing half to float |
| 86 | ASSERT_EQ(pinp[i * arr[0].layout[1] + idx] + 0, val + 0); |
| 87 | pidx[i * n + j] = 0; |
| 88 | } |
| 89 | } |
| 90 | } |
| 91 | }; |
| 92 | |
| 93 | auto run = [&](int k, size_t m, size_t n, Mode mode, int lda = 0) { |
| 94 | if (::testing::Test::HasFailure()) { |
| 95 | return; |
| 96 | } |
| 97 | cur_mode = mode; |
| 98 | checker.set_proxy(k); |
| 99 | checker.set_param(mode); |
| 100 | TensorLayout layout{{m, n}, Dtype{}}; |
| 101 | if (lda) { |
| 102 | layout.stride[0] = lda; |
| 103 | } |
| 104 | |
| 105 | checker.set_output_canonizer(output_canonizer); |