Add convolution cases to codegen (#2920)

Summary:
Pull Request resolved: #2920

TSIA
ghstack-source-id: 221926168
exported-using-ghexport
bypass-github-export-checks

Reviewed By: SS-JIA

Differential Revision: D55829466

fbshipit-source-id: 48b4f15c41141093dd061c43e6b769eb4c25c81b

Author: Jorge Pineda

backends/vulkan/test/op_tests/cases.py

@@ -54,13 +54,78 @@ 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(),
     "aten.div.Tensor": get_binary_elementwise_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"}}

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<bool>({ref.src_cpp_name}); \n"
+        elif ref.src_cpp_type == INT:
+            ret_str += f"add_scalar<int64_t>({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:

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<at::Tensor>"
 BOOL = "bool"
+INT = "int64_t"
 TENSOR_TUPLE = "::std::tuple<at::Tensor,at::Tensor>"
 
 ###########################
@@ -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"