[Vulkan] scalar_tensor uploads float bytes to integer scalar shaders

[mergennachin]

🐛 Describe the bug

While implementing the scalar-tensor delegation fixes in #23156, I found a separate runtime parameter-type mismatch in ScalarTensor.cpp.

The runtime always extracts the scalar as float and uploads a float parameter buffer, but selects the shader's scalar parameter type from graph.dtype_of(scalar_in). An integer literal therefore selects an integer parameter declaration while receiving float bytes. For example, 3.0f has the bit pattern 0x40400000, which an int32 parameter reads as 1077936128.

Converting every integer scalar through float also loses precision for exactly representable int32 values above 2**24, such as 16777217.

The relevant code is in ScalarTensor.cpp, and the parameter declaration comes from scalar_tensor.glsl.

This becomes reachable through ordinary exported models after fixing the raw ATen registration and serialized operator name described in item 2 of #23156. Without those exporter fixes, the scalar tensors fall back to CPU and hide this runtime defect.

The following regression cases exercise mixed scalar types and integer precision. They require a Vulkan-enabled runtime and the scalar-tensor exporter fixes from #23156:

import torch
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
from executorch.exir import to_edge_transform_and_lower
from executorch.extension.pybindings.portable_lib import (
    _load_for_executorch_from_buffer,
)


class WhereScalars(torch.nn.Module):
    def __init__(self, positive, negative):
        super().__init__()
        self.positive = positive
        self.negative = negative

    def forward(self, x):
        return torch.where(x, self.positive, self.negative)


inputs = (torch.tensor([True, False, True, False]),)
for positive, negative in ((3, -7.0), (3.0, -7), (16777217, -7)):
    model = WhereScalars(positive, negative)
    edge = to_edge_transform_and_lower(
        torch.export.export(model, inputs),
        partitioner=[VulkanPartitioner()],
    )
    program_buffer = edge.to_executorch().buffer
    module = _load_for_executorch_from_buffer(program_buffer)
    actual = module.run_method("forward", inputs)[0]
    torch.testing.assert_close(actual, model(*inputs), atol=0, rtol=0)

The parameter buffer's type must match the shader declaration. Integer outputs should preserve the integer value without an intermediate float conversion. A local fix using int32 parameters for integer outputs and float parameters for floating-point outputs passes all three cases on MoltenVK. The surrounding Vulkan regression suite also passes: 36 tests, including dynamic eager attention and SDPA.

Versions

ExecuTorch base commit 0c7ce72758c0cbb832c5ac5970626beb347ca745, with local changes for #23156; PyTorch 2.13.0; Python 3.11; macOS arm64 on Apple M1 Pro using MoltenVK.

Drafted with OpenAI Codex. The parameter-type diagnosis is based on the runtime source; the local fix and regression cases were executed on the GPU.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin