[ET-VK] Serialize all scalar types

Pull Request resolved: #2414

Adds all the [non-quantized scalar types from ATen Vulkan](https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/vulkan/api/Types.h#L21-L26).
```
  _(uint8_t, VK_FORMAT_R8G8B8A8_UINT, Byte)              \
  _(int8_t, VK_FORMAT_R8G8B8A8_SINT, Char)               \
  _(int32_t, VK_FORMAT_R32G32B32A32_SINT, Int)           \
  _(bool, VK_FORMAT_R8G8B8A8_SINT, Bool)                 \
  _(unsigned short, VK_FORMAT_R16G16B16A16_SFLOAT, Half) \
  _(float, VK_FORMAT_FLOAT4, Float)                      \
```

Tensor images and buffers will now be created according to this `dtype`.

It's up to contributors to associate their shader to one of these dtypes, e.g., in [`all_shaders.yaml`](https://github.com/pytorch/executorch/blob/main/backends/vulkan/runtime/graph/ops/glsl/all_shaders.yaml#L32-L46).
```
image_to_nchw:
  ...
  generate_variant_forall:
    DTYPE:
      - VALUE: "half"
        SUFFIX: "half"
      - VALUE: "float"
        SUFFIX: "float"
  ...
```

Differential Revision: [D54873255](https://our.internmc.facebook.com/intern/diff/D54873255/)
ghstack-source-id: 218582139

Author: jorgep31415

backends/vulkan/runtime/VulkanBackend.cpp

@@ -57,9 +57,18 @@ const uint8_t* getConstantDataPtr(
 
 api::ScalarType get_scalar_type(const vkgraph::VkDataType& vk_datatype) {
   switch (vk_datatype) {
-    case (vkgraph::VkDataType::fp32): {
+    case vkgraph::VkDataType::BOOL:
+      return api::kBool;
+    case vkgraph::VkDataType::UINT8:
+      return api::kByte;
+    case vkgraph::VkDataType::INT8:
+      return api::kChar;
+    case vkgraph::VkDataType::INT32:
+      return api::kInt;
+    case vkgraph::VkDataType::FLOAT16:
+      return api::kHalf;
+    case vkgraph::VkDataType::FLOAT32:
       return api::kFloat;
-    }
   }
 }
 

backends/vulkan/serialization/schema.fbs

@@ -10,9 +10,13 @@ table OperatorCall {
   args:[int];
 }
 
-enum VkDataType : short {
-  // IEEE754 single-precision floating-point.
-  fp32 = 0,
+enum VkDataType : byte {
+  BOOL = 0,
+  UINT8 = 1,
+  INT8 = 2,
+  INT32 = 3,
+  FLOAT16 = 4,
+  FLOAT32 = 5,
 }
 
 table VkTensor {

backends/vulkan/serialization/vulkan_graph_builder.py

@@ -35,9 +35,21 @@ def __init__(self, program: ExportedProgram) -> None:
 
     @staticmethod
     def get_vk_datatype(torch_dtype: torch.dtype) -> vk_graph_schema.VkDataType:
-        # TODO(T182302927): Support more dtypes including float16, int(32|64).
-        if torch_dtype == torch.float32:
-            return vk_graph_schema.VkDataType.fp32
+        if torch_dtype == torch.bool:
+            return vk_graph_schema.VkDataType.BOOL
+        elif torch_dtype == torch.uint8:
+            return vk_graph_schema.VkDataType.UINT8
+        elif torch_dtype == torch.int8:
+            return vk_graph_schema.VkDataType.INT8
+        elif torch_dtype == torch.int32:
+            return vk_graph_schema.VkDataType.INT32
+        elif torch_dtype == torch.float16:
+            return vk_graph_schema.VkDataType.FLOAT16
+        elif torch_dtype == torch.float32:
+            return vk_graph_schema.VkDataType.FLOAT32
+        # Narrowing conversion for index tensor produced by max_poolNd_with_indices.
+        elif torch_dtype == torch.int64:
+            return vk_graph_schema.VkDataType.INT32
         else:
             raise AssertionError(f"Invalid dtype for vulkan_preprocess ({torch_dtype})")
 

backends/vulkan/serialization/vulkan_graph_schema.py

@@ -22,7 +22,12 @@ class OperatorCall:
 
 
 class VkDataType(IntEnum):
-    fp32 = 0
+    BOOL = 0
+    UINT8 = 1
+    INT8 = 2
+    INT32 = 3
+    FLOAT16 = 4
+    FLOAT32 = 5
 
 
 @dataclass

backends/vulkan/vulkan_preprocess.py

@@ -6,8 +6,6 @@
 
 from typing import final, List
 
-import executorch.backends.vulkan.serialization.vulkan_graph_schema as vk_graph_schema
-
 from executorch.backends.vulkan.serialization.vulkan_graph_builder import VkGraphBuilder
 from executorch.backends.vulkan.serialization.vulkan_graph_serialize import (
     serialize_vulkan_graph,
@@ -25,20 +23,12 @@
 from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass
 
 from executorch.exir.program._program import _copy_module
-from torch import dtype, float32
 
 DEFAULT_DEBUG_HANDLE = 65535
 
 
 @final
 class VulkanBackend(BackendDetails):
-    @staticmethod
-    def get_vk_datatype(torch_dtype: dtype) -> vk_graph_schema.VkDataType:
-        if torch_dtype == float32:
-            return vk_graph_schema.VkDataType.fp32
-        else:
-            raise AssertionError(f"Invalid dtype for vulkan_preprocess ({torch_dtype})")
-
     @classmethod
     # pyre-ignore
     def preprocess(  # noqa: C901