[ET-VK] double, short, and uint16 dtype runtime support

backends/vulkan/runtime/gen_vulkan_spv.py

@@ -56,52 +56,97 @@
 TYPE_MAPPINGS: Dict[str, Any] = {
     "IMAGE_T": {
         3: {
+            "double": "image3D",
             "float": "image3D",
             "half": "image3D",
-            "int": "iimage3D",
-            "uint": "uimage3D",
+            # integer dtypes
             "int8": "iimage3D",
             "uint8": "uimage3D",
+            "int16": "iimage3D",
+            "uint16": "uimage3D",
+            "int32": "iimage3D",
+            "uint32": "uimage3D",
+            "int64": "iimage3D",
+            "uint64": "uimage3D",
+            # common dtype aliases
             "bool": "uimage3D",
+            "int": "iimage3D",
+            "uint": "uimage3D",
         },
         2: {
+            "double": "image2D",
             "float": "image2D",
             "half": "image2D",
-            "int": "iimage2D",
-            "uint": "uimage2D",
+            # integer dtypes
             "int8": "iimage2D",
             "uint8": "uimage2D",
+            "int16": "iimage2D",
+            "uint16": "uimage2D",
+            "int32": "iimage2D",
+            "uint32": "uimage2D",
+            "int64": "iimage2D",
+            "uint64": "uimage2D",
+            # common dtype aliases
             "bool": "uimage2D",
+            "int": "iimage2D",
+            "uint": "uimage2D",
         },
     },
     "SAMPLER_T": {
         3: {
+            "double": "sampler3D",
             "float": "sampler3D",
             "half": "sampler3D",
-            "int": "isampler3D",
-            "uint": "usampler3D",
+            # integer dtypes
             "int8": "isampler3D",
             "uint8": "usampler3D",
+            "int16": "isampler3D",
+            "uint16": "usampler3D",
+            "int32": "isampler3D",
+            "uint32": "usampler3D",
+            "int64": "isampler3D",
+            "uint64": "usampler3D",
+            # common dtype aliases
             "bool": "usampler3D",
+            "int": "isampler3D",
+            "uint": "usampler3D",
         },
         2: {
+            "double": "sampler2D",
             "float": "sampler2D",
             "half": "sampler2D",
-            "int": "isampler2D",
-            "uint": "usampler2D",
+            # integer dtypes
             "int8": "isampler2D",
             "uint8": "usampler2D",
+            "int16": "isampler2D",
+            "uint16": "usampler2D",
+            "int32": "isampler2D",
+            "uint32": "usampler2D",
+            "int64": "isampler2D",
+            "uint64": "usampler2D",
+            # common dtype aliases
             "bool": "usampler2D",
+            "int": "isampler2D",
+            "uint": "usampler2D",
         },
     },
     "IMAGE_FORMAT": {
+        "double": "rgba32f",
         "float": "rgba32f",
         "half": "rgba16f",
-        "int": "rgba32i",
-        "uint": "rgba32ui",
+        # integer dtypes
         "int8": "rgba8i",
         "uint8": "rgba8ui",
+        "int16": "rgba16i",
+        "uint16": "rgba16ui",
+        "int32": "rgba32i",
+        "uint32": "rgba32ui",
+        "int64": "rgba32i",
+        "uint64": "rgba32ui",
+        # common dtype aliases
         "bool": "rgba8ui",
+        "int": "rgba32i",
+        "uint": "rgba32ui",
     },
 }
 
@@ -118,31 +163,48 @@ def define_variable(name: str) -> str:
 def buffer_scalar_type(dtype: str) -> str:
     if dtype == "half":
         return "float16_t"
-    elif dtype[-1] == "8":
-        return dtype + "_t"
+    elif dtype == "float":
+        return "float"
+    elif dtype == "double":
+        return "float64_t"
+    # integer dtype alias conversion
     elif dtype == "bool":
         return "uint8_t"
+    # we don't want to append _t for int32 or uint32 as int is already 32bit
+    elif dtype == "int32" or dtype == "uint32":
+        return "int" if dtype == "int32" else "uint"
+    elif dtype[-1].isdigit():
+        return dtype + "_t"
     return dtype
 
 
 def buffer_gvec_type(dtype: str, n: int) -> str:
     if n == 1:
         return buffer_scalar_type(dtype)
 
-    if dtype == "float":
-        return f"vec{n}"
-    if dtype == "uint":
-        return f"uvec{n}"
-    elif dtype == "half":
+    if dtype == "half":
         return f"f16vec{n}"
-    elif dtype == "int":
-        return f"ivec{n}"
+    elif dtype == "float":
+        return f"vec{n}"
+    elif dtype == "double":
+        return f"f64vec{n}"
+    # integer dtype
     elif dtype == "int8":
         return f"i8vec{n}"
     elif dtype == "uint8":
         return f"u8vec{n}"
-    elif dtype == "bool":
-        return f"u8vec{n}"
+    elif dtype == "int16":
+        return f"i16vec{n}"
+    elif dtype == "uint16":
+        return f"u16vec{n}"
+    elif dtype == "int32" or dtype == "int":
+        return f"ivec{n}"
+    elif dtype == "uint32" or dtype == "uint":
+        return f"uvec{n}"
+    elif dtype == "int64":
+        return f"i64vec{n}"
+    elif dtype == "uint64":
+        return f"u64vec{n}"
 
     raise AssertionError(f"Invalid dtype: {dtype}")
 
@@ -360,20 +422,19 @@ def define_required_extensions(dtypes: Union[str, List[str]]):
     dtype_list = dtypes if isinstance(dtypes, list) else [dtypes]
 
     for dtype in dtype_list:
-        nbit = None
         glsl_type = None
         if dtype == "half":
-            nbit = "16bit"
             glsl_type = "float16"
-        elif dtype == "int16" or dtype == "uint16":
-            nbit = "16bit"
-            glsl_type = "int16"
-        elif dtype == "int8" or dtype == "uint8" or dtype == "bool":
-            nbit = "8bit"
+        elif dtype == "double":
+            glsl_type = "float64"
+        elif dtype in ["int8", "uint8"]:
             glsl_type = "int8"
+        elif dtype in ["int16", "uint16"]:
+            glsl_type = "int16"
+        elif dtype in ["int64", "uint64"]:
+            glsl_type = "int64"
 
-        if nbit is not None and glsl_type is not None:
-            out_str += f"#extension GL_EXT_shader_{nbit}_storage : require\n"
+        if glsl_type is not None:
             out_str += f"#extension GL_EXT_shader_explicit_arithmetic_types_{glsl_type} : require\n"
 
     return out_str
@@ -629,6 +690,8 @@ def generateVariantCombinations(
 
                     elif "VALUE" in value:
                         suffix = value.get("SUFFIX", value["VALUE"])
+                        if value["VALUE"] in ["int", "uint"]:
+                            raise ValueError(f"Use int32 or uint32 instead of {value['VALUE']}")
                         param_values.append((param_name, suffix, value["VALUE"]))
 
                     else:

backends/vulkan/runtime/graph/ops/glsl/arange.yaml

@@ -14,6 +14,6 @@ arange:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: int32
   shader_variants:
     - NAME: arange

backends/vulkan/runtime/graph/ops/glsl/avg_pool2d.yaml

@@ -13,6 +13,6 @@ avg_pool2d:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: int32
   shader_variants:
     - NAME: avg_pool2d

backends/vulkan/runtime/graph/ops/glsl/binary_op.yaml

@@ -17,7 +17,7 @@ binary_op:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: int32
   shader_variants:
     - NAME: binary_add
     - NAME: binary_sub

backends/vulkan/runtime/graph/ops/glsl/buffer_to_buffer.yaml

@@ -12,8 +12,10 @@ buffer_to_buffer:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: double
       - VALUE: int8
       - VALUE: uint8
+      - VALUE: int32
+      - VALUE: int64
   shader_variants:
     - NAME: buffer_to_buffer

backends/vulkan/runtime/graph/ops/glsl/buffer_to_nchw.yaml

@@ -13,9 +13,11 @@ buffer_to_nchw:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: double
       - VALUE: int8
       - VALUE: uint8
+      - VALUE: int32
+      - VALUE: int64
   shader_variants:
     - NAME: buffer_to_nchw
     - NAME: buffer_to_nchw_no_pc

backends/vulkan/runtime/graph/ops/glsl/copy_channel_offset.yaml

@@ -7,6 +7,6 @@ copy_channel_offset:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: int32
   shader_variants:
     - NAME: copy_channel_offset

backends/vulkan/runtime/graph/ops/glsl/copy_offset.yaml

@@ -7,7 +7,7 @@ copy_offset:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: int32
       - VALUE: int8
       - VALUE: uint8
     STORAGE:

backends/vulkan/runtime/graph/ops/glsl/copy_packed_dim_offset.yaml

@@ -7,6 +7,6 @@ copy_packed_dim_offset:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: int32
   shader_variants:
     - NAME: copy_packed_dim_offset

backends/vulkan/runtime/graph/ops/glsl/embedding.yaml

@@ -7,6 +7,6 @@ embedding:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: int32
   shader_variants:
     - NAME: embedding

backends/vulkan/runtime/graph/ops/glsl/flip.yaml

@@ -6,8 +6,10 @@ flip:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: double
       - VALUE: int8
       - VALUE: uint8
+      - VALUE: int32
+      - VALUE: int64
   shader_variants:
     - NAME: flip

backends/vulkan/runtime/graph/ops/glsl/image_to_nchw.yaml

@@ -14,9 +14,11 @@ image_to_nchw:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: double
       - VALUE: int8
       - VALUE: uint8
+      - VALUE: int32
+      - VALUE: int64
   shader_variants:
     - NAME: image_to_nchw_texture3d
     - NAME: image_to_nchw_texture2d

backends/vulkan/runtime/graph/ops/glsl/index_select.yaml

@@ -7,6 +7,6 @@ index_select:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: int32
   shader_variants:
     - NAME: index_select

backends/vulkan/runtime/graph/ops/glsl/index_select_channel.yaml

@@ -7,6 +7,6 @@ index_select_channel:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: int32
   shader_variants:
     - NAME: index_select_channel

backends/vulkan/runtime/graph/ops/glsl/nchw_to_buffer.yaml

@@ -13,9 +13,11 @@ nchw_to_buffer:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: double
       - VALUE: int8
       - VALUE: uint8
+      - VALUE: int32
+      - VALUE: int64
   shader_variants:
     - NAME: nchw_to_buffer
     - NAME: nchw_to_buffer_no_pc

backends/vulkan/runtime/graph/ops/glsl/nchw_to_image.glsl

@@ -87,5 +87,9 @@ void main() {
     return;
   }
 
-  write_texel(t_out, lpos_to_pos(lpos, axis_map), read_texel(tidx));
+  $if DTYPE == "double" and DTYPE == "int64":
+    VEC4_T texel = read_texel(tidx);
+    write_texel(t_out, lpos_to_pos(lpos, axis_map), texel);
+  $else:
+    write_texel(t_out, lpos_to_pos(lpos, axis_map), read_texel(tidx));
 }

backends/vulkan/runtime/graph/ops/glsl/nchw_to_image.yaml

@@ -14,9 +14,11 @@ nchw_to_image:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: double
       - VALUE: int8
       - VALUE: uint8
+      - VALUE: int32
+      - VALUE: int64
   shader_variants:
     - NAME: nchw_to_image_texture3d
     - NAME: nchw_to_image_texture2d

backends/vulkan/runtime/graph/ops/glsl/no_op.yaml

@@ -12,7 +12,7 @@ no_op:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: int32
       - VALUE: int8
       - VALUE: uint8
     STORAGE:

backends/vulkan/runtime/graph/ops/glsl/permute.yaml

@@ -7,6 +7,6 @@ permute:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: int32
   shader_variants:
     - NAME: permute

backends/vulkan/runtime/graph/ops/glsl/repeat.yaml

@@ -7,7 +7,7 @@ repeat:
     DTYPE:
       - VALUE: half
       - VALUE: float
-      - VALUE: int
+      - VALUE: int32
       - VALUE: int8
       - VALUE: uint8
   shader_variants: