From 623089a36c59dc66243fd2c69e61b0151383651f Mon Sep 17 00:00:00 2001 From: Jorge Pineda Date: Mon, 8 Apr 2024 13:01:46 -0700 Subject: [PATCH] [ET-VK] Add convolution cases to codegen TSIA Differential Revision: [D55829466](https://our.internmc.facebook.com/intern/diff/D55829466/) [ghstack-poisoned] --- backends/vulkan/test/op_tests/cases.py | 65 +++++++++++++++++++ .../vulkan/test/op_tests/utils/codegen.py | 23 ++++++- .../test/op_tests/utils/codegen_base.py | 13 +++- 3 files changed, 98 insertions(+), 3 deletions(-) diff --git a/backends/vulkan/test/op_tests/cases.py b/backends/vulkan/test/op_tests/cases.py index 91b36a368b8..47c8dfdf5f9 100644 --- a/backends/vulkan/test/op_tests/cases.py +++ b/backends/vulkan/test/op_tests/cases.py @@ -54,6 +54,70 @@ def get_pool2d_inputs(): return test_suite +def get_conv2d_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, + ), + ] + ) + test_suite.supports["layouts"] = ["api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED"] + return test_suite + + test_suites = { "aten.add.Tensor": get_binary_elementwise_inputs(), "aten.sub.Tensor": get_binary_elementwise_inputs(), @@ -61,6 +125,7 @@ def get_pool2d_inputs(): "aten.mul.Tensor": get_binary_elementwise_inputs(), "aten.mm.default": get_mm_inputs(), "aten.max_pool2d_with_indices.default": get_pool2d_inputs(), + "aten.convolution.default": get_conv2d_inputs(), } prepacked_args = {"aten.mm.default": {"mat2"}} diff --git a/backends/vulkan/test/op_tests/utils/codegen.py b/backends/vulkan/test/op_tests/utils/codegen.py index 76e59f53952..fc6b5217f0c 100644 --- a/backends/vulkan/test/op_tests/utils/codegen.py +++ b/backends/vulkan/test/op_tests/utils/codegen.py @@ -12,8 +12,10 @@ AT_INT_ARRAY_REF, AT_SCALAR, AT_TENSOR, + AT_TENSOR_OPT, BOOL, CppTestFileGen, + INT, TENSOR_TUPLE, TestSuite, TestSuiteGen, @@ -96,7 +98,7 @@ def __init__(self, op_reg_name: str, f: NativeFunction, suite_def: TestSuite): ATenArg(name=arg.name, cpp_type=cpp_type, default=arg.default) ) - requires_prepack = "weight" in arg.name + requires_prepack = "weight" in arg.name or "bias" in arg.name supports_prepack = False if arg.name in self.suite_def.prepacked_args: supports_prepack = True @@ -173,6 +175,23 @@ def create_value_for(self, ref: ValueRefList) -> str: prepack = self.prepack_ref(ref) cpp_type = "IOValueRef" if (ref.is_in and not prepack) else "ValueRef" + + if ref.src_cpp_type == AT_TENSOR_OPT: + ret_str = f"{cpp_type} {ref.name} = " + ret_str += f"!{ref.src_cpp_name}.has_value() ? " + ret_str += f"{self.graph}{self.dot}add_none() : " + if not prepack: + 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" + elif prepack: + ret_str += f"{self.graph}{self.dot}" + ret_str += f"add_tensorref({ref.src_cpp_name}->sizes().vec(), " + ret_str += f"from_at_scalartype({ref.src_cpp_name}->scalar_type()), " + ret_str += f"{ref.src_cpp_name}->const_data_ptr()); \n" + return ret_str + ret_str = f"{cpp_type} {ref.name} = {self.graph}{self.dot}" if ref.src_cpp_type == AT_TENSOR and not prepack: ret_str += "add_input_tensor(" if ref.is_in else "add_tensor(" @@ -189,6 +208,8 @@ def create_value_for(self, ref: ValueRefList) -> str: ret_str += f"add_scalar_list({ref.src_cpp_name}.vec()); \n" elif ref.src_cpp_type == BOOL: ret_str += f"add_scalar({ref.src_cpp_name}); \n" + elif ref.src_cpp_type == INT: + ret_str += f"add_scalar({ref.src_cpp_name}); \n" elif ref.src_cpp_type == TENSOR_TUPLE: ret_str += f"add_value_list({{{ref.name}_first, {ref.name}_second}}); \n" else: diff --git a/backends/vulkan/test/op_tests/utils/codegen_base.py b/backends/vulkan/test/op_tests/utils/codegen_base.py index 2af45030ec5..1a8ebc6c6b9 100644 --- a/backends/vulkan/test/op_tests/utils/codegen_base.py +++ b/backends/vulkan/test/op_tests/utils/codegen_base.py @@ -15,10 +15,12 @@ ## ATen code patterns ## ######################## -AT_TENSOR = "at::Tensor" -AT_SCALAR = "at::Scalar" AT_INT_ARRAY_REF = "at::IntArrayRef" +AT_SCALAR = "at::Scalar" +AT_TENSOR = "at::Tensor" +AT_TENSOR_OPT = "::std::optional" BOOL = "bool" +INT = "int64_t" TENSOR_TUPLE = "::std::tuple" ########################### @@ -116,12 +118,19 @@ def create_input_data(self, arg: Argument, data: Any) -> str: if cpp_type == AT_TENSOR: ret_str += f"make_rand_tensor({init_list_str(data)}, test_dtype);" + elif cpp_type == AT_TENSOR_OPT: + if str(data) == "None": + ret_str += "std::nullopt;" + else: + ret_str += f"make_rand_tensor({init_list_str(data)}, test_dtype);" elif cpp_type == AT_SCALAR: ret_str += f"{data};" elif cpp_type == AT_INT_ARRAY_REF: ret_str += f"{init_list_str(data)};" elif cpp_type == BOOL: ret_str += f"{str(data).lower()};" + elif cpp_type == INT: + ret_str += f"{str(data).lower()};" else: raise RuntimeError(f"Unsupported cpp type {cpp_type}") return ret_str + "\n"