[Vulkan] Pass typed fill values to full and make it resizable (pytorch#23250)

The full shaders always received the fill value as a float, which loses
int32 fills above 2^24 and is ill-defined for bools. Full.cpp now passes
the fill value in the output's accumulator type. Both full shaders
explicitly declare the matching fill UBO type for each output variant,
independent of the accumulator-type helper. full, full_like, zeros, ones
and their _like variants are marked resizable so they can live inside
dynamic-shape delegates, and fill values the output dtype cannot
represent stay on CPU.

Part 11/15 of the Vulkan transformer and operator-conformance stack.
Depends on pytorch#23249; review against the selected base branch. Integration
PR: pytorch#23254.

Validation: The earlier MoltenVK/portable-kernel run had 6 passes and 2
warnings in 113.26s (0:01:53). For this follow-up, all 8 full shader
variants compiled with glslc; commit-hook lint and git diff --check
pass. Native tests have not been rerun since this edit; hardware and
SwiftShader CI are pending.

Recreates pytorch#23211 through ghstack. Prior review discussion remains on
that PR.

Authored with OpenAI Codex; split planned with Claude Code.


cc @SS-JIA @manuelcandales @digantdesai @cbilgin

Author: mergennachin

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,