Arm backend: Clean up grid_sampler custom payload

Refactor the grid_sampler_2d TOSA CUSTOM payload into a helper
module so the custom shader contract, constants, and validation
logic live in one place.

Emit the ML SDK Vulkan custom-shader attribute format, including
the converter-recognized domain, entry point, workgroup sizes,
shader language/code, tensor storage-buffer bindings, descriptor
sets, and Vulkan format metadata.

Add focused tests for payload round-trip behavior and for the
grid_sampler_2d rewrite pass.

Signed-off-by: Baris Demir <baris.demir@arm.com>
Change-Id: I8f13f91b97e86b3b32bf2c25479684bc09faeddc

Author: Baris Demir

backends/arm/scripts/generate_grid_sampler_spirv.py

@@ -0,0 +1,75 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import base64
+import shutil
+import subprocess  # nosec B404 - required to invoke the shader compiler.
+import tempfile
+from pathlib import Path
+
+
+SHADER_DIR = Path(__file__).resolve().parents[1] / "vgf" / "shaders"
+DEFAULT_SOURCE = SHADER_DIR / "grid_sampler.glsl"
+DEFAULT_OUTPUT = SHADER_DIR / "grid_sampler.spirv.b64"
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Compile the VGF grid_sampler GLSL shader to SPIR-V and write the "
+            "base64-encoded payload consumed by the ExecuTorch custom-shader "
+            "lowering."
+        )
+    )
+    parser.add_argument(
+        "--source",
+        type=Path,
+        default=DEFAULT_SOURCE,
+        help=f"GLSL source file. Defaults to {DEFAULT_SOURCE}",
+    )
+    parser.add_argument(
+        "--output",
+        type=Path,
+        default=DEFAULT_OUTPUT,
+        help=f"Base64 SPIR-V output file. Defaults to {DEFAULT_OUTPUT}",
+    )
+    parser.add_argument(
+        "--glslc",
+        default="glslc",
+        help="Path to glslc. Defaults to resolving glslc from PATH.",
+    )
+    return parser.parse_args()
+
+
+def _resolve_glslc(glslc: str) -> str:
+    resolved = shutil.which(glslc)
+    if resolved is None:
+        raise RuntimeError(
+            f"Could not find {glslc}. Install the Vulkan SDK or pass --glslc."
+        )
+    return resolved
+
+
+def _write_base64_spirv(spirv_path: Path, output_path: Path) -> None:
+    encoded = base64.b64encode(spirv_path.read_bytes()).decode("ascii")
+    output_path.write_text(encoded + "\n", encoding="utf-8")
+
+
+def main() -> None:
+    args = _parse_args()
+    glslc = _resolve_glslc(args.glslc)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        spirv_path = Path(tmpdir) / "grid_sampler.spirv"
+        subprocess.run(  # nosec B603 - glslc path is resolved explicitly.
+            [glslc, str(args.source), "-o", str(spirv_path)],
+            check=True,
+        )
+        _write_base64_spirv(spirv_path, args.output)
+
+
+if __name__ == "__main__":
+    main()

backends/arm/test/misc/test_custom_shader_payload.py

@@ -0,0 +1,79 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import base64
+
+import pytest
+from executorch.backends.arm.vgf.shaders.grid_sampler import (
+    build_grid_sampler_2d_payload,
+    decode_payload,
+    encode_payload,
+    GRID_SAMPLER_2D_SHADER_BINARY,
+    GRID_SAMPLER_2D_SHADER_ENTRY_POINT,
+    GRID_SAMPLER_2D_SHADER_LANGUAGE,
+    GRID_SAMPLER_2D_SHADER_SOURCE,
+    GRID_SAMPLER_2D_VK_FORMAT,
+    GRID_SAMPLER_2D_WORKGROUP_SIZES,
+)
+
+
+def test_grid_sampler_2d_custom_shader_payload_no_target_round_trip():
+    payload = build_grid_sampler_2d_payload(
+        interpolation_mode=0,
+        padding_mode=2,
+        align_corners=True,
+    )
+    decoded = decode_payload(encode_payload(payload))
+
+    assert decoded["entry_point"] == GRID_SAMPLER_2D_SHADER_ENTRY_POINT
+    assert decoded["workgroup_sizes"] == GRID_SAMPLER_2D_WORKGROUP_SIZES
+    assert decoded["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE
+    assert base64.b64decode(decoded["shader_code"])[:4] == b"\x03\x02\x23\x07"
+    assert decoded["input_0_type"] == "Tensor"
+    assert decoded["input_0_vkformat"] == GRID_SAMPLER_2D_VK_FORMAT
+    assert decoded["input_0_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER"
+    assert decoded["input_0_binding"] == 0
+    assert decoded["input_1_type"] == "Tensor"
+    assert decoded["input_1_vkformat"] == GRID_SAMPLER_2D_VK_FORMAT
+    assert decoded["input_1_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER"
+    assert decoded["input_1_binding"] == 1
+    assert decoded["output_0_type"] == "Tensor"
+    assert decoded["output_0_vkformat"] == GRID_SAMPLER_2D_VK_FORMAT
+    assert decoded["output_0_vkdescriptortype"] == "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER"
+    assert decoded["output_0_binding"] == 2
+
+
+def test_grid_sampler_2d_custom_shader_payload_no_target_uses_spirv():
+    payload = build_grid_sampler_2d_payload(
+        interpolation_mode=0,
+        padding_mode=0,
+        align_corners=False,
+    )
+
+    shader_binary = base64.b64decode(payload["shader_code"])
+
+    assert payload["shader_language"] == "SPIR-V"
+    assert shader_binary[:4] == b"\x03\x02\x23\x07"
+
+
+def test_grid_sampler_2d_custom_shader_payload_no_target_has_shader_resources():
+    assert GRID_SAMPLER_2D_SHADER_SOURCE == "grid_sampler.glsl"
+    assert GRID_SAMPLER_2D_SHADER_BINARY == "grid_sampler.spirv.b64"
+
+
+def test_grid_sampler_2d_custom_shader_payload_no_target_rejects_bad_modes():
+    with pytest.raises(ValueError, match="Unsupported interpolation_mode"):
+        build_grid_sampler_2d_payload(
+            interpolation_mode=99,
+            padding_mode=0,
+            align_corners=False,
+        )
+
+    with pytest.raises(ValueError, match="Unsupported padding_mode"):
+        build_grid_sampler_2d_payload(
+            interpolation_mode=0,
+            padding_mode=99,
+            align_corners=False,
+        )

backends/arm/test/passes/test_rewrite_grid_sampler_to_tosa_custom_pass.py

@@ -3,8 +3,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import json
-
 import executorch.backends.arm.tosa.dialect  # noqa: F401
 import torch
 import torch.nn.functional as F
@@ -15,6 +13,15 @@
 from executorch.backends.arm.vgf._passes.rewrite_grid_sampler_to_tosa_custom import (
     RewriteGridSamplerToTosaCustomPass,
 )
+from executorch.backends.arm.vgf.shaders.grid_sampler import (
+    CUSTOM_SHADER_DOMAIN_NAME,
+    decode_payload,
+    GRID_SAMPLER_2D_OPERATOR_NAME,
+    GRID_SAMPLER_2D_SHADER_ENTRY_POINT,
+    GRID_SAMPLER_2D_SHADER_LANGUAGE,
+    GRID_SAMPLER_2D_VK_FORMAT,
+    GRID_SAMPLER_2D_WORKGROUP_SIZES,
+)
 from executorch.exir import to_edge
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export import export
@@ -62,11 +69,22 @@ def test_rewrite_grid_sampler_to_tosa_custom_no_target():
     custom_node = next(
         node for node in nodes if node.target == exir_ops.backend.tosa.CUSTOM.default
     )
-    assert custom_node.kwargs["operator_name"] == "grid_sampler_2d"
-    assert custom_node.kwargs["domain_name"] == "arm.custom_shader"
+    assert custom_node.kwargs["operator_name"] == GRID_SAMPLER_2D_OPERATOR_NAME
+    assert custom_node.kwargs["domain_name"] == CUSTOM_SHADER_DOMAIN_NAME
 
-    payload = json.loads(
-        bytes(custom_node.kwargs["implementation_attrs"]).decode("utf-8")
-    )
-    assert payload["op"] == "grid_sampler_2d"
-    assert payload["shader"]["encoding"] == "placeholder"
+    payload = decode_payload(custom_node.kwargs["implementation_attrs"])
+    assert payload["entry_point"] == GRID_SAMPLER_2D_SHADER_ENTRY_POINT
+    assert payload["workgroup_sizes"] == GRID_SAMPLER_2D_WORKGROUP_SIZES
+    assert payload["shader_language"] == GRID_SAMPLER_2D_SHADER_LANGUAGE
+    assert payload["input_0_type"] == "Tensor"
+    assert payload["input_0_vkformat"] == GRID_SAMPLER_2D_VK_FORMAT
+    assert payload["input_0_binding"] == 0
+    assert payload["input_0_descriptorset"] == 0
+    assert payload["input_1_type"] == "Tensor"
+    assert payload["input_1_vkformat"] == GRID_SAMPLER_2D_VK_FORMAT
+    assert payload["input_1_binding"] == 1
+    assert payload["input_1_descriptorset"] == 0
+    assert payload["output_0_type"] == "Tensor"
+    assert payload["output_0_vkformat"] == GRID_SAMPLER_2D_VK_FORMAT
+    assert payload["output_0_binding"] == 2
+    assert payload["output_0_descriptorset"] == 0

backends/arm/vgf/_passes/rewrite_grid_sampler_to_tosa_custom.py

@@ -3,19 +3,24 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import json
 import operator
 from typing import Set, Type
 
 import torch
 from executorch.backends.arm._passes import ArmPass
 from executorch.backends.arm._passes.arm_pass_utils import create_node
 from executorch.backends.arm.tosa.dialect.ops.custom import register_fake_tosa
+from executorch.backends.arm.vgf.shaders.grid_sampler import (
+    build_grid_sampler_2d_payload,
+    CUSTOM_SHADER_DOMAIN_NAME,
+    encode_payload,
+    GRID_SAMPLER_2D_OPERATOR_NAME,
+)
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
 
 
-@register_fake_tosa("grid_sampler_2d")
+@register_fake_tosa(GRID_SAMPLER_2D_OPERATOR_NAME)
 def _grid_sampler_2d_custom_fake_impl(
     inputs, operator_name, domain_name, implementation_attrs
 ) -> list[torch.Tensor]:
@@ -46,16 +51,12 @@ class RewriteGridSamplerToTosaCustomPass(ArmPass):
     def _encode_payload(
         interpolation_mode: int, padding_mode: int, align_corners: bool
     ) -> list[int]:
-        payload = {
-            "version": 1,
-            "type": "arm_custom_shader",
-            "op": "grid_sampler_2d",
-            "interpolation_mode": int(interpolation_mode),
-            "padding_mode": int(padding_mode),
-            "align_corners": bool(align_corners),
-            "shader": {"encoding": "placeholder", "entry_point": "main"},
-        }
-        return list(json.dumps(payload, sort_keys=True).encode("utf-8"))
+        payload = build_grid_sampler_2d_payload(
+            interpolation_mode=interpolation_mode,
+            padding_mode=padding_mode,
+            align_corners=align_corners,
+        )
+        return encode_payload(payload)
 
     def call(self, graph_module):
         modified = False
@@ -83,8 +84,8 @@ def call(self, graph_module):
                     op_target=exir_ops.backend.tosa.CUSTOM.default,
                     args=([input_tensor, grid],),
                     kwargs={
-                        "operator_name": "grid_sampler_2d",
-                        "domain_name": "arm.custom_shader",
+                        "operator_name": GRID_SAMPLER_2D_OPERATOR_NAME,
+                        "domain_name": CUSTOM_SHADER_DOMAIN_NAME,
                         "implementation_attrs": implementation_attrs,
                     },
                     from_node=node,

backends/arm/vgf/shaders/__init__.py

@@ -0,0 +1,4 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.

backends/arm/vgf/shaders/grid_sampler.glsl

@@ -0,0 +1,20 @@
+#version 450
+
+layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;
+
+layout(set = 0, binding = 0) readonly buffer Input0 {
+    float input0[];
+};
+
+layout(set = 0, binding = 1) readonly buffer Input1 {
+    float input1[];
+};
+
+layout(set = 0, binding = 2) writeonly buffer Output0 {
+    float output0[];
+};
+
+void main() {
+    uint index = gl_GlobalInvocationID.x;
+    output0[index] = input0[index];
+}

backends/arm/vgf/shaders/grid_sampler.py

@@ -0,0 +1,93 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+from importlib.resources import files
+from typing import Any
+
+CUSTOM_SHADER_DOMAIN_NAME = "com.arm.VulkanCustomShader"
+GRID_SAMPLER_2D_OPERATOR_NAME = "torch.nn.functional.grid_sample"
+GRID_SAMPLER_2D_WORKGROUP_SIZES = [8, 8, 1]
+GRID_SAMPLER_2D_SHADER_ENTRY_POINT = "main"
+GRID_SAMPLER_2D_SHADER_LANGUAGE = "SPIR-V"
+GRID_SAMPLER_2D_VK_FORMAT = "VK_FORMAT_R32_SFLOAT"
+GRID_SAMPLER_2D_SHADER_SOURCE = "grid_sampler.glsl"
+GRID_SAMPLER_2D_SHADER_BINARY = "grid_sampler.spirv.b64"
+
+_INTERPOLATION_MODE_NAMES = {
+    0: "bilinear",
+    1: "nearest",
+    2: "bicubic",
+}
+_PADDING_MODE_NAMES = {
+    0: "zeros",
+    1: "border",
+    2: "reflection",
+}
+
+
+def _mode_name(
+    mode: int,
+    names: dict[int, str],
+    mode_kind: str,
+) -> str:
+    if mode not in names:
+        raise ValueError(
+            f"Unsupported {mode_kind} {mode} for {GRID_SAMPLER_2D_OPERATOR_NAME}"
+        )
+    return names[mode]
+
+
+def build_grid_sampler_2d_payload(
+    interpolation_mode: int,
+    padding_mode: int,
+    align_corners: bool,
+) -> dict[str, Any]:
+    _mode_name(
+        int(interpolation_mode),
+        _INTERPOLATION_MODE_NAMES,
+        "interpolation_mode",
+    )
+    _mode_name(
+        int(padding_mode),
+        _PADDING_MODE_NAMES,
+        "padding_mode",
+    )
+    shader_code = "".join(
+        files(__package__)
+        .joinpath(GRID_SAMPLER_2D_SHADER_BINARY)
+        .read_text(encoding="utf-8")
+        .split()
+    )
+
+    return {
+        "entry_point": GRID_SAMPLER_2D_SHADER_ENTRY_POINT,
+        "workgroup_sizes": GRID_SAMPLER_2D_WORKGROUP_SIZES,
+        "shader_language": GRID_SAMPLER_2D_SHADER_LANGUAGE,
+        "shader_code": shader_code,
+        "input_0_type": "Tensor",
+        "input_0_vkformat": GRID_SAMPLER_2D_VK_FORMAT,
+        "input_0_vkdescriptortype": "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER",
+        "input_0_binding": 0,
+        "input_0_descriptorset": 0,
+        "input_1_type": "Tensor",
+        "input_1_vkformat": GRID_SAMPLER_2D_VK_FORMAT,
+        "input_1_vkdescriptortype": "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER",
+        "input_1_binding": 1,
+        "input_1_descriptorset": 0,
+        "output_0_type": "Tensor",
+        "output_0_vkformat": GRID_SAMPLER_2D_VK_FORMAT,
+        "output_0_vkdescriptortype": "VK_DESCRIPTOR_TYPE_STORAGE_BUFFER",
+        "output_0_binding": 2,
+        "output_0_descriptorset": 0,
+    }
+
+
+def encode_payload(payload: dict[str, Any]) -> list[int]:
+    return list(json.dumps(payload, sort_keys=True).encode("utf-8"))
+
+
+def decode_payload(implementation_attrs: list[int]) -> dict[str, Any]:
+    return json.loads(bytes(implementation_attrs).decode("utf-8"))