1# Copyright 2020-21 Huawei Technologies Co., Ltd 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================ 15 16import numpy as np 17import pytest 18 19import mindspore.context as context 20from mindspore import Tensor 21from mindspore.ops import operations as P 22 23 24@pytest.mark.level0 25@pytest.mark.platform_x86_gpu_training 26@pytest.mark.env_onecard 27def test_topk_small_2d(): 28 context.set_context(mode=context.GRAPH_MODE, device_target="GPU") 29 30 x_np = np.random.rand(3, 4).astype(np.float32) 31 k = 4 32 ms_output = P.TopK(True)(Tensor(x_np), k) 33 np_output = np.sort(x_np, axis=-1)[..., ::-1][..., 0:k] 34 assert np.allclose(ms_output[0].asnumpy(), np_output) 35 36 x_np = np.random.rand(3, 4).astype(np.float32) 37 k = 4 38 ms_output = P.TopK(False)(Tensor(x_np), k) 39 np_output = np.sort(x_np, axis=-1)[..., ::-1][..., 0:k] 40 assert np.allclose(ms_output[0].asnumpy(), np_output) 41 42 43@pytest.mark.level0 44@pytest.mark.platform_x86_gpu_training 45@pytest.mark.env_onecard 46def test_topk_3d(): 47 context.set_context(mode=context.GRAPH_MODE, device_target="GPU") 48 x_np = np.random.rand(2, 256, 128).astype(np.float32) 49 k = 4 50 ms_output = P.TopK(True)(Tensor(x_np), k) 51 np_output = np.sort(x_np, axis=-1)[..., ::-1][..., 0:k] 52 assert np.allclose(ms_output[0].asnumpy(), np_output) 53 54 x_np = np.random.rand(2, 3, 4).astype(np.float32) 55 k = 2 56 ms_output = P.TopK(True)(Tensor(x_np), k) 57 np_output = np.sort(x_np, axis=-1)[..., ::-1][..., 0:k] 58 assert np.allclose(ms_output[0].asnumpy(), np_output) 59 60 61@pytest.mark.level0 62@pytest.mark.platform_x86_gpu_training 63@pytest.mark.env_onecard 64def test_topk_big_2d(): 65 context.set_context(mode=context.GRAPH_MODE, device_target="GPU") 66 x_np = np.random.rand(512, 1024).astype(np.float32) 67 k = 512 68 ms_output = P.TopK(True)(Tensor(x_np), k) 69 np_output = np.sort(x_np, axis=-1)[..., ::-1][..., 0:k] 70 assert np.allclose(ms_output[0].asnumpy(), np_output) 71 72 # sorted elements num greater than max thread per block 73 x_np = np.random.rand(128, 2048).astype(np.float32) 74 k = 1 75 ms_output = P.TopK(True)(Tensor(x_np), k) 76 np_output = np.sort(x_np, axis=-1)[..., ::-1][..., 0:k] 77 assert np.allclose(ms_output[0].asnumpy(), np_output) 78 79 x_np = np.random.rand(32, 2048).astype(np.float32) 80 k = 2048 81 ms_output = P.TopK(True)(Tensor(x_np), k) 82 np_output = np.sort(x_np, axis=-1)[..., ::-1][..., 0:k] 83 assert np.allclose(ms_output[0].asnumpy(), np_output) 84 85 # sorted elements num greater than max share memory per block 86 x_np = np.random.rand(16, 40960).astype(np.float32) 87 k = 1 88 ms_output = P.TopK(True)(Tensor(x_np), k) 89 np_output = np.sort(x_np, axis=-1)[..., ::-1][..., 0:k] 90 assert np.allclose(ms_output[0].asnumpy(), np_output) 91 92 93@pytest.mark.level0 94@pytest.mark.platform_x86_gpu_training 95@pytest.mark.env_onecard 96def test_topk_big_k(): 97 context.set_context(mode=context.GRAPH_MODE, device_target="GPU") 98 x_np = np.random.rand(8, 40960).astype(np.float32) 99 k = 4096 100 ms_output = P.TopK(True)(Tensor(x_np), k) 101 np_output = np.sort(x_np, axis=-1)[..., ::-1][..., 0:k] 102 assert np.allclose(ms_output[0].asnumpy(), np_output) 103 104 105@pytest.mark.level0 106@pytest.mark.platform_x86_gpu_training 107@pytest.mark.env_onecard 108def test_topk_1d(): 109 context.set_context(mode=context.GRAPH_MODE, device_target="GPU") 110 x_np = np.random.rand(12).astype(np.float32) 111 k = 4 112 ms_output = P.TopK(True)(Tensor(x_np), k) 113 np_output = np.sort(x_np)[::-1][0:k] 114 115 assert np.allclose(ms_output[0].asnumpy(), np_output) 116 x_np = np.random.rand(1200).astype(np.float32) 117 k = 256 118 ms_output = P.TopK(True)(Tensor(x_np), k) 119 np_output = np.sort(x_np)[::-1][0:k] 120 assert np.allclose(ms_output[0].asnumpy(), np_output) 121 122 x_np = np.random.rand(250000).astype(np.float32) 123 k = 2000 124 ms_output = P.TopK(True)(Tensor(x_np), k) 125 np_output = np.sort(x_np)[::-1][0:k] 126 assert np.allclose(ms_output[0].asnumpy(), np_output) 127 128 x_np = np.random.rand(10240).astype(np.float32) 129 k = 4096 130 ms_output = P.TopK(True)(Tensor(x_np), k) 131 np_output = np.sort(x_np)[::-1][0:k] 132 assert np.allclose(ms_output[0].asnumpy(), np_output) 133 134 x_np = np.random.rand(720).astype(np.float32) 135 k = 720 136 ms_output = P.TopK(True)(Tensor(x_np), k) 137 np_output = np.sort(x_np)[::-1][0:k] 138 assert np.allclose(ms_output[0].asnumpy()[:k], np_output) 139