Update on "[ET-VK] Enable FP16 type in operators"

Differential Revision: [D56189470](https://our.internmc.facebook.com/intern/diff/D56189470/)

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/runtime/api/Adapter.cpp

@@ -6,6 +6,8 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+// NOLINTBEGIN(clang-diagnostic-missing-field-initializers)
+
 #include <executorch/backends/vulkan/runtime/api/Adapter.h>
 
 #include <bitset>

backends/vulkan/runtime/api/Tensor.cpp

@@ -228,7 +228,14 @@ vTensor::vTensor(
           memory_layout_,
           gpu_sizes_,
           dtype_,
-          allocate_memory)) {}
+          allocate_memory)) {
+  if (dtype == api::kHalf) {
+    VK_CHECK_COND(
+        api::context()->adapter_ptr()->has_16bit_storage(),
+        "Half dtype is only available if the physical device supports float16 "
+        "storage buffers!");
+  }
+}
 
 vTensor::vTensor(
     api::Context* const context,

backends/vulkan/runtime/api/gen_vulkan_spv.py

@@ -92,9 +92,8 @@ def define_variable(name: str) -> str:
 def get_buffer_scalar_type(dtype: str) -> str:
     if dtype == "half":
         return "float16_t"
-    # TODO(ssjia): use int8_t for int8 types
     elif dtype[-1] == "8":
-        return dtype[:-1]
+        return dtype + "_t"
 
     return dtype
 

backends/vulkan/test/op_tests/cases.py

@@ -21,7 +21,7 @@
 
 
 def get_binary_elementwise_inputs():
-    return VkTestSuite(
+    test_suite = VkTestSuite(
         [
             ((M1, M2), (M1, M2)),
             ((M1, M2), (M1, 1), 2.0),
@@ -31,6 +31,11 @@ def get_binary_elementwise_inputs():
             ((S, S1, S2), (S, 1, S2), 2.0),
         ]
     )
+    test_suite.layouts = [
+        "api::kWidthPacked",
+        "api::kChannelsPacked",
+    ]
+    return test_suite
 
 
 def get_mm_inputs():
@@ -41,7 +46,12 @@ def get_mm_inputs():
         ],
     )
     test_suite.prepacked_args = ["mat2"]
+    # ATen matmul doesn't support half
     test_suite.dtypes = ["at::kFloat"]
+    test_suite.layouts = [
+        "api::kWidthPacked",
+        "api::kChannelsPacked",
+    ]
     return test_suite
 
 
@@ -51,7 +61,6 @@ def get_pool2d_inputs():
             ((S, M1, M2), [2, 2], [1, 1], [0, 0], [1, 1]),
         ]
     )
-    test_suite.supports["layouts"] = ["api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED"]
     return test_suite
 
 
@@ -115,7 +124,6 @@ def get_conv2d_inputs():
             ),
         ]
     )
-    test_suite.supports["layouts"] = ["api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED"]
     return test_suite
 
 
@@ -127,7 +135,6 @@ def get_native_layer_norm_inputs():
             ((S, XL, M1, M2), [M2], (M2), (M2), 0.001),
         ]
     )
-    test_suite.supports["layouts"] = ["api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED"]
     return test_suite
 
 
@@ -139,7 +146,6 @@ def get_full_inputs():
             ([L, M, M1, M2], 2.72),
         ]
     )
-    test_suite.supports["layouts"] = ["api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED"]
     return test_suite
 
 
@@ -162,7 +168,6 @@ def get_select_int_inputs():
             ((8, 6, 1, 1), 1, 4),
         ]
     )
-    test_suite.supports["layouts"] = ["api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED"]
     return test_suite
 
 

backends/vulkan/test/op_tests/targets.bzl

@@ -92,6 +92,6 @@ def define_common_targets(is_fbcode = False):
         deps = [
             "//third-party/googletest:gtest_main",
             "//executorch/backends/vulkan:vulkan_graph_runtime",
-            ":all_aten_ops_lib",
+            runtime.external_dep_location("libtorch"),
         ],
     )

backends/vulkan/test/op_tests/utils/codegen.py

@@ -39,13 +39,10 @@
 
 @dataclass
 class VkTestSuite(TestSuite):
-    supports = {
-        "storage_types": ["api::StorageType::TEXTURE_3D"],
-        "layouts": [
-            "api::GPUMemoryLayout::TENSOR_WIDTH_PACKED",
-            "api::GPUMemoryLayout::TENSOR_CHANNELS_PACKED",
-        ],
-    }
+    def __init__(self, input_cases: List[Any]):
+        super().__init__(input_cases)
+        self.storage_types: List[str] = ["api::kTexture3D"]
+        self.layouts: List[str] = ["api::kChannelsPacked"]
 
 
 ##########################
@@ -376,7 +373,7 @@ def gen_graph_exec_code(self) -> str:
         return graph_exec
 
     def gen_conditional_skips(self) -> str:
-        skips = f"if (test_dtype == at::kHalf && "
+        skips = "if (test_dtype == at::kHalf && "
         skips += f"!{self.graph}{self.dot}context()->adapter_ptr()->has_16bit_storage()) {{\n"
         skips += "  GTEST_SKIP();"
         skips += "}\n"
@@ -458,14 +455,13 @@ def generate_fixture_cpp(self) -> str:
         )
 
     def gen_parameterization(self) -> str:
-        # pyre-ignore
         dtypes = self.suite_def.dtypes
-        storage_types = self.suite_def.supports["storage_types"]
-        layouts = self.suite_def.supports["layouts"]
+        storage_types = self.suite_def.storage_types
+        layouts = self.suite_def.layouts
 
         return f"""
         INSTANTIATE_TEST_SUITE_P(
-            StorageLayoutCombos_{self.op_name},
+            Combos_{self.op_name},
             GeneratedOpsTest_{self.op_name},
             ::testing::Combine(
                 ::testing::Values({', '.join(dtypes)}),
@@ -513,9 +509,11 @@ def gen_parameterization(self) -> str:
         return true;
     }
     bool is_close = at::allclose(t1, t2, rtol, atol);
-    if (!is_close) {
-        std::cout << "t1:" << t1 << std::endl;
-        std::cout << "t2:" << t2 << std::endl;
+    if (!is_close && t1.numel() < 500) {
+        std::cout << "reference: " << std::endl;
+        print(t1, 150);
+        std::cout << "vulkan: " << std::endl;
+        print(t2, 150);
     }
     return is_close;
 }

backends/vulkan/test/op_tests/utils/codegen_base.py

@@ -34,12 +34,12 @@
 ###########################
 
 
-@dataclass
 class TestSuite:
-    input_cases: List[Any]
-    prepacked_args = []
-    requires_prepack = False
-    dtypes = ["at::kFloat", "at::kHalf"]
+    def __init__(self, input_cases: List[Any]):
+        self.input_cases: List[Any] = input_cases
+        self.prepacked_args: List[str] = []
+        self.requires_prepack: bool = False
+        self.dtypes: List[str] = ["at::kFloat", "at::kHalf"]
 
     def supports_prepack(self):
         return len(self.prepacked_args) > 0

backends/vulkan/test/test_vulkan_delegate.py

@@ -54,8 +54,6 @@ def assert_outputs_equal(
                         )
                     )
         else:
-            print(model_output[0])
-            print(ref_output)
             # If one output, eager returns tensor while executor tuple of size 1
             self.assertTrue(
                 torch.allclose(model_output[0], ref_output, atol=atol, rtol=rtol)
@@ -200,8 +198,8 @@ def forward(self, x, y):
 
         sub_module = SubModule()
         sample_inputs = (
-            torch.rand(size=(2, 3), dtype=torch.float16),
-            torch.rand(size=(2, 3), dtype=torch.float16),
+            torch.rand(size=(2, 3), dtype=torch.float32),
+            torch.rand(size=(2, 3), dtype=torch.float32),
         )
 
         self.lower_module_and_test_output(sub_module, sample_inputs)