From bebd775373a8b4f22bdfca8ffbaac3bd99b589b7 Mon Sep 17 00:00:00 2001 From: Mergen Nachin Date: Tue, 29 Sep 2026 11:02:06 -0400 Subject: [PATCH] [INITIAL] Recreate the Vulkan transformer and conformance stack with ghstack [ghstack-poisoned] --- backends/vulkan/op_registry.py | 4 + .../runtime/graph/ops/glsl/full_buffer.glsl | 2 +- .../runtime/graph/ops/glsl/full_texture.glsl | 5 +- .../vulkan/runtime/graph/ops/impl/Full.cpp | 14 +- backends/vulkan/test/test_vulkan_dynamic.py | 157 ++++++++++++++++++ 5 files changed, 178 insertions(+), 4 deletions(-) diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index da7d8f0cde0..b5ba2774e8e 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1594,6 +1594,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), ) diff --git a/backends/vulkan/runtime/graph/ops/glsl/full_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/full_buffer.glsl index 8243d13325c..075ccc51488 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/full_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/full_buffer.glsl @@ -22,7 +22,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, accum_scalar_type(DTYPE), "fill_value")} layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; diff --git a/backends/vulkan/runtime/graph/ops/glsl/full_texture.glsl b/backends/vulkan/runtime/graph/ops/glsl/full_texture.glsl index 0124ac2d0b9..b23172525c2 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/full_texture.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/full_texture.glsl @@ -19,10 +19,11 @@ ${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, accum_scalar_type(DTYPE), "fill_value")} layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; @@ -37,6 +38,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); diff --git a/backends/vulkan/runtime/graph/ops/impl/Full.cpp b/backends/vulkan/runtime/graph/ops/impl/Full.cpp index e5e577da7ef..7c29e362984 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Full.cpp +++ b/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(fill_value); + vkapi::BufferBindInfo fill_value_buffer; + if (graph.dtype_of(out) == vkapi::kInt) { + fill_value_buffer = + graph.create_params_buffer(graph.extract_scalar(fill_value)); + } else if (graph.dtype_of(out) == vkapi::kBool) { + const uint32_t value = graph.extract_scalar(fill_value); + fill_value_buffer = graph.create_params_buffer(value); + } else { + fill_value_buffer = + graph.create_params_buffer(graph.extract_scalar(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 diff --git a/backends/vulkan/test/test_vulkan_dynamic.py b/backends/vulkan/test/test_vulkan_dynamic.py index a44f5f1a253..8e10060b426 100644 --- a/backends/vulkan/test/test_vulkan_dynamic.py +++ b/backends/vulkan/test/test_vulkan_dynamic.py @@ -228,6 +228,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): @@ -303,6 +324,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): @@ -362,6 +406,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): @@ -592,6 +715,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_reduction_special_values(self): class Reduce(torch.nn.Module): def __init__(self, op):