| # Copyright (c) Meta Platforms, Inc. and affiliates. |
| # All rights reserved. |
| # |
| # This source code is licensed under the BSD-style license found in the |
| # LICENSE file in the root directory of this source tree. |
| |
| |
| from collections import namedtuple |
| from typing import Callable |
| |
| from executorch.backends.vulkan.test.op_tests.utils.test_suite import VkTestSuite |
| |
| |
| # Prime numbers dim sizes for testing |
| XL = 113 |
| L = 89 |
| M2 = 41 |
| M1 = 37 |
| M = 29 |
| S2 = 11 |
| S1 = 7 |
| S = 5 |
| XS = 3 |
| |
| test_suites = {} |
| |
| |
| def register_test_suite(aten_op): |
| def test_suite_decorator(fn: Callable) -> Callable: |
| if isinstance(aten_op, str): |
| test_suites[aten_op] = fn() |
| elif isinstance(aten_op, list): |
| for op in aten_op: |
| test_suites[op] = fn() |
| return fn |
| |
| return test_suite_decorator |
| |
| |
| @register_test_suite( |
| ["aten.add.Tensor", "aten.sub.Tensor", "aten.div.Tensor", "aten.mul.Tensor"] |
| ) |
| def get_binary_elementwise_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ((M1, M2), (M1, M2)), |
| ((M1, M2), (M1, 1), 2.0), |
| ((M1, M2), (1, M2)), |
| ((S, S1, S2), (S, S1, S2)), |
| ((S, S1, S2), (S, S1, 1), 2.0), |
| ((S, S1, S2), (S, 1, S2), 2.0), |
| ((XS, S, S1, S2), (XS, S, 1, 1), 2.0), |
| ] |
| ) |
| test_suite.layouts = [ |
| "utils::kWidthPacked", |
| "utils::kChannelsPacked", |
| ] |
| return test_suite |
| |
| |
| @register_test_suite("aten.mm.default") |
| def get_mm_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ((M1, L), (L, M2)), |
| ((S1, S2), (S2, M)), |
| ((6, 32), (32, 64)), |
| ], |
| ) |
| test_suite.prepacked_args = ["mat2"] |
| # ATen matmul doesn't support half |
| test_suite.dtypes = ["at::kFloat"] |
| test_suite.storage_types = ["utils::kTexture3D", "utils::kBuffer"] |
| test_suite.layouts = [ |
| "utils::kWidthPacked", |
| "utils::kChannelsPacked", |
| ] |
| return test_suite |
| |
| |
| @register_test_suite("aten.bmm.default") |
| def get_bmm_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ((S, M1, L), (S, L, M2)), |
| ((M, S1, S2), (M, S2, M)), |
| ((4, 6, 32), (4, 32, 16)), |
| ], |
| ) |
| test_suite.prepacked_args = ["mat2"] |
| # ATen matmul doesn't support half |
| test_suite.dtypes = ["at::kFloat"] |
| test_suite.layouts = [ |
| "utils::kWidthPacked", |
| "utils::kChannelsPacked", |
| ] |
| return test_suite |
| |
| |
| @register_test_suite("aten.addmm.default") |
| def get_addmm_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ((1, S), (S1, S), (S, S), 1.0, 1.5), |
| ((S, 1), (S, S1), (S1, S1), 1.0, 1.0), |
| ((M1, M2), (M1, M2), (M2, M2)), |
| ((M1, M2), (M1, M2), (M2, M2), 4.2, 2.3), |
| ((M1, 1), (M1, L), (L, L), 2.0, 3.0), |
| ((M2), (M1, M2), (M2, M2)), |
| ((6, M2), (6, M2), (M2, M2)), |
| ] |
| ) |
| # ATen matmul doesn't support half |
| test_suite.dtypes = ["at::kFloat"] |
| test_suite.layouts = [ |
| "utils::kWidthPacked", |
| "utils::kChannelsPacked", |
| ] |
| return test_suite |
| |
| |
| common_MKN_list = [ |
| (S2, M2, M1), |
| (L, L, M1), |
| ] |
| |
| |
| def get_linear_texture_inputs(): |
| MKN_list = common_MKN_list |
| |
| inputs_list = [((M, K), (N, K), None) for M, K, N in MKN_list] |
| inputs_list += [((M, K), (N, K), (N)) for M, K, N in MKN_list] |
| inputs_list += [((3, M, K), (N, K), None) for M, K, N in MKN_list] |
| inputs_list += [((3, M, K), (N, K), (N)) for M, K, N in MKN_list] |
| inputs_list += [((3, 6, K), (N, K), (N)) for M, K, N in MKN_list] |
| |
| test_suite = VkTestSuite(inputs_list) |
| test_suite.dtypes = ["at::kFloat"] |
| test_suite.layouts = [ |
| "utils::kWidthPacked", |
| "utils::kChannelsPacked", |
| ] |
| test_suite.test_name_suffix = "texture" |
| return test_suite |
| |
| |
| def get_linear_buffer_inputs(): |
| MKN_list = common_MKN_list |
| |
| inputs_list = [((M, K), (N, K), None) for M, K, N in MKN_list] |
| inputs_list += [((3, M, K), (N, K), None) for M, K, N in MKN_list] |
| |
| test_suite = VkTestSuite(inputs_list) |
| test_suite.dtypes = ["at::kFloat"] |
| test_suite.layouts = [ |
| "utils::kWidthPacked", |
| "utils::kChannelsPacked", |
| ] |
| test_suite.storage_types = ["utils::kBuffer"] |
| test_suite.test_name_suffix = "buffer" |
| return test_suite |
| |
| |
| @register_test_suite("aten.linear.default") |
| def get_linear_test_suites(): |
| return [get_linear_texture_inputs(), get_linear_buffer_inputs()] |
| |
| |
| @register_test_suite("aten._weight_int8pack_mm.default") |
| def get_weight_int8pack_mm_inputs(): |
| MKN_list = common_MKN_list |
| |
| inputs_list = [((M, K), (N, K), (N)) for M, K, N in MKN_list] |
| |
| test_suite = VkTestSuite(inputs_list) |
| test_suite.dtypes = ["at::kFloat", "at::kHalf"] |
| test_suite.layouts = ["utils::kWidthPacked"] |
| test_suite.storage_types = ["utils::kTexture3D", "utils::kBuffer"] |
| test_suite.prepacked_args = ["mat2", "scales"] |
| test_suite.requires_prepack = True |
| |
| test_suite.arg_dtype["mat2"] = "at::kChar" |
| test_suite.arg_data_range["mat2"] = (0, 100) |
| |
| test_suite.arg_data_range["scales"] = (0.0008, 0.001) |
| |
| return test_suite |
| |
| |
| @register_test_suite("aten.avg_pool2d.default") |
| def get_avg_pool2d_inputs(): |
| Test = namedtuple( |
| "VkAvgPoolTest", |
| [ |
| "self", |
| "kernel_size", |
| "stride", |
| "padding", |
| "ceil_mode", |
| "count_include_pad", |
| "divisor_override", |
| ], |
| ) |
| |
| test_cases = [] |
| for ceil_mode in [True, False]: |
| for count_include_pad in [True, False]: |
| for divisor_override in [None, 5]: |
| test_cases += [ |
| Test( |
| self=(S, M1, M2), |
| kernel_size=[2, 2], |
| stride=[1, 1], |
| padding=[0, 0], |
| ceil_mode=ceil_mode, |
| count_include_pad=count_include_pad, |
| divisor_override=divisor_override, |
| ), |
| ] |
| |
| test_suite = VkTestSuite([tuple(tc) for tc in test_cases]) |
| test_suite.dtypes = ["at::kFloat"] |
| return test_suite |
| |
| |
| @register_test_suite("aten.max_pool2d_with_indices.default") |
| def get_max_pool2d_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ((S, M1, M2), [2, 2], [1, 1], [0, 0], [1, 1]), |
| ] |
| ) |
| return test_suite |
| |
| |
| @register_test_suite("aten.convolution.default") |
| def get_conv_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ( |
| (1, 6, 40, 50), |
| (8, 6, 3, 3), |
| (8,), |
| [1, 2], |
| [2, 3], |
| [1, 1], |
| False, |
| [0, 0], |
| 1, |
| ), |
| ( |
| (1, 6, 40, 50), |
| (6, 8, 3, 3), |
| (8,), |
| [1, 2], |
| [2, 3], |
| [1, 1], |
| True, |
| [0, 1], |
| 1, |
| ), |
| ( |
| (1, 8, 72, 96), |
| (8, 1, 3, 3), |
| (8,), |
| [1, 1], |
| [1, 1], |
| [1, 1], |
| False, |
| [0, 0], |
| 8, |
| ), |
| ( |
| (1, 8, 72, 96), |
| (8, 8, 1, 1), |
| (8,), |
| [1, 1], |
| [1, 1], |
| [1, 1], |
| False, |
| [0, 0], |
| 1, |
| ), |
| ( |
| (1, 6, 40, 50), |
| (8, 6, 3, 3), |
| None, |
| [1, 2], |
| [2, 3], |
| [1, 1], |
| False, |
| [0, 0], |
| 1, |
| ), |
| ( |
| (1, 6, 7), |
| (6, 1, 3), |
| (6,), |
| [1], |
| [0], |
| [1], |
| False, |
| [0], |
| 6, |
| ), |
| ( |
| (2, 20, 30), |
| (10, 4, 6), |
| (10,), |
| [5], |
| [5], |
| [3], |
| False, |
| [0], |
| 5, |
| ), |
| ( |
| (1, 9, 11), |
| (9, 1, 3), |
| None, |
| [1], |
| [0], |
| [1], |
| False, |
| [0], |
| 9, |
| ), |
| ( |
| (5, 15, 30), |
| (20, 3, 3), |
| None, |
| [3], |
| [5], |
| [7], |
| False, |
| [0], |
| 5, |
| ), |
| ( |
| (1, 16, 672, 512), |
| (64, 16, 1, 1), |
| (64,), |
| [1, 1], |
| [0, 0], |
| [1, 1], |
| False, |
| [0, 0], |
| 1, |
| ), |
| ] |
| ) |
| return test_suite |
| |
| |
| @register_test_suite("aten.native_layer_norm.default") |
| def get_native_layer_norm_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ((S1, S2), [S2], (S2), (S2), 0.001), |
| ((M, M1, M2), [M2], (M2), (M2), 0.001), |
| ((S, XL, M1, M2), [M2], (M2), (M2), 0.001), |
| ] |
| ) |
| return test_suite |
| |
| |
| @register_test_suite("aten.upsample_nearest2d.vec") |
| def get_upsample_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| # (input tensor shape, output 2D image size (H, W), output scaling factors) |
| ((2, 2, 2, 2), None, [1, 1]), |
| ((1, 1, 2, 2), None, [2, 2]), |
| ((1, 1, 2, 2), None, [2, 4]), |
| ((1, 1, 2, 2), None, [4, 2]), |
| ((1, 1, 2, 2), [2, 2], None), |
| ((1, 1, 2, 2), [2, 4], None), |
| ((1, 1, 2, 2), [3, 2], None), |
| ] |
| ) |
| return test_suite |
| |
| |
| @register_test_suite(["aten.full.default", "aten.full_like.default"]) |
| def get_full_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ([S1, S2], 42.0), |
| ([M, M1, M2], 3.14), |
| ([L, M, M1, M2], 2.72), |
| ] |
| ) |
| return test_suite |
| |
| |
| @register_test_suite( |
| [ |
| "aten.zeros.default", |
| "aten.zeros_like.default", |
| "aten.ones.default", |
| "aten.ones_like.default", |
| ] |
| ) |
| def get_ones_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ([S1, S2]), |
| ([M, M1, M2]), |
| ([L, M, M1, M2]), |
| ] |
| ) |
| return test_suite |
| |
| |
| @register_test_suite(["aten.select.int", "aten.select_copy.int"]) |
| def get_select_int_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ((6, 2, 7), 0, 3), |
| ((6, 2, 7), 1, 0), |
| ((6, 2, 7), 2, 3), |
| ((6, 10, 7), 0, 3), |
| ((6, 10, 7), 1, 0), |
| ((6, 10, 7), 1, 9), |
| ((6, 10, 7), 2, 6), |
| ((9, 2, 9, 4), 0, 8), |
| ((9, 2, 9, 4), 1, 1), |
| ((9, 2, 9, 4), 2, 0), |
| ((9, 2, 9, 4), 2, 8), |
| ((9, 2, 9, 4), 3, 3), |
| ((8, 6, 1, 1), 0, 4), |
| ((8, 6, 1, 1), 1, 4), |
| ] |
| ) |
| return test_suite |
| |
| |
| @register_test_suite(["aten.permute.default", "aten.permute_copy.default"]) |
| def get_permute_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ((9, 2, 9, 4), [0, 1, 2, 3]), |
| ((9, 2, 9, 4), [0, 1, 3, 2]), |
| ((9, 2, 9, 4), [0, 2, 1, 3]), |
| ((9, 2, 9, 4), [0, 2, 3, 1]), |
| ((9, 2, 9, 4), [0, 3, 1, 2]), |
| ((9, 2, 9, 4), [0, 3, 2, 1]), |
| ((9, 2, 9, 4), [3, 0, 1, 2]), |
| ((9, 2, 9, 4), [3, 2, 0, 1]), |
| ((9, 2, 9, 4), [2, 3, 0, 1]), |
| ((9, 2, 9, 4), [2, 0, 3, 1]), |
| ((9, 2, 9), [2, 0, 1]), |
| ((9, 2, 9), [1, 2, 0]), |
| ((9, 2), [0, 1]), |
| ((9, 2), [1, 0]), |
| ] |
| ) |
| |
| test_suite.layouts = ["utils::kChannelsPacked"] |
| return test_suite |
| |
| |
| @register_test_suite("aten.view_copy.default") |
| def get_view_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ((3, 4, 5), [1, 1, -1]), |
| ((3, 4, 5), [1, -1, 1]), |
| ((3, 4, 5), [-1, 1, 1]), |
| ((8, 7, 2, 3), [4, 3, 7, 4]), |
| ((8, 7, 2, 3), [7, -1, 2, 1]), |
| ((8, 7, 2, 3), [1, 1, 1, -1]), |
| ((8, 7, 2, 3), [-1]), |
| ((2, 3, 3, 7), [2, -1, 1, 1]), |
| ((3, 5, 2, 7), [7, -1, 2, 1]), |
| ((2, 2, 8, 6), [2, 6, -1, 1]), |
| ((2, 2, 8, 6), [6, -1, 1]), |
| ((S1, S2, S1, S2), [S2, -1, 1, S1]), |
| ((S1, S2, S1, S2), [S1, 1, -1, S2]), |
| ((S1, S2, S1, S2), [-1, 1, S1, S2]), |
| ] |
| ) |
| test_suite.layouts = [ |
| "utils::kWidthPacked", |
| "utils::kHeightPacked", |
| "utils::kChannelsPacked", |
| ] |
| return test_suite |
| |
| |
| @register_test_suite("aten.slice_copy.Tensor") |
| def get_slice_out_inputs(): |
| Test = namedtuple("VkSliceTest", ["self", "dim", "start", "end", "step"]) |
| Test.__new__.__defaults__ = (None, 0, None, None, 1) |
| |
| # Slice by width and height |
| test_cases = [ |
| Test(self=[1, 1, 4, 10], dim=3, start=3), |
| Test(self=[1, 1, 4, 10], dim=3, start=3, step=2), |
| Test(self=[1, 1, 4, 10], dim=3, start=3, end=4, step=2), |
| Test(self=[1, 1, 4, 10], dim=2, start=3), |
| Test(self=[9, 9, 9, 9], dim=2, start=0, end=9, step=1), |
| Test(self=[9, 9, 9, 9], dim=2, start=1, end=8, step=1), |
| Test(self=[9, 9, 9, 9], dim=2, start=1, end=2, step=1), |
| Test(self=[9, 9, 9, 9], dim=3, start=1, end=5, step=1), |
| Test(self=[9, 9, 9, 9], dim=3, start=1, end=5, step=2), |
| Test(self=[9, 9, 9, 9], dim=-1, start=1, end=5, step=2), |
| Test(self=[9, 9, 9, 9], dim=-2, start=1, end=5, step=2), |
| Test(self=[9, 9, 9], dim=1, start=2, step=1), |
| Test(self=[9, 9, 9], dim=1, start=2, step=2), |
| Test(self=[9, 9, 9], dim=2, start=2, step=1), |
| Test(self=[9, 9, 9], dim=2, start=2, step=2), |
| Test(self=[9, 9], dim=0, start=2, step=1), |
| Test(self=[9, 9], dim=0, start=2, step=2), |
| Test(self=[9, 9], dim=1, start=2, step=1), |
| Test(self=[9, 9], dim=1, start=2, step=2), |
| ] |
| |
| # Slice by batch |
| test_cases += [ |
| Test(self=[6, 5, 3, 2], dim=0), |
| Test(self=[6, 5, 3, 2], dim=0, step=2), |
| Test(self=[13, 13, 3, 2], dim=0, step=2), |
| Test(self=[13, 13, 3, 2], dim=0, start=1, step=2), |
| Test(self=[13, 13, 3, 2], dim=0, start=1, step=5), |
| Test(self=[13, 13, 3, 2], dim=0, start=1, step=20), |
| Test(self=[13, 2, 3, 2], dim=0, start=1, step=2), |
| Test(self=[13, 2, 3, 2], dim=0, start=1, step=5), |
| Test(self=[13, 2, 3, 2], dim=0, start=1, step=20), |
| ] |
| |
| # Slice by channel |
| test_cases += [ |
| Test(self=[2, 5, 1, 10], dim=1), |
| Test(self=[2, 5, 1, 10], dim=1, start=1), |
| Test(self=[2, 5, 1, 10], dim=1, start=1, step=2), |
| Test(self=[5, 13, 1, 10], dim=1), |
| Test(self=[5, 13, 1, 10], dim=1, start=1), |
| Test(self=[5, 13, 1, 10], dim=1, start=1, step=2), |
| Test(self=[5, 13, 1, 10], dim=1, start=1, step=5), |
| Test(self=[5, 13, 1, 10], dim=1, start=1, step=20), |
| Test(self=[13, 1, 10], dim=0), |
| Test(self=[13, 1, 10], dim=0, start=1), |
| Test(self=[13, 1, 10], dim=0, start=1, step=2), |
| Test(self=[13, 1, 10], dim=0, start=1, step=5), |
| Test(self=[13, 1, 10], dim=0, start=1, step=20), |
| ] |
| |
| # Slice by negative/unspecified indices |
| INT64_MAX = 9223372036854775807 # represents arr[:] |
| test_cases += [ |
| Test(self=[8, 9], dim=0, start=-2, step=1), |
| Test(self=[8, 9], dim=0, start=-2, step=2), |
| Test(self=[8, 9], dim=0, end=-2, step=1), |
| Test(self=[8, 9], dim=0, end=-2, step=2), |
| Test(self=[8, 9], dim=0, end=INT64_MAX, step=1), |
| Test(self=[8, 9], dim=0, end=INT64_MAX, step=2), |
| Test(self=[8, 9], dim=1, start=-2, step=1), |
| Test(self=[8, 9], dim=1, start=-2, step=2), |
| Test(self=[8, 9], dim=1, end=-2, step=1), |
| Test(self=[8, 9], dim=1, end=-2, step=2), |
| Test(self=[8, 9], dim=1, end=INT64_MAX, step=1), |
| Test(self=[8, 9], dim=1, end=INT64_MAX, step=2), |
| ] |
| |
| test_suite = VkTestSuite([tuple(tc) for tc in test_cases]) |
| |
| test_suite.dtypes = ["at::kFloat", "at::kHalf"] |
| test_suite.layouts = ["utils::kChannelsPacked"] |
| test_suite.data_gen = "make_seq_tensor" |
| return test_suite |
| |
| |
| def get_slice_view_inputs(): |
| Test = namedtuple("VkSliceTest", ["self", "dim", "start", "end", "step"]) |
| Test.__new__.__defaults__ = (None, 0, None, None, 1) |
| |
| # Slice by channel |
| test_cases = [ |
| Test(self=[1, 17, 1, 10], dim=1, start=0, end=4), |
| Test(self=[1, 17, 1, 10], dim=1, start=0, end=8), |
| Test(self=[1, 17, 3, 7], dim=1, start=0, end=12), |
| ] |
| |
| test_suite = VkTestSuite([tuple(tc) for tc in test_cases]) |
| |
| test_suite.dtypes = ["at::kFloat"] |
| test_suite.storage_types = ["utils::kBuffer", "utils::kTexture3D"] |
| test_suite.layouts = ["utils::kWidthPacked"] |
| test_suite.data_gen = "make_seq_tensor" |
| test_suite.is_view_op = True |
| |
| return test_suite |
| |
| |
| @register_test_suite(["aten.slice.Tensor"]) |
| def get_slice_inputs(): |
| texture_test_suite = get_slice_out_inputs() |
| texture_test_suite.test_name_suffix = "no_view" |
| |
| view_test_suite = get_slice_view_inputs() |
| view_test_suite.test_name_suffix = "view" |
| |
| return [view_test_suite, texture_test_suite] |
| |
| |
| @register_test_suite(["aten.transpose.int"]) |
| def get_transpose_inputs(): |
| Test = namedtuple("VkTransposeViewTest", ["self", "dim0", "dim1"]) |
| Test.__new__.__defaults__ = (None, 0, 1) |
| |
| test_cases = [ |
| Test(self=[M1, M2], dim0=0, dim1=1), |
| Test(self=[M1, S2, M], dim0=0, dim1=1), |
| Test(self=[M1, S2, M], dim0=0, dim1=2), |
| Test(self=[M1, S2, M], dim0=2, dim1=1), |
| Test(self=[S, M, S2, M2], dim0=3, dim1=2), |
| Test(self=[S, M, S2, M2], dim0=1, dim1=2), |
| Test(self=[S, M, S2, M2], dim0=3, dim1=1), |
| ] |
| |
| test_suite = VkTestSuite([tuple(tc) for tc in test_cases]) |
| |
| test_suite.dtypes = ["at::kFloat"] |
| test_suite.storage_types = ["utils::kBuffer", "utils::kTexture3D"] |
| test_suite.layouts = ["utils::kWidthPacked", "utils::kChannelsPacked"] |
| test_suite.data_gen = "make_seq_tensor" |
| test_suite.is_view_op = True |
| return test_suite |
| |
| |
| @register_test_suite("aten.index_select.default") |
| def get_index_select_inputs(): |
| Test = namedtuple("VkIndexSelectTest", ["self", "dim", "index"]) |
| Test.__new__.__defaults__ = (None, 0, None) |
| |
| test_cases = [] |
| |
| for i in range(4): |
| test_cases += [ |
| Test(self=[9, 9, 9, 9], dim=i, index=[0]), |
| Test(self=[9, 9, 9, 9], dim=i, index=[2]), |
| Test(self=[9, 9, 9, 9], dim=i, index=[0, 2]), |
| Test(self=[9, 9, 9, 9], dim=i, index=[3, 1]), |
| Test(self=[9, 9, 9, 9], dim=i, index=[5, 5]), |
| Test(self=[9, 9, 9, 9], dim=i, index=[2, 3, 4, 5, 7]), |
| ] |
| |
| test_suite = VkTestSuite([tuple(tc) for tc in test_cases]) |
| |
| test_suite.dtypes = ["at::kFloat"] |
| test_suite.layouts = ["utils::kChannelsPacked"] |
| return test_suite |
| |
| |
| @register_test_suite("aten.embedding.default") |
| def get_embedding_inputs(): |
| Test = namedtuple("VkEmbeddingTest", ["weight", "indices"]) |
| Test.__new__.__defaults__ = (None, None) |
| |
| test_cases = [ |
| Test(weight=[10, 9], indices=[0, 2]), |
| Test(weight=[10, 9], indices=[2, 3, 4, 5, 7]), |
| Test(weight=[10, 9], indices=[[0, 2], [1, 4], [7, 7]]), |
| Test(weight=[10, 9], indices=[[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]), |
| Test(weight=[10, 9], indices=[[[3, 1, 4], [1, 5, 9]], [[2, 6, 5], [3, 5, 8]]]), |
| ] |
| |
| test_suite = VkTestSuite([tuple(tc) + (-1, "false", "false") for tc in test_cases]) |
| |
| test_suite.dtypes = ["at::kFloat"] |
| test_suite.layouts = ["utils::kChannelsPacked"] |
| return test_suite |
| |
| |
| @register_test_suite("aten.unsqueeze_copy.default") |
| def get_unsqueeze_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ((2, 3, 4), 0), |
| ((1, 1, 1), 0), |
| ((1, 1, 1), 1), |
| ((1, 1, 1), 2), |
| ((1, 1, 1), 3), |
| ((9, 9, 9), 0), |
| ((9, 9, 9), 1), |
| ((9, 9, 9), 2), |
| ((9, 9, 9), 3), |
| ((9, 9), 0), |
| ((9, 9), 1), |
| ((9, 9), 2), |
| ((9,), 0), |
| ((9,), 1), |
| ] |
| ) |
| test_suite.layouts = [ |
| "utils::kChannelsPacked", |
| ] |
| test_suite.data_gen = "make_seq_tensor" |
| return test_suite |
| |
| |
| @register_test_suite("aten.clone.default") |
| def get_clone_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ((S2, S1, S2, S1),), |
| ((S2, S1, S2),), |
| ((S2, S1),), |
| ((S2,),), |
| ((XS, S1, XS, S1),), |
| ((XS, S1, XS),), |
| ((S1, XS, S1),), |
| ((XS, S1),), |
| ((S1, XS),), |
| ((S1,),), |
| ((XS,),), |
| ] |
| ) |
| test_suite.layouts = [ |
| "utils::kChannelsPacked", |
| ] |
| test_suite.data_gen = "make_seq_tensor" |
| return test_suite |
| |
| |
| @register_test_suite("aten.repeat.default") |
| def get_repeat_inputs(): |
| test_suite_2d = VkTestSuite( |
| [ |
| ((2, 3), [1, 4]), |
| ((2, 3), [4, 1]), |
| ((2, 3), [4, 4]), |
| ((2, 3), [3, 1, 4]), |
| ] |
| ) |
| test_suite_2d.layouts = ["utils::kChannelsPacked"] |
| test_suite_2d.storage_types = ["utils::kTexture2D"] |
| test_suite_2d.data_gen = "make_seq_tensor" |
| test_suite_2d.dtypes = ["at::kFloat"] |
| test_suite_2d.test_name_suffix = "2d" |
| |
| test_suite_3d = VkTestSuite( |
| [ |
| # Repeat channels only (most challenging case) |
| ((3, XS, S), [2, 1, 1]), |
| ((7, XS, S), [4, 1, 1]), |
| ((1, 7, XS, S), [1, 4, 1, 1]), |
| ((3, 7, XS, S), [1, 4, 1, 1]), |
| # Repat channels with other dims |
| ((1, 7, XS, S), [1, 4, 1, 3]), |
| ((3, 7, XS, S), [1, 4, 1, 3]), |
| ((3, 7, XS, S), [1, 4, 3, 1]), |
| ((3, 7, XS, S), [1, 4, 3, 3]), |
| # Repeat Batch |
| ((3, 7, XS, S), [3, 4, 3, 3]), |
| ((3, 7, XS, S), [3, 1, 3, 3]), |
| # More other cases |
| ((3, 7, 1, 1), [1, 4, 1, 1]), |
| ((2, 3), [1, 4]), |
| ((2, 3), [4, 1]), |
| ((2, 3), [4, 4]), |
| ((S1, S2, S2), [1, 3, 1]), |
| ((S1, S2, S2), [1, 3, 3]), |
| ((S1, S2, S2), [3, 3, 1]), |
| ((S1, S2, S2), [3, 3, 3]), |
| ((S1, S2, S2, S2), [1, 1, 3, 1]), |
| ((S1, S2, S2, S2), [1, 1, 1, 3]), |
| ((S1, S2, S2, S2), [1, 1, 3, 3]), |
| ((S1, S2, S2, S2), [1, 3, 1, 3]), |
| ((S1, S2, S2, S2), [3, 3, 3, 3]), |
| ((S1, S2, S2, S2), [3, 3, 1, 1]), |
| # Expanding cases |
| ((2, 3), [3, 1, 4]), |
| ((2, 3), [3, 3, 2, 4]), |
| ] |
| ) |
| test_suite_3d.layouts = ["utils::kChannelsPacked"] |
| test_suite_3d.storage_types = ["utils::kTexture3D"] |
| test_suite_3d.data_gen = "make_seq_tensor" |
| test_suite_3d.dtypes = ["at::kFloat"] |
| test_suite_3d.test_name_suffix = "3d" |
| |
| return [test_suite_2d, test_suite_3d] |
| |
| |
| @register_test_suite("aten.repeat_interleave.self_int") |
| def get_repeat_interleave_inputs(): |
| test_suite_W = VkTestSuite( |
| [ |
| ((4, 32, 256), 3, -2), |
| # Test repeat on each non-packed dim |
| ((16, 32, 64), 5, -2), |
| ((16, 32, 64), 5, -3), |
| # Test batched inputs |
| ((3, 5, 32, 64), 4, -2), |
| ((3, 5, 32, 64), 4, -3), |
| ] |
| ) |
| test_suite_W.layouts = [ |
| "utils::kWidthPacked", |
| ] |
| test_suite_W.data_gen = "make_seq_tensor" |
| test_suite_W.dtypes = ["at::kFloat"] |
| test_suite_W.test_name_suffix = "W_packed" |
| |
| test_suite_C = VkTestSuite( |
| [ |
| # Test repeat on each non-packed dim |
| ((32, 32, 16), 5, -1), |
| ((32, 32, 16), 5, -2), |
| # Test batched inputs |
| ((3, 16, 8, 64), 4, -1), |
| ((3, 16, 8, 64), 4, -2), |
| ] |
| ) |
| test_suite_C.layouts = [ |
| "utils::kChannelsPacked", |
| ] |
| test_suite_C.data_gen = "make_seq_tensor" |
| test_suite_C.dtypes = ["at::kFloat"] |
| test_suite_C.test_name_suffix = "C_packed" |
| |
| return [test_suite_W, test_suite_C] |
| |
| |
| @register_test_suite("aten.cat.default") |
| def get_cat_inputs(): |
| # TensorList must be specified as list of tuples |
| test_suite = VkTestSuite( |
| [ |
| # Cat on Height |
| ([(S1, S1, 3, 5), (S1, S1, 0, 5)], 2), |
| ([(S1, S1, 3, 5), (S1, S1, 4, 5)], 2), |
| ([(S1, 3, 5), (S1, 4, 5)], 1), |
| ([(3, 5), (4, 5)], 0), |
| ([(3, 5), (4, 5), (1, 5)], 0), |
| ( |
| [(3, 5)], |
| 0, |
| ), |
| # Cat on Width |
| ([(S1, S1, 5, 3), (S1, S1, 5, 4)], 3), |
| ([(S1, 5, 3), (S1, 5, 4)], 2), |
| ([(5, 0), (5, 4)], 1), |
| ([(5, 3), (5, 4)], 1), |
| ([(5, 3), (5, 4), (5, 1)], 1), |
| ( |
| [(5, 4)], |
| 1, |
| ), |
| ([(5,), (6,)], 0), |
| # Cat on Batch |
| ([(S, S1, 5, 4), (S1, S1, 5, 4)], 0), |
| ([(S, XS, 5, 4), (S1, XS, 5, 4)], 0), |
| ([(S, S2, 5, 4), (S1, S2, 5, 4)], 0), |
| ( |
| [ |
| (3, 1, 2, 5), |
| (3, 1, 2, 5), |
| (3, 1, 2, 5), |
| ], |
| 0, |
| ), |
| # Cat on Channel |
| ([(S, 5, 4), (0, 5, 4), (S2, 5, 4)], 0), |
| ([(S, 5, 4), (S1, 5, 4), (S2, 5, 4)], 0), |
| ([(XS, 5, 4), (XS, 5, 4), (S2, 5, 4)], 0), |
| ([(XS, S, 5, 4), (XS, S1, 5, 4), (XS, S2, 5, 4)], 1), |
| ([(XS, XS, 5, 4), (XS, XS, 5, 4), (XS, S2, 5, 4)], 1), |
| ( |
| [ |
| (XS, 1, 2, 5), |
| (XS, 1, 2, 5), |
| (XS, 1, 2, 5), |
| ], |
| 1, |
| ), |
| ] |
| ) |
| test_suite.layouts = [ |
| "utils::kChannelsPacked", |
| ] |
| test_suite.data_gen = "make_seq_tensor" |
| test_suite.dtypes = ["at::kFloat"] |
| return test_suite |
| |
| |
| @register_test_suite("aten.split_with_sizes_copy.default") |
| def get_split_with_sizes_inputs(): |
| Test = namedtuple("VkSliceTest", ["self", "sizes", "dim"]) |
| test_cases = [ |
| # Split on Width |
| Test(self=(S1, 7, 10, 10), sizes=[1, 2, 3, 4], dim=3), |
| Test(self=(7, 10, 10), sizes=[1, 2, 3, 4], dim=2), |
| Test(self=(7, 10, 10), sizes=[1, 9], dim=2), |
| Test(self=(10, 10), sizes=[1, 9], dim=1), |
| Test(self=(10,), sizes=[1, 9], dim=0), |
| # Split on Height |
| Test(self=(S1, 7, 10, 10), sizes=[1, 2, 3, 4], dim=2), |
| Test(self=(7, 10, 10), sizes=[1, 2, 3, 4], dim=1), |
| Test(self=(7, 10, 10), sizes=[10], dim=1), |
| Test(self=(7, 6, 10), sizes=[1, 1, 1, 1, 1, 1], dim=1), |
| Test(self=(10, 10), sizes=[1, 2, 3, 4], dim=0), |
| # Split on Batch |
| Test(self=(10, 7, 10, 10), sizes=[3, 6, 1], dim=0), |
| Test(self=(10, 7, 10, 10), sizes=[10], dim=0), |
| # Split on Channel |
| Test(self=(7, 13, 4, 8), sizes=[3, 6, 1, 3], dim=1), |
| Test(self=(7, 13, 4, 8), sizes=[3, 3, 3, 3, 1], dim=1), |
| Test(self=(13, 4, 8), sizes=[3, 3, 3, 3, 1], dim=0), |
| Test(self=(13, 4, 8), sizes=[2, 9, 2], dim=0), |
| Test(self=(13, 4, 8), sizes=[13], dim=0), |
| ] |
| test_suite = VkTestSuite([tuple(tc) for tc in test_cases]) |
| |
| test_suite.layouts = [ |
| "utils::kChannelsPacked", |
| ] |
| test_suite.data_gen = "make_seq_tensor" |
| test_suite.dtypes = ["at::kFloat"] |
| return test_suite |
| |
| |
| @register_test_suite("aten.split.Tensor") |
| def get_split_tensor_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| # Split on Width |
| ((S1, 7, 10, 12), 12, 3), |
| ((S1, 7, 10, 12), 3, 3), |
| ((S1, 7, 10, 12), 1, 3), |
| ((7, 10, 12), 12, 2), |
| ((7, 10, 12), 3, 2), |
| ((7, 10, 12), 1, 2), |
| ((10, 12), 12, 1), |
| ((10, 12), 3, 1), |
| ((10, 12), 1, 1), |
| ((12,), 12, 0), |
| ((12,), 3, 0), |
| ((12,), 1, 0), |
| # Split on Height |
| ((S1, 7, 12, 8), 12, 2), |
| ((S1, 7, 12, 8), 3, 2), |
| ((S1, 7, 12, 8), 1, 2), |
| ((7, 12, 8), 12, 1), |
| ((7, 12, 8), 3, 1), |
| ((7, 12, 8), 1, 1), |
| ((12, 8), 12, 0), |
| ((12, 8), 3, 0), |
| ((12, 8), 1, 0), |
| # Split on Batch |
| ((12, 7, 10, 10), 12, 0), |
| ((12, 7, 10, 10), 3, 0), |
| ((12, 7, 10, 10), 1, 0), |
| # Split on Channel |
| ((7, 15, 10, 10), 15, 1), |
| ((7, 15, 10, 10), 5, 1), |
| ((7, 15, 10, 10), 3, 1), |
| ((7, 15, 10, 10), 1, 1), |
| ((15, 10, 10), 15, 0), |
| ((15, 10, 10), 5, 0), |
| ((15, 10, 10), 3, 0), |
| ((15, 10, 10), 1, 0), |
| ] |
| ) |
| |
| test_suite.layouts = [ |
| "utils::kChannelsPacked", |
| ] |
| test_suite.data_gen = "make_seq_tensor" |
| test_suite.dtypes = ["at::kFloat"] |
| return test_suite |
| |
| |
| def get_reduce_inputs(is_softmax: bool = False): |
| bool_arg = False if is_softmax else True |
| return [ |
| ((L), 0, bool_arg), |
| ((L), -1, bool_arg), |
| ((M, L), 0, bool_arg), |
| ((M, L), 1, bool_arg), |
| ((L, M), -1, bool_arg), |
| ((M, L), -2, bool_arg), |
| ((S, S1, S2), 0, bool_arg), |
| ((S, S1, S2), 1, bool_arg), |
| ((S, S1, S2), 2, bool_arg), |
| ((S, S1, S2), -1, bool_arg), |
| ((S, S1, S2), -2, bool_arg), |
| ((S, S1, S2), -3, bool_arg), |
| ((1, S, S1, S2), 1, bool_arg), |
| ((1, S, S1, S2), 2, bool_arg), |
| ((1, S, S1, S2), 3, bool_arg), |
| ((1, S, S1, S2), -1, bool_arg), |
| ((1, S, S1, S2), -2, bool_arg), |
| ((1, S, S1, S2), -3, bool_arg), |
| # Test batches > 1 where the reduction dim is not the concat dim |
| ((S, S2, S1, 128), -1, bool_arg), |
| ] |
| |
| |
| @register_test_suite(["aten._softmax.default", "aten._log_softmax.default"]) |
| def get_softmax_inputs(): |
| test_suite = VkTestSuite(get_reduce_inputs(is_softmax=True)) |
| test_suite.layouts = [ |
| "utils::kWidthPacked", |
| "utils::kChannelsPacked", |
| ] |
| return test_suite |
| |
| |
| @register_test_suite( |
| ["aten.amax.default", "aten.amin.default", "aten.sum.dim_IntList", "aten.mean.dim"] |
| ) |
| def get_reduce_op_inputs(): |
| test_suite = VkTestSuite(get_reduce_inputs()) |
| test_suite.layouts = [ |
| "utils::kChannelsPacked", |
| "utils::kWidthPacked", |
| ] |
| return test_suite |
| |
| |
| @register_test_suite( |
| [ |
| "aten.sqrt.default", |
| "aten.rsqrt.default", |
| "aten.exp.default", |
| "aten.hardshrink.default", |
| "aten.sin.default", |
| "aten.neg.default", |
| "aten.cos.default", |
| "aten.hardswish.default", |
| "aten.hardsigmoid.default", |
| ] |
| ) |
| def get_unary_ops_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| (M1,), |
| (M1, M2), |
| (S1, M1, M2), |
| (S1, S2, S2, M2), |
| ] |
| ) |
| test_suite.storage_types = ["utils::kTexture3D", "utils::kBuffer"] |
| test_suite.atol = "1e-4" |
| test_suite.rtol = "1e-4" |
| return test_suite |
| |
| |
| @register_test_suite("aten._native_batch_norm_legit_no_training.default") |
| def get_native_batch_norm_inputs(): |
| Test = namedtuple( |
| "VkSliceTest", ["self", "weight", "bias", "mean", "var", "momentum", "eps"] |
| ) |
| |
| test_cases = [ |
| Test( |
| self=(1, 1, 2, 5), |
| weight=(1,), |
| bias=(1,), |
| mean=(1,), |
| var=(1,), |
| momentum=0.0, |
| eps=0.001, |
| ), |
| Test( |
| self=(S2, 1, 2, 5), |
| weight=(1,), |
| bias=(1,), |
| mean=(1,), |
| var=(1,), |
| momentum=0.0, |
| eps=0.001, |
| ), |
| Test( |
| self=(1, S2, 2, 5), |
| weight=(S2,), |
| bias=(S2,), |
| mean=(S2,), |
| var=(S2,), |
| momentum=0.0, |
| eps=0.001, |
| ), |
| Test( |
| self=(9, S1, 2, 5), |
| weight=(S1,), |
| bias=(S1,), |
| mean=(S1,), |
| var=(S1,), |
| momentum=0.0, |
| eps=0.01, |
| ), |
| Test( |
| self=(3, S1, 2, 5), |
| weight=(S1,), |
| bias=(S1,), |
| mean=(S1,), |
| var=(S1,), |
| momentum=0.0, |
| eps=0.001, |
| ), |
| Test( |
| self=(3, S2, 2, 5), |
| weight=(S2,), |
| bias=(S2,), |
| mean=(S2,), |
| var=(S2,), |
| momentum=0.0, |
| eps=0.001, |
| ), |
| Test( |
| self=(3, S2, 2, 5), |
| weight=(S2,), |
| bias=(S2,), |
| mean=(S2,), |
| var=(S2,), |
| momentum=0.0, |
| eps=0.000, |
| ), |
| ] |
| |
| test_suite = VkTestSuite(test_cases) |
| test_suite.requires_prepack = True |
| test_suite.prepacked_args = ["weight", "bias", "mean", "var"] |
| |
| return test_suite |
| |
| |
| @register_test_suite("aten.gelu.default") |
| def get_gelu_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ((M1), "tanh"), |
| ((M1, M2), "tanh"), |
| ((S1, M1, M2), "tanh"), |
| ((S1, S2, S2, M2), "tanh"), |
| ] |
| ) |
| return test_suite |
| |
| |
| @register_test_suite("aten.arange.start_step") |
| def get_arange_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| (1, 13), |
| (1.0, 11), |
| (-13, 3), |
| (-11.0, 2), |
| (3, 15, 3), |
| (3, 23, 2), |
| (3, 23.0, 4), |
| (13, 1, -1), |
| (-3, -13, -2), |
| (13, -2.0, -4), |
| ], |
| ) |
| |
| test_suite.layouts = [ |
| "utils::kChannelsPacked", |
| ] |
| return test_suite |
| |
| |
| @register_test_suite("aten.constant_pad_nd.default") |
| def get_constant_pad_nd_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ([S1, S2], [1, 1], 24.0), |
| ([M, M1, M2], [2, 2], 23.2), |
| ([L, M, M1, M2], [3, 5], 12.2), |
| ([S1, S2], [1, 1, 1, 1], 24.0), |
| ([M, M1, M2], [2, 2, 2, 2], 23.2), |
| ([L, M, M1, M2], [3, 5, 3, 5], 12.2), |
| ([M, M1, M2], [1, 2, 3, 4, 5, 6], 23.2), |
| ([L, M, M1, M2], [3, 3, 3, 3, 3, 3], 12.2), |
| ] |
| ) |
| return test_suite |
| |
| |
| @register_test_suite("aten.minimum.default") |
| def get_minimum_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ((M1, M2), (M2)), |
| ((M1, M2), (M1, M2)), |
| ((M1, M2, M), (M2, M)), |
| ((M1, M1, S1, S2), (M1, M1, S1, S2)), |
| ((S1, S1, S2, S), (S1, S2, S)), |
| ((M1, S1, S2), (L, M1, S1, S2)), |
| ((S1, S2), (L, M1, S1, S2)), |
| ] |
| ) |
| return test_suite |
| |
| |
| @register_test_suite("aten.squeeze_copy.dims") |
| def get_squeeze_copy_dim_inputs(): |
| test_suite = VkTestSuite( |
| [ |
| ([S, S, S, 1], 3), |
| ([S, 1, S, S], 1), |
| ([S, 1, 1, S], [1, 2]), |
| ([1, S, S, S], 0), |
| ([S, S, S, S], 3), |
| ([S, S, S, S], 2), |
| ([S, S, S, S], 1), |
| ([M, M1, 1], 2), |
| ([M, 1, M1], 1), |
| ([1, M1, M1], 0), |
| ] |
| ) |
| return test_suite |
| |
| |
| @register_test_suite("aten.flip.default") |
| def get_flip_inputs(): |
| Test = namedtuple("Flip", ["self", "dim"]) |
| Test.__new__.__defaults__ = (None, 0) |
| |
| test_cases = [ |
| Test(self=[9], dim=[0]), |
| Test(self=[9, 9], dim=[0, 1]), |
| Test(self=[9, 9, 9], dim=[0, 2]), |
| Test(self=[9, 9, 9], dim=[0, 1, 2]), |
| Test(self=[9, 9, 9, 9], dim=[0]), |
| Test(self=[9, 9, 9, 9], dim=[0, 2, 3]), |
| Test(self=[9, 9, 9, 9], dim=[1, 3]), |
| Test(self=[9, 9, 9, 9], dim=[0, 1, 2, 3]), |
| ] |
| |
| test_suite = VkTestSuite([tuple(tc) for tc in test_cases]) |
| return test_suite |
