Skip to content
This repository has been archived by the owner on Jan 24, 2024. It is now read-only.

op unittets for top_k/transpose/triangular_solve #1472

Merged
merged 5 commits into from
Jun 1, 2023
Merged

op unittets for top_k/transpose/triangular_solve #1472

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
569db44
op unittest for topk
zzk0 May 25, 2023
60e3048
op unittest for transpose
zzk0 May 25, 2023
ed0df7d
op unittest for triangular_solve
zzk0 May 25, 2023
c0f5343
triangular solve op: use higher precision
zzk0 May 26, 2023
8dbdb47
Merge branch 'PaddlePaddle:develop' into op_test_top_k
zzk0 May 31, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion cinn/hlir/op/contrib/sort.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -364,7 +364,7 @@ std::vector<std::vector<int>> InferShapeForTopK(const std::vector<std::vector<in
}
CHECK_GE(axis, 0);
CHECK_LT(axis, res[0].size());
res[0][axis] = k;
res[0][axis] = std::min(res[0][axis], k);
return {res[0], res[0]};
}

Expand Down
302 changes: 260 additions & 42 deletions python/tests/ops/test_top_k_op.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,76 +14,294 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
import numpy as np
from op_test import OpTest, OpTestTool
import paddle
import paddle.nn.functional as F
import cinn
from cinn.frontend import *
from cinn.common import *
from op_test import OpTest, OpTestTool
from op_test_helper import TestCaseHelper, run_test


@OpTestTool.skip_if(not is_compiled_with_cuda(),
"x86 test will be skipped due to timeout.")
class TestTopKOp(OpTest):
def setUp(self):
self.init_case()
print(f"\nRunning {self.__class__.__name__}: {self.case}")
self.inputs = {}
self.prepare_inputs()

def init_case(self):
def prepare_inputs(self):
self.inputs = {
"x1": np.random.random([
2,
4,
]).astype("float32")
"x": self.random(self.case["shape"], self.case["dtype"])
}
self.k = 1
self.axis = 1
self.k = self.case["k"]
self.axis = self.case["axis"]
self.largest = self.case["largest"]

def build_paddle_program(self, target):
axis = -1
x1 = paddle.to_tensor(self.inputs["x1"], stop_gradient=True)
out = paddle.topk(x1, self.k, self.axis)

x = paddle.to_tensor(self.inputs["x"], stop_gradient=True)
if x.shape[self.axis] < self.k:
self.k = x.shape[self.axis]
out = paddle.topk(x, self.k, self.axis)
self.paddle_outputs = [out[0], out[1]]

def build_cinn_program(self, target):
builder = NetBuilder("sum")
x1 = builder.create_input(Float(32), self.inputs["x1"].shape, "x1")
out = builder.top_k(x1, self.k, self.axis, False)
builder = NetBuilder("topk")
x = builder.create_input(
self.nptype2cinntype(self.inputs["x"].dtype),
self.inputs["x"].shape, "x")
out = builder.top_k(x, self.k, self.axis, self.largest)
prog = builder.build()
forward_res = self.get_cinn_output(
prog, target, [x1], [self.inputs["x1"]], [out[0], out[1]])

prog, target, [x], [self.inputs["x"]], [out[0], out[1]])
self.cinn_outputs = forward_res

def test_check_results(self):
self.check_outputs_and_grads()


class TestTopKCase1(TestTopKOp):
def init_case(self):
self.inputs = {
"x1": np.random.random([
2,
4,
]).astype("float32")
}
self.k = 2
self.axis = 1
class TestTopKOpDumpicateElement(TestTopKOp):
def setUp(self):
self.inputs = {}
self.prepare_inputs()

def prepare_inputs(self):
self.inputs = {"x": self.random([128], "int64", -10, 10)}
self.axis = 0
self.largest = False
self.k = 5

class TestTopKCase2(TestTopKOp):
def init_case(self):
self.inputs = {
"x1": np.random.random([
2,
4,
]).astype("float32")
}
self.k = 1

# known issue: same as sort op
# This test case will cause CINN to allocate a large amount of GPU memory, nearly 10 GB.
class TestTopKOpLargeCudaMemoryOccupation(TestTopKOp):
def setUp(self):
self.inputs = {}
self.prepare_inputs()

def prepare_inputs(self):
self.inputs = {"x": self.random([8192], "float64")}
self.axis = 0
self.largest = False
self.k = 5


class TestTopKOpShapeTest(TestCaseHelper):
def init_attrs(self):
self.class_name = "TestTopKOpShapeTest"
self.cls = TestTopKOp
self.inputs = [
{
"shape": [512],
"k": 3
},
{
"shape": [1024],
"k": 10
},
{
"shape": [1200],
"k": 1024
},
{
"shape": [64, 16],
"k": 3
},
{
"shape": [4, 32, 8],
"k": 4
},
{
"shape": [16, 8, 4, 2],
"k": 5
},
{
"shape": [2, 8, 4, 2, 5],
"k": 1
},
{
"shape": [4, 8, 1, 2, 16],
"k": 3
},
{
"shape": [1],
"k": 1
},
{
"shape": [1, 1, 1, 1],
"k": 1
},
{
"shape": [1, 1, 1, 1, 1],
"k": 1
},
]
self.dtypes = [{"dtype": "float32"}]
self.attrs = [{"axis": 0, "largest": True}]


class TestTopKOpDtypeTest(TestCaseHelper):
def init_attrs(self):
self.class_name = "TestTopKOpDtypeTest"
self.cls = TestTopKOp
self.inputs = [
{
"shape": [1024],
},
{
"shape": [64, 16],
},
{
"shape": [4, 32, 8],
},
{
"shape": [16, 8, 4, 2],
},
]
self.dtypes = [
{
"dtype": "float32"
},
{
"dtype": "float64"
},
{
"dtype": "int32"
},
{
"dtype": "int64"
},
]
self.attrs = [{"axis": 0, "largest": True, "k": 3}]


class TestTopKOpAxisTest(TestCaseHelper):
def init_attrs(self):
self.class_name = "TestTopKOpAxisTest"
self.cls = TestTopKOp
self.inputs = [
{
"shape": [16, 8, 4, 8],
},
]
self.dtypes = [{"dtype": "float32"}]
self.attrs = [
{
"axis": 0,
"largest": True,
"k": 3
},
{
"axis": 1,
"largest": True,
"k": 3
},
{
"axis": 2,
"largest": True,
"k": 3
},
{
"axis": 3,
"largest": True,
"k": 3
},
]


class TestTopKOpKTest(TestCaseHelper):
def init_attrs(self):
self.class_name = "TestTopKOpKTest"
self.cls = TestTopKOp
self.inputs = [
{
"shape": [16, 8, 4, 2],
},
]
self.dtypes = [{"dtype": "float32"}]
self.attrs = [
{
"axis": 0,
"largest": True,
"k": 8
},
{
"axis": 1,
"largest": True,
"k": 4
},
{
"axis": 2,
"largest": True,
"k": 2
},
{
"axis": 3,
"largest": True,
"k": 1
},
{
"axis": 0,
"largest": True,
"k": 20
},
{
"axis": 1,
"largest": True,
"k": 10
},
{
"axis": 2,
"largest": True,
"k": 10
},
{
"axis": 3,
"largest": True,
"k": 5
},
]


class TestTopKOpAscendingTest(TestTopKOpShapeTest):
def init_attrs(self):
self.class_name = "TestTopKOpAscendingTest"
self.cls = TestTopKOp
self.inputs = [
{
"shape": [16, 8, 4, 8],
},
]
self.dtypes = [{"dtype": "float32"}]
self.attrs = [
{
"axis": 0,
"largest": False,
"k": 3
},
{
"axis": 1,
"largest": False,
"k": 3
},
{
"axis": 2,
"largest": False,
"k": 3
},
{
"axis": 3,
"largest": False,
"k": 3
},
]


if __name__ == "__main__":
unittest.main()
run_test(TestTopKOpDumpicateElement)
# run_test(TestTopKOpLargeCudaMemoryOccupation)

TestTopKOpShapeTest().run()
TestTopKOpDtypeTest().run()
TestTopKOpAxisTest().run()
TestTopKOpKTest().run()
TestTopKOpAscendingTest().run()
Loading