[Vulkan] Pass typed fill values to full and make it resizable

backends/vulkan/op_registry.py

@@ -1612,6 +1612,10 @@ def register_full_cpp_ops():
     return OpFeatures(
         inputs_storage=utils.ANY_STORAGE,
         inputs_dtypes=utils.FP_INT_BOOL_T,
+        supports_resize=True,
+        are_node_inputs_supported_fn=lambda node: node.target
+        not in (exir_ops.edge.aten.full.default, exir_ops.edge.aten.full_like.default)
+        or is_scalar_value_supported(node.args[1], node.meta["val"].dtype),
     )
 
 

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

@@ -13,6 +13,14 @@ ${define_required_extensions("buffer", DTYPE)}
 #define PRECISION ${PRECISION}
 
 #define T ${buffer_scalar_type(DTYPE)}
+$if DTYPE == "half" or DTYPE == "float":
+  #define FILL_DTYPE float
+$elif DTYPE == "int32":
+  #define FILL_DTYPE int
+$elif DTYPE == "uint8":
+  #define FILL_DTYPE uint
+$else:
+  #error Unsupported full dtype
 
 ${define_active_storage_type("buffer")}
 
@@ -22,7 +30,7 @@ layout(std430) buffer;
 
 ${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")}
 ${layout_declare_ubo(B, "BufferMetadata", "outp")}
-${layout_declare_ubo(B, "float", "fill_value")}
+${layout_declare_ubo(B, "FILL_DTYPE", "fill_value")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 

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

@@ -13,16 +13,25 @@ ${define_required_extensions("texture3d", DTYPE)}
 #define PRECISION ${PRECISION}
 
 #define VEC4_T ${texel_load_type(DTYPE, "texture3d")}
+$if DTYPE == "half" or DTYPE == "float":
+  #define FILL_DTYPE float
+$elif DTYPE == "int32":
+  #define FILL_DTYPE int
+$elif DTYPE == "uint8":
+  #define FILL_DTYPE uint
+$else:
+  #error Unsupported full dtype
 
 ${define_active_storage_type("texture3d")}
 
 layout(std430) buffer;
 
 #include "indexing.glslh"
+#include "convert.glslh"
 
 ${layout_declare_tensor(B, "w", "t_out", DTYPE, "texture3d")}
 ${layout_declare_ubo(B, "TextureMetadata", "outp")}
-${layout_declare_ubo(B, "float", "fill_value")}
+${layout_declare_ubo(B, "FILL_DTYPE", "fill_value")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
@@ -37,6 +46,8 @@ void main() {
   }
 
   VEC4_T outtex = VEC4_T(fill_value);
+  $if DTYPE == "half":
+    outtex = round_to_half_rte(outtex);
 
   TensorIndex4D tidx =
       texture_pos_to_tensor4d_idx_simple(outp, pos, out_layout);

backends/vulkan/runtime/graph/ops/impl/Full.cpp

@@ -36,7 +36,17 @@ void add_full_node(
     const ValueRef size_or_in,
     const ValueRef fill_value,
     const ValueRef out) {
-  float fill_value_val = graph.extract_scalar<float>(fill_value);
+  vkapi::BufferBindInfo fill_value_buffer;
+  if (graph.dtype_of(out) == vkapi::kInt) {
+    fill_value_buffer =
+        graph.create_params_buffer(graph.extract_scalar<int32_t>(fill_value));
+  } else if (graph.dtype_of(out) == vkapi::kBool) {
+    const uint32_t value = graph.extract_scalar<bool>(fill_value);
+    fill_value_buffer = graph.create_params_buffer(value);
+  } else {
+    fill_value_buffer =
+        graph.create_params_buffer(graph.extract_scalar<float>(fill_value));
+  }
 
   std::string kernel_name("full");
   kernel_name.reserve(kShaderNameReserve);
@@ -52,7 +62,7 @@ void add_full_node(
       // Inputs and Outputs
       {{out, vkapi::kWrite}},
       // Shader params buffers
-      {graph.meta_ubo(out), graph.create_params_buffer(fill_value_val)},
+      {graph.meta_ubo(out), fill_value_buffer},
       // Push Constants
       {},
       // Specialization Constants

backends/vulkan/test/test_vulkan_dynamic.py

@@ -233,6 +233,27 @@ def forward(self, x):
                 edge = self._lower(model, inputs[0])
                 self._run(edge, model, inputs, atol=0, rtol=0)
 
+    def test_dynamic_full(self):
+        class Full(torch.nn.Module):
+            def forward(self, x):
+                return (
+                    torch.full(x.shape, 2.5),
+                    torch.zeros(x.shape),
+                    torch.ones(x.shape),
+                    torch.full_like(x, -1.5),
+                    torch.zeros_like(x),
+                    torch.ones_like(x),
+                )
+
+        for storage in (VkStorageType.TEXTURE_3D, VkStorageType.BUFFER):
+            with self.subTest(storage=storage):
+                model = Full()
+                inputs = [(torch.randn(2, s, 3),) for s in (16, 3, 31, 2, 16)]
+                edge = self._lower(
+                    model, inputs[0], ({1: Dim("s", min=2, max=32)},), storage
+                )
+                self._run(edge, model, inputs, atol=0, rtol=0)
+
     def test_dynamic_logical_not(self):
         class LogicalNot(torch.nn.Module):
             def forward(self, x):
@@ -308,6 +329,29 @@ def forward(self, x):
                 edge = self._lower(model, inputs[0], fully_delegated=False)
                 self._run(edge, model, inputs, atol=0, rtol=0, equal_nan=True)
 
+    def test_integer_fill_values(self):
+        class IntegerFill(torch.nn.Module):
+            def __init__(self, dtype):
+                super().__init__()
+                self.dtype = dtype
+
+            def forward(self, x):
+                return (
+                    torch.full(x.shape, 16777217, dtype=self.dtype),
+                    torch.full_like(x, -(2**31), dtype=self.dtype),
+                    torch.full(x.shape, 2**31 - 1, dtype=self.dtype),
+                )
+
+        inputs = [(torch.randn(2, s, 3),) for s in (7, 2, 15, 7)]
+        for dtype in (torch.int32, torch.int64):
+            for storage in (VkStorageType.TEXTURE_3D, VkStorageType.BUFFER):
+                with self.subTest(dtype=dtype, storage=storage):
+                    model = IntegerFill(dtype)
+                    edge = self._lower(
+                        model, inputs[0], ({1: Dim("s", min=2, max=16)},), storage
+                    )
+                    self._run(edge, model, inputs, atol=0, rtol=0)
+
     def test_integer_scalar_range_fallback(self):
         class LargeScalar(torch.nn.Module):
             def __init__(self, kind, value):
@@ -367,6 +411,85 @@ def forward(self, x):
                             )
                 self._run(edge, model, inputs, atol=0, rtol=0)
 
+    def test_integer_factories_without_downcasting(self):
+        class IntegerFactories(torch.nn.Module):
+            def __init__(self, dtype):
+                super().__init__()
+                self.dtype = dtype
+
+            def forward(self, x):
+                return (
+                    torch.scalar_tensor(16777217, dtype=self.dtype),
+                    torch.full(x.shape, 2**31 - 1, dtype=self.dtype),
+                    torch.full_like(x, -(2**31), dtype=self.dtype),
+                )
+
+        inputs = [(torch.zeros(3, s),) for s in (7, 2, 15, 7)]
+        for dtype in (torch.int32, torch.int64):
+            for downcast in (True, False):
+                with self.subTest(dtype=dtype, downcast=downcast):
+                    model = IntegerFactories(dtype)
+                    delegated = dtype == torch.int32 or downcast
+                    edge = self._lower(
+                        model,
+                        inputs[0],
+                        ({1: Dim("s", min=2, max=16)},),
+                        fully_delegated=delegated,
+                        downcast_64_bit=downcast,
+                    )
+                    self.assertEqual(bool(_vulkan_graphs(edge)), delegated)
+                    self._run(edge, model, inputs, atol=0, rtol=0)
+
+    def test_64_bit_inputs_without_downcasting(self):
+        class Input64Bit(torch.nn.Module):
+            def forward(self, x):
+                return (
+                    x + x,
+                    torch.full_like(x, 3, dtype=torch.int32),
+                    torch.ones(x.shape, dtype=torch.float32),
+                )
+
+        model = Input64Bit()
+        for dtype in (torch.int64, torch.float64):
+            with self.subTest(dtype=dtype):
+                inputs = [
+                    (torch.arange(3 * s, dtype=dtype).reshape(3, s),)
+                    for s in (7, 2, 15, 7)
+                ]
+                edge = self._lower(
+                    model,
+                    inputs[0],
+                    ({1: Dim("s", min=2, max=16)},),
+                    fully_delegated=False,
+                    downcast_64_bit=False,
+                )
+                graphs = _vulkan_graphs(edge)
+                self.assertTrue(graphs)
+                for graph in graphs:
+                    for value in graph.values:
+                        if isinstance(value.value, VkTensor):
+                            self.assertNotIn(
+                                value.value.datatype,
+                                (VkDataType.INT64, VkDataType.FLOAT64),
+                            )
+                self._run(edge, model, inputs, atol=0, rtol=0)
+
+    def test_bool_fill_values(self):
+        class BoolFill(torch.nn.Module):
+            def forward(self, x):
+                return (
+                    torch.full_like(x, 0.5, dtype=torch.bool),
+                    torch.full_like(x, -1.5, dtype=torch.bool),
+                    torch.full(x.shape, 0, dtype=torch.bool),
+                )
+
+        model = BoolFill()
+        inputs = [(torch.zeros(3, 7),)]
+        for storage in (VkStorageType.TEXTURE_3D, VkStorageType.BUFFER):
+            with self.subTest(storage=storage):
+                edge = self._lower(model, inputs[0], storage=storage)
+                self._run(edge, model, inputs, atol=0, rtol=0)
+
     def test_64_bit_fusion_inputs_without_downcasting(self):
         class SelectScalar(torch.nn.Module):
             def __init__(self, narrow):
@@ -636,6 +759,40 @@ def forward(self, x):
                         self.assertEqual(_vulkan_graphs(edge), [])
                         self._run(edge, model, [(x,)])
 
+    def test_fp16_scalar_rounding(self):
+        class CreateTensor(torch.nn.Module):
+            def __init__(self, value, scalar):
+                super().__init__()
+                self.value = value
+                self.scalar = scalar
+
+            def forward(self, x):
+                if self.scalar:
+                    return torch.scalar_tensor(self.value, dtype=x.dtype)
+                return torch.full_like(x, self.value)
+
+        x = torch.ones(3, 7, dtype=torch.float16)
+        for value in (
+            0.3,
+            -1.00075,
+            2**-24,
+            -(2**-24),
+            2**-25,
+            3 * 2**-25,
+            65519.0,
+            65520.0,
+            1e5,
+            -1e5,
+        ):
+            for scalar in (False, True):
+                if not scalar and abs(value) > 65504:
+                    continue
+                for storage in (VkStorageType.TEXTURE_3D, VkStorageType.BUFFER):
+                    with self.subTest(value=value, scalar=scalar, storage=storage):
+                        model = CreateTensor(value, scalar)
+                        edge = self._lower(model, (x,), storage=storage)
+                        self._run(edge, model, [(x,)], atol=0, rtol=0)
+
     def test_int32_buffer_reduction_shader_range(self):
         from executorch.extension.pybindings.portable_lib import (
             _load_for_executorch_from_buffer,