diff --git a/backends/vulkan/test/op_tests/utils/gen_computegraph.py b/backends/vulkan/test/op_tests/utils/gen_computegraph.py index b24879f660a..5989468100e 100644 --- a/backends/vulkan/test/op_tests/utils/gen_computegraph.py +++ b/backends/vulkan/test/op_tests/utils/gen_computegraph.py @@ -58,6 +58,8 @@ class ValueRef: src_cpp_type: str is_in: bool = False is_out: bool = False + fixed_storage_type: Optional[str] = None + fixed_memory_layout: Optional[str] = None requires_prepack: bool = False supports_prepack: bool = False # When is_dynamic_size is true, the underlying object size is not known @@ -137,20 +139,43 @@ def __init__( if arg.name in self.suite_def.prepacked_args: supports_prepack = True + fixed_storage_type = None + if arg.name in self.suite_def.arg_storage_types: + fixed_storage_type = self.suite_def.arg_storage_types[arg.name] + + fixed_memory_layout = None + if arg.name in self.suite_def.arg_memory_layouts: + fixed_memory_layout = self.suite_def.arg_memory_layouts[arg.name] + self.refs[arg.name] = ValueRef( name=f"{arg.name}_ref", src_cpp_name=arg.name, src_cpp_type=cpp_type, is_in=(cpp_type in InableCppType), + fixed_storage_type=fixed_storage_type, + fixed_memory_layout=fixed_memory_layout, requires_prepack=requires_prepack, supports_prepack=supports_prepack, ) ret_type = cpp.returns_type(self.f.func.returns, symint=False).cpp_type() self.out = ATenArg(name="out", cpp_type=ret_type, default=None) + + fixed_storage_type = None + if "out" in self.suite_def.arg_storage_types: + fixed_storage_type = self.suite_def.arg_storage_types["out"] + fixed_memory_layout = None + if "out" in self.suite_def.arg_memory_layouts: + fixed_memory_layout = self.suite_def.arg_memory_layouts["out"] + if ret_type == AT_TENSOR: self.refs["out"] = ValueRef( - name="out_ref", src_cpp_name="out", src_cpp_type=ret_type, is_out=True + name="out_ref", + src_cpp_name="out", + src_cpp_type=ret_type, + is_out=True, + fixed_storage_type=fixed_storage_type, + fixed_memory_layout=fixed_memory_layout, ) elif ret_type == TWO_TENSOR_TUPLE: self.refs["out"] = [ @@ -159,12 +184,24 @@ def __init__( src_cpp_name="std::get<0>(out)", src_cpp_type="at::Tensor", is_out=True, + fixed_storage_type=( + fixed_storage_type[0] if fixed_storage_type else None + ), + fixed_memory_layout=( + fixed_memory_layout[0] if fixed_memory_layout else None + ), ), ValueRef( name="out_ref_second", src_cpp_name="std::get<1>(out)", src_cpp_type="at::Tensor", is_out=True, + fixed_storage_type=( + fixed_storage_type[1] if fixed_storage_type else None + ), + fixed_memory_layout=( + fixed_memory_layout[1] if fixed_memory_layout else None + ), ), ValueRef( name="out_ref", @@ -180,18 +217,36 @@ def __init__( src_cpp_name="std::get<0>(out)", src_cpp_type="at::Tensor", is_out=True, + fixed_storage_type=( + fixed_storage_type[0] if fixed_storage_type else None + ), + fixed_memory_layout=( + fixed_memory_layout[0] if fixed_memory_layout else None + ), ), ValueRef( name="out_ref_second", src_cpp_name="std::get<1>(out)", src_cpp_type="at::Tensor", is_out=True, + fixed_storage_type=( + fixed_storage_type[1] if fixed_storage_type else None + ), + fixed_memory_layout=( + fixed_memory_layout[1] if fixed_memory_layout else None + ), ), ValueRef( name="out_ref_third", src_cpp_name="std::get<2>(out)", src_cpp_type="at::Tensor", is_out=True, + fixed_storage_type=( + fixed_storage_type[2] if fixed_storage_type else None + ), + fixed_memory_layout=( + fixed_memory_layout[2] if fixed_memory_layout else None + ), ), ValueRef( name="out_ref", @@ -289,6 +344,7 @@ def create_value_for( # noqa: C901 if not include_declarations: cpp_type = "" + print(ref) if ref.src_cpp_type == OPT_AT_TENSOR: ret_str = f"{cpp_type} {ref.name} = " if prepack: @@ -302,7 +358,12 @@ def create_value_for( # noqa: C901 ret_str += f"{self.graph}{self.dot}" ret_str += "add_input_tensor(" if ref.is_in else "add_tensor(" ret_str += f"{ref.src_cpp_name}->sizes().vec(), " - ret_str += f"from_at_scalartype({ref.src_cpp_name}->scalar_type())); \n" + ret_str += f"from_at_scalartype({ref.src_cpp_name}->scalar_type()" + if ref.fixed_storage_type: + ret_str += f", {ref.fixed_storage_type}" + if ref.fixed_memory_layout: + ret_str += f", {ref.fixed_memory_layout}" + ret_str += "));\n" elif prepack: ret_str += f"{self.graph}{self.dot}" ret_str += f"add_tensorref({ref.src_cpp_name}->sizes().vec(), " @@ -385,7 +446,12 @@ def create_value_for( # noqa: C901 elif ref.src_cpp_type == AT_TENSOR and not prepack: ret_str += "add_input_tensor(" if ref.is_in else "add_tensor(" ret_str += f"{ref.src_cpp_name}.sizes().vec(), " - ret_str += f"from_at_scalartype({ref.src_cpp_name}.scalar_type())); \n" + ret_str += f"from_at_scalartype({ref.src_cpp_name}.scalar_type())" + if ref.fixed_storage_type: + ret_str += f", {ref.fixed_storage_type}" + if ref.fixed_memory_layout: + ret_str += f", {ref.fixed_memory_layout}" + ret_str += ");\n" elif ref.src_cpp_type == AT_TENSOR and prepack: ret_str += f"add_tensorref({ref.src_cpp_name}.sizes().vec(), " ret_str += f"from_at_scalartype({ref.src_cpp_name}.scalar_type()), " diff --git a/backends/vulkan/test/op_tests/utils/gen_correctness_base.py b/backends/vulkan/test/op_tests/utils/gen_correctness_base.py index 5be4ddba6bf..250edf333bc 100644 --- a/backends/vulkan/test/op_tests/utils/gen_correctness_base.py +++ b/backends/vulkan/test/op_tests/utils/gen_correctness_base.py @@ -140,7 +140,13 @@ def call_data_gen_fn(self, arg: Argument, data: Any, terminate: bool = True) -> else self.suite_def.arg_data_range[arg.name] ) - ret_str = f"{self.suite_def.data_gen}({init_list_str(data)}, {tensor_dtype}, {data_range[0]}, {data_range[1]})" + data_gen_fn = ( + self.suite_def.data_gen + if arg.name not in self.suite_def.arg_data_gen_fn + else self.suite_def.arg_data_gen_fn[arg.name] + ) + + ret_str = f"{data_gen_fn}({init_list_str(data)}, {tensor_dtype}, {data_range[0]}, {data_range[1]})" if terminate: ret_str += ";" @@ -288,13 +294,29 @@ def generate_suite_cpp(self) -> str: if (dtype == at::kBool) return at::rand(sizes, at::device(at::kCPU)) > 0.5; - + if (high == 1.0 && low == 0.0) return at::rand(sizes, at::device(at::kCPU).dtype(dtype)); return at::rand(sizes, at::device(at::kCPU).dtype(dtype)) * (high - low) + low; }} +at::Tensor make_zeros_tensor( + std::vector sizes, + at::ScalarType dtype = at::kFloat, + float low = 0.0, + float high = 1.0) {{ + return at::zeros(sizes, at::device(at::kCPU).dtype(dtype)); +}} + +at::Tensor make_ones_tensor( + std::vector sizes, + at::ScalarType dtype = at::kFloat, + float low = 0.0, + float high = 1.0) {{ + return at::ones(sizes, at::device(at::kCPU).dtype(dtype)); +}} + at::Tensor make_seq_tensor( std::vector sizes, at::ScalarType dtype = at::kFloat, diff --git a/backends/vulkan/test/op_tests/utils/gen_correctness_vk.py b/backends/vulkan/test/op_tests/utils/gen_correctness_vk.py index e7cf5ba92a5..c368c23c539 100644 --- a/backends/vulkan/test/op_tests/utils/gen_correctness_vk.py +++ b/backends/vulkan/test/op_tests/utils/gen_correctness_vk.py @@ -29,7 +29,6 @@ class GeneratedOpsTest_{op_name} : public ::testing::TestWithParam< ::std::tuple void SetUp() override {{ GraphConfig config; - config.expect_dynamic_shapes = true; utils::StorageType default_storage_type; utils::GPUMemoryLayout default_memory_layout; std::tie(test_dtype, default_storage_type, default_memory_layout) = GetParam(); diff --git a/backends/vulkan/test/op_tests/utils/test_suite.py b/backends/vulkan/test/op_tests/utils/test_suite.py index 72ba457b5af..427864b0d5d 100644 --- a/backends/vulkan/test/op_tests/utils/test_suite.py +++ b/backends/vulkan/test/op_tests/utils/test_suite.py @@ -5,7 +5,7 @@ # LICENSE file in the root directory of this source tree. from dataclasses import dataclass -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional ################################### ## Generic Test Suite definition ## @@ -23,6 +23,7 @@ def __init__(self, input_cases: List[Any]): self.data_range = (0, 1) self.arg_dtype = {} + self.arg_data_gen_fn: Dict[str, str] = {} self.arg_data_range = {} self.atol: str = "1e-5" @@ -48,3 +49,5 @@ def __init__(self, input_cases: List[Any]): self.layouts: List[str] = ["utils::kChannelsPacked"] self.data_gen: str = "make_rand_tensor" self.force_io: bool = True + self.arg_storage_types: Dict[str, str] = {} + self.arg_memory_layouts: Dict[str, str] = {}