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

Commit

Permalink
op unittets for top_k/transpose/triangular_solve (#1472)
Browse files Browse the repository at this point in the history 
* op unittest for topk

* op unittest for transpose

* op unittest for triangular_solve

* triangular solve op: use higher precision
  • Loading branch information
@zzk0
zzk0 authored Jun 1, 2023
1 parent d21cbef commit 977407d
Show file tree
Hide file tree
Showing 4 changed files with 720 additions and 258 deletions.
2 changes: 1 addition & 1 deletion cinn/hlir/op/contrib/sort.cc
View file
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
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

0 comments on commit 977407d

Please sign in to comment.