Generate int variants for binary_ops (#2472)
Summary:
Pull Request resolved: https://github.com/pytorch/executorch/pull/2472
Add int variant for binary_ops shader.
Reviewed By: SS-JIA, jorgep31415
Differential Revision: D54963877
fbshipit-source-id: 6a156579b84ccf9a69be4f1a406a4d8af2464408
diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_op.glsl b/backends/vulkan/runtime/graph/ops/glsl/binary_op.glsl
index 69b87b1..acb672a 100644
--- a/backends/vulkan/runtime/graph/ops/glsl/binary_op.glsl
+++ b/backends/vulkan/runtime/graph/ops/glsl/binary_op.glsl
@@ -52,21 +52,21 @@
}
ivec4 in_coord = out_coord_to_in_coord(coord, in_sizes.data);
- vec4 in_texel = texelFetch(
+ ${VEC4_T[DTYPE]} in_texel = ${VEC4_T[DTYPE]}(texelFetch(
image_in,
COORD_TO_POS_${PACKING}(in_coord, in_sizes.data),
- 0);
+ 0));
ivec4 other_coord = out_coord_to_in_coord(coord, other_sizes.data);
- vec4 other_texel = texelFetch(
+ ${VEC4_T[DTYPE]} other_texel = ${VEC4_T[DTYPE]}(texelFetch(
image_other,
COORD_TO_POS_${PACKING}(other_coord, other_sizes.data),
- 0);
+ 0));
// Detect broadcasting
if (PACKED_DIM_${PACKING}(other_sizes.data) < PACKED_DIM_${PACKING}(in_sizes.data)) {
other_texel = other_texel.xxxx;
}
- imageStore(image_out, pos, OP(in_texel, other_texel, alpha.data));
+ imageStore(image_out, pos, ${VEC4_T[DTYPE]}(OP(in_texel, other_texel, alpha.data)));
}
diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_op.yaml b/backends/vulkan/runtime/graph/ops/glsl/binary_op.yaml
index 720ec7f..28f65ee 100644
--- a/backends/vulkan/runtime/graph/ops/glsl/binary_op.yaml
+++ b/backends/vulkan/runtime/graph/ops/glsl/binary_op.yaml
@@ -23,6 +23,8 @@
SUFFIX: half
- VALUE: float
SUFFIX: float
+ - VALUE: int
+ SUFFIX: int
shader_variants:
- NAME: binary_add
- NAME: binary_sub
