diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index 466f9d69bde..105bcea89a6 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -327,6 +327,17 @@ def register_pow_tensor_scalar(): ) +@update_features(exir_ops.edge.aten.eq.Scalar) +def register_eq_scalar(): + return OpFeatures( + inputs_storage=utils.ANY_STORAGE, + inputs_dtypes=utils.FP_INT_T, + outputs_dtypes=utils.BOOL_T, + supports_resize=True, + supports_highdim=True, + ) + + # ============================================================================= # ToCopy.cpp # ============================================================================= diff --git a/backends/vulkan/runtime/gen_vulkan_spv.py b/backends/vulkan/runtime/gen_vulkan_spv.py index 69c87563bbd..45988df8d62 100644 --- a/backends/vulkan/runtime/gen_vulkan_spv.py +++ b/backends/vulkan/runtime/gen_vulkan_spv.py @@ -255,6 +255,61 @@ def texel_load_component_type(dtype: str, storage_type: str) -> str: return texel_component_type(dtype) +def get_higher_precision_dtype(dtype_a: str, dtype_b: str) -> str: + """Return the higher-precision of two dtypes, intended for picking the type + in which to perform a mixed-dtype computation (e.g. a tensor of one dtype + combined with a scalar of another). The result is the operand dtype that can + represent the other without loss for the value ranges these shaders handle. + + Ranking rule (highest first): + - The float family ALWAYS outranks the int family. So any float dtype + beats any int dtype (e.g. int32 + float -> float), because integer + values in these shaders are small enough to be exactly representable in + the float compute type. + - Within the float family: double > float > half. + - Within the int family: rank by bit width + (int64/uint64 > int32/uint32 > int16/uint16 > int8/uint8/bool). + - Signed-vs-unsigned tie at the same bit width resolves to the SIGNED + dtype (matches PyTorch's same-width promotion preference and keeps the + compute type able to hold negative operands). + + Dtype aliases are normalized first: int<->int32, uint<->uint32, and bool is + treated as the lowest-ranked 8-bit integer. + """ + alias_map = {"int": "int32", "uint": "uint32"} + a = alias_map.get(dtype_a, dtype_a) + b = alias_map.get(dtype_b, dtype_b) + + if a == b: + return a + + # (family_rank, bit_width, signed_rank). Higher tuple wins. family_rank 1 for + # the float family ensures it always outranks the int family (family_rank 0). + # signed_rank breaks same-width int ties in favor of signed. + float_ranks = {"half": 16, "float": 32, "double": 64} + int_ranks = { + "bool": 8, + "uint8": 8, + "int8": 8, + "uint16": 16, + "int16": 16, + "uint32": 32, + "int32": 32, + "uint64": 64, + "int64": 64, + } + + def rank(dtype: str) -> tuple: + if dtype in float_ranks: + return (1, float_ranks[dtype], 0) + if dtype in int_ranks: + signed_rank = 0 if dtype in ("bool",) or dtype.startswith("uint") else 1 + return (0, int_ranks[dtype], signed_rank) + raise AssertionError(f"Cannot rank precision of dtype: {dtype}") + + return a if rank(a) >= rank(b) else b + + def get_access_qualifier(access_type: Optional[str]) -> str: if access_type is None: return "" @@ -497,6 +552,7 @@ def define_required_extensions(storage_type: str, dtypes: Union[str, List[str]]) "texel_component_type": texel_component_type, "texel_load_type": texel_load_type, "texel_load_component_type": texel_load_component_type, + "get_higher_precision_dtype": get_higher_precision_dtype, "layout_declare_buffer": layout_declare_buffer, "layout_declare_image": layout_declare_image, "layout_declare_sampler": layout_declare_sampler, diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_op_defs.glslh b/backends/vulkan/runtime/graph/ops/glsl/binary_op_defs.glslh index e2bdec703ca..83ba3e4e0d3 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/binary_op_defs.glslh +++ b/backends/vulkan/runtime/graph/ops/glsl/binary_op_defs.glslh @@ -19,38 +19,44 @@ // - Uses standard pow() for x > 0 // +// Operands are evaluated in the promoted compute type COMPUTE_T (and its vector +// form COMPUTE_VEC4_T), which the including shader sets from +// get_higher_precision_dtype(DTYPE, SCALAR_VALUE_TYPE) so mixed tensor/scalar +// dtype combinations compute without loss before narrowing back to the output +// dtype. + // Scalar overload -T power_of(T x, T y) { +COMPUTE_T power_of(COMPUTE_T x, COMPUTE_T y) { if (x == 0.0) { // Handle 0^y: 0^0 = 1, 0^y = 0 for y > 0 - return (y == 0.0) ? T(1.0) : T(0.0); + return (y == 0.0) ? COMPUTE_T(1.0) : COMPUTE_T(0.0); } // Use absolute value to avoid undefined behavior - float result = pow(abs(x), y); + float result = pow(abs(float(x)), float(y)); // For negative bases with odd integer exponents, preserve the negative sign if (x < 0.0) { - float int_y = round(y); - if (abs(y - int_y) < 1e-5 && int(int_y) % 2 == 1) { + float int_y = round(float(y)); + if (abs(float(y) - int_y) < 1e-5 && int(int_y) % 2 == 1) { result = -result; } } - return T(result); + return COMPUTE_T(result); } -#ifdef VEC4_T +#ifdef COMPUTE_VEC4_T // Vector overload -VEC4_T power_of(VEC4_T x, VEC4_T y) { - VEC4_T result; +COMPUTE_VEC4_T power_of(COMPUTE_VEC4_T x, COMPUTE_VEC4_T y) { + COMPUTE_VEC4_T result; for (int i = 0; i < 4; i++) { result[i] = power_of(x[i], y[i]); } return result; } -#endif // VEC4_T +#endif // COMPUTE_VEC4_T #endif // BINARY_OP_DEFS_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_buffer.glsl index 9e3a35bf4f1..1de98128159 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_buffer.glsl @@ -6,15 +6,32 @@ * LICENSE file in the root directory of this source tree. */ +// Binary comparison ops write a bool/uint8 output dtype, which differs from +// the input dtype. IS_COMPARISON_OP is set explicitly per shader variant in the +// .yaml. + #version 450 core +$PROMOTED_DTYPE = get_higher_precision_dtype(DTYPE, SCALAR_VALUE_TYPE) + ${define_required_extensions(STORAGE, DTYPE)} +${define_required_extensions(STORAGE, PROMOTED_DTYPE)} +${define_explicit_type_extensions(SCALAR_VALUE_TYPE)} +${define_explicit_type_extensions(PROMOTED_DTYPE)} +$if IS_COMPARISON_OP: + ${define_required_extensions(STORAGE, "uint8")} #define PRECISION ${PRECISION} #define NAME ${VARIANT_NAME} #define T ${buffer_scalar_type(DTYPE)} +#define SCALAR_T ${buffer_scalar_type(SCALAR_VALUE_TYPE)} +#define COMPUTE_T ${buffer_scalar_type(PROMOTED_DTYPE)} +$if IS_COMPARISON_OP: + #define OUT_T ${buffer_scalar_type("uint8")} +$else: + #define OUT_T ${buffer_scalar_type(DTYPE)} #define op(X, Y) ${OPERATOR} @@ -24,19 +41,24 @@ layout(std430) buffer; #include "indexing.glslh" -${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)} +$if IS_COMPARISON_OP: + ${layout_declare_tensor(B, "w", "t_out", "uint8", STORAGE)} +$else: + ${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)} + ${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)} ${layout_declare_ubo(B, "BufferMetadata", "outp")} ${layout_declare_ubo(B, "BufferMetadata", "inp")} layout(push_constant) uniform restrict Block { - float scalar_value; + SCALAR_T scalar_value; }; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; -#include "binary_op_defs.glslh" +$if not IS_COMPARISON_OP: + #include "binary_op_defs.glslh" void main() { const uint out_bufi = gl_GlobalInvocationID.x; @@ -44,5 +66,6 @@ void main() { return; } - t_out[out_bufi] = T(op(t_in[out_bufi], T(scalar_value))); + t_out[out_bufi] = + OUT_T(op(COMPUTE_T(t_in[out_bufi]), COMPUTE_T(scalar_value))); } diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_buffer.yaml index b818132cf9b..8a9e3dff3f1 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_buffer.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_buffer.yaml @@ -7,14 +7,20 @@ binary_scalar_buffer: parameter_names_with_default_values: OPERATOR: power_of(X, Y) - NDIM: 3 + IS_COMPARISON_OP: false DTYPE: float - PACKING: C_packed + SCALAR_VALUE_TYPE: float STORAGE: buffer generate_variant_forall: - DTYPE: - - VALUE: half - - VALUE: float - - VALUE: int32 + combination: + parameter_names: [DTYPE, SCALAR_VALUE_TYPE] + combos: + - parameter_values: [half, float] + - parameter_values: [float, float] + - parameter_values: [int32, int32] + - parameter_values: [int32, float] shader_variants: - NAME: pow_scalar_buffer + - NAME: eq_scalar_buffer + OPERATOR: X == Y + IS_COMPARISON_OP: true diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_texture.glsl b/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_texture.glsl index 651dfdd7b5d..acd523a67ba 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_texture.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_texture.glsl @@ -6,9 +6,20 @@ * LICENSE file in the root directory of this source tree. */ +// Binary comparison ops write a bool/uint8 output dtype, which differs from +// the input dtype. IS_COMPARISON_OP is set explicitly per shader variant in the +// .yaml. + #version 450 core +$PROMOTED_DTYPE = get_higher_precision_dtype(DTYPE, SCALAR_VALUE_TYPE) + ${define_required_extensions(STORAGE, DTYPE)} +${define_required_extensions(STORAGE, PROMOTED_DTYPE)} +${define_explicit_type_extensions(SCALAR_VALUE_TYPE)} +${define_explicit_type_extensions(PROMOTED_DTYPE)} +$if IS_COMPARISON_OP: + ${define_required_extensions(STORAGE, "uint8")} #define PRECISION ${PRECISION} @@ -16,6 +27,13 @@ ${define_required_extensions(STORAGE, DTYPE)} #define VEC4_T ${texel_load_type(DTYPE, STORAGE)} #define T ${texel_load_component_type(DTYPE, STORAGE)} +#define SCALAR_T ${buffer_scalar_type(SCALAR_VALUE_TYPE)} +#define COMPUTE_VEC4_T ${texel_load_type(PROMOTED_DTYPE, STORAGE)} +#define COMPUTE_T ${texel_load_component_type(PROMOTED_DTYPE, STORAGE)} +$if IS_COMPARISON_OP: + #define VEC4_OUT_T ${texel_load_type("uint8", STORAGE)} +$else: + #define VEC4_OUT_T VEC4_T #define op(X, Y) ${OPERATOR} @@ -25,19 +43,24 @@ layout(std430) buffer; #include "indexing.glslh" -${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)} +$if IS_COMPARISON_OP: + ${layout_declare_tensor(B, "w", "t_out", "uint8", STORAGE)} +$else: + ${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)} + ${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)} ${layout_declare_ubo(B, "TextureMetadata", "outp")} ${layout_declare_ubo(B, "TextureMetadata", "inp")} layout(push_constant) uniform restrict Block { - float scalar_value; + SCALAR_T scalar_value; }; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; -#include "binary_op_defs.glslh" +$if not IS_COMPARISON_OP: + #include "binary_op_defs.glslh" void main() { const ivec3 pos = ivec3(gl_GlobalInvocationID); @@ -47,7 +70,8 @@ void main() { } VEC4_T in_texel = texelFetch(t_in, pos, 0); - VEC4_T out_texel = VEC4_T(op(in_texel, VEC4_T(scalar_value))); + VEC4_OUT_T out_texel = VEC4_OUT_T( + op(COMPUTE_VEC4_T(in_texel), COMPUTE_VEC4_T(scalar_value))); imageStore(t_out, pos, out_texel); } diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_texture.yaml b/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_texture.yaml index 3e731bf7a15..b5dfe31f350 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_texture.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/binary_scalar_texture.yaml @@ -7,14 +7,20 @@ binary_scalar_texture: parameter_names_with_default_values: OPERATOR: power_of(X, Y) - NDIM: 3 + IS_COMPARISON_OP: false DTYPE: float - PACKING: C_packed + SCALAR_VALUE_TYPE: float STORAGE: texture3d generate_variant_forall: - DTYPE: - - VALUE: half - - VALUE: float - - VALUE: int32 + combination: + parameter_names: [DTYPE, SCALAR_VALUE_TYPE] + combos: + - parameter_values: [half, float] + - parameter_values: [float, float] + - parameter_values: [int32, int32] + - parameter_values: [int32, float] shader_variants: - NAME: pow_scalar_texture3d + - NAME: eq_scalar_texture3d + OPERATOR: equal(X, Y) + IS_COMPARISON_OP: true diff --git a/backends/vulkan/runtime/graph/ops/impl/BinaryScalarOp.cpp b/backends/vulkan/runtime/graph/ops/impl/BinaryScalarOp.cpp index 15553706494..f804725e828 100644 --- a/backends/vulkan/runtime/graph/ops/impl/BinaryScalarOp.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/BinaryScalarOp.cpp @@ -16,8 +16,47 @@ #include +#include + +#include + namespace vkcompute { +namespace { + +vkapi::ScalarType resolve_scalar_extract_dtype( + ComputeGraph& graph, + const ValueRef scalar, + const vkapi::ScalarType tensor_dtype) { + // For float tensors, ensure that the scalar argument is extracted as a float + // to avoid having to generate additional shader variants for float/half + // tensor + int scalar. + if (tensor_dtype == vkapi::kFloat || tensor_dtype == vkapi::kHalf) { + return vkapi::kFloat; + } + if (graph.val_is_bool(scalar) || graph.val_is_symint(scalar)) { + return vkapi::kInt; + } + return graph.dtype_of(scalar); +} + +int32_t extract_int32_scalar(ComputeGraph& graph, const ValueRef scalar) { + if (graph.val_is_int(scalar)) { + return utils::safe_downcast(graph.get_int(scalar)); + } + if (graph.val_is_bool(scalar)) { + return graph.get_bool(scalar) ? 1 : 0; + } + if (graph.val_is_symint(scalar)) { + return graph.read_symint(scalar); + } + VK_THROW( + "Expected int, bool, or SymInt scalar, got: ", + graph.get_val_type(scalar)); +} + +} // namespace + void resize_binary_scalar_op_node( ComputeGraph* graph, const std::vector& args, @@ -39,14 +78,25 @@ void add_binary_scalar_op_node( const std::string& op_name) { ValueRef arg = prepack_standard_like(graph, in, out, true); - // Extract scalar value - float scalar_val = graph.extract_scalar(scalar); + const vkapi::ScalarType tensor_dtype = graph.dtype_of(in); + const vkapi::ScalarType scalar_dtype = + resolve_scalar_extract_dtype(graph, scalar, tensor_dtype); + std::vector push_constants; + if (scalar_dtype == vkapi::kInt) { + const int32_t scalar_val = extract_int32_scalar(graph, scalar); + push_constants.emplace_back(&scalar_val, sizeof(scalar_val)); + } else if (scalar_dtype == vkapi::kFloat) { + const float scalar_val = graph.extract_scalar(scalar); + push_constants.emplace_back(&scalar_val, sizeof(scalar_val)); + } else { + VK_THROW("Unsupported tensor-scalar op scalar dtype: ", scalar_dtype); + } - // Pick shader std::string kernel_name = op_name + "_scalar"; kernel_name.reserve(kShaderNameReserve); add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); - add_dtype_suffix(kernel_name, graph.dtype_of(in)); + add_dtype_suffix(kernel_name, tensor_dtype); + add_dtype_suffix(kernel_name, scalar_dtype); vkapi::ParamsBindList param_ubos = {graph.meta_ubo(out), graph.meta_ubo(in)}; @@ -60,7 +110,7 @@ void add_binary_scalar_op_node( // Shader params buffers param_ubos, // Push Constants - {PushConstantDataInfo(&scalar_val, sizeof(scalar_val))}, + push_constants, // Specialization Constants {}, // Resize Args @@ -73,8 +123,13 @@ void pow_tensor_scalar(ComputeGraph& graph, const std::vector& args) { return add_binary_scalar_op_node(graph, args[0], args[1], args[2], "pow"); } +void eq_tensor_scalar(ComputeGraph& graph, const std::vector& args) { + return add_binary_scalar_op_node(graph, args[0], args[1], args[2], "eq"); +} + REGISTER_OPERATORS { VK_REGISTER_OP(aten.pow.Tensor_Scalar, pow_tensor_scalar); + VK_REGISTER_OP(aten.eq.Scalar, eq_tensor_scalar); } } // namespace vkcompute diff --git a/backends/vulkan/test/op_tests/cases.py b/backends/vulkan/test/op_tests/cases.py index 7a3b6943653..7e08bda27e3 100644 --- a/backends/vulkan/test/op_tests/cases.py +++ b/backends/vulkan/test/op_tests/cases.py @@ -2217,3 +2217,36 @@ def get_pow_tensor_scalar_inputs(): ] test_suite.dtypes = ["at::kFloat"] return test_suite + + +@register_test_suite("aten.eq.Scalar") +def get_eq_scalar_inputs(): + # Scalars are chosen to fall within the make_seq_tensor range (1..numel), + # so each case exercises a genuine mix of equal / not-equal elements rather + # than a trivially all-false comparison. + test_suite = VkTestSuite( + [ + ((M1,), 5), + ((M2, M1), 100), + ((S1, M1, M2), 1000), + ((S1, S2, S2, M2), 2000), + ((S, S1, S2), 50), + ((M1, M2), 700), + ((S1, S2), 20), + # Int tensor (dtype below) vs non-integer float scalar exercises the + # mixed int32/float shader variant: comparing in the promoted float + # type yields all-false (no integer equals 3.5). + ((M1,), 3.5), + ] + ) + test_suite.storage_types = [ + "utils::kBuffer", + "utils::kTexture3D", + ] + test_suite.layouts = [ + "utils::kWidthPacked", + "utils::kChannelsPacked", + ] + test_suite.dtypes = ["at::kInt", "at::kFloat"] + test_suite.data_gen = "make_seq_tensor" + return test_suite diff --git a/backends/vulkan/test/op_tests/utils/gen_benchmark_vk.py b/backends/vulkan/test/op_tests/utils/gen_benchmark_vk.py index 07d355c92de..24f3919f065 100644 --- a/backends/vulkan/test/op_tests/utils/gen_benchmark_vk.py +++ b/backends/vulkan/test/op_tests/utils/gen_benchmark_vk.py @@ -198,6 +198,19 @@ def generate_benchmark_fixture(self) -> str: }} }} +ValueRef add_scalar_to_graph(ComputeGraph& graph, const at::Scalar& scalar) {{ + if (scalar.isBoolean()) {{ + return graph.add_scalar(scalar.toBool()); + }} + if (scalar.isIntegral(/*includeBool=*/false)) {{ + return graph.add_scalar(scalar.toLong()); + }} + if (scalar.isFloatingPoint()) {{ + return graph.add_scalar(scalar.toDouble()); + }} + VK_THROW("Unsupported at::Scalar!"); +}} + at::Tensor make_casted_randint_tensor( std::vector sizes, at::ScalarType dtype = at::kFloat, diff --git a/backends/vulkan/test/op_tests/utils/gen_computegraph.py b/backends/vulkan/test/op_tests/utils/gen_computegraph.py index 507719b8555..b6374a890e5 100644 --- a/backends/vulkan/test/op_tests/utils/gen_computegraph.py +++ b/backends/vulkan/test/op_tests/utils/gen_computegraph.py @@ -476,8 +476,8 @@ def create_value_for( # noqa: C901 ret_str += f"from_at_scalartype({ref.src_cpp_name}.scalar_type()), " ret_str += f"{ref.src_cpp_name}.const_data_ptr()); \n" elif ref.src_cpp_type == AT_SCALAR: - # TODO(ssjia): generalize this to work with all scalar types - ret_str += f"add_scalar({ref.src_cpp_name}.toDouble()); \n" + ret_str = f"{cpp_type} {ref.name} = " + ret_str += f"add_scalar_to_graph(*{self.graph}, {ref.src_cpp_name}); \n" elif ref.src_cpp_type == AT_INT_ARRAY_REF: ret_str += f"add_scalar_list({ref.src_cpp_name}.vec()); \n" elif ref.src_cpp_type == BOOL: diff --git a/backends/vulkan/test/op_tests/utils/gen_correctness_vk.py b/backends/vulkan/test/op_tests/utils/gen_correctness_vk.py index 08bc502f964..83ff19100b2 100644 --- a/backends/vulkan/test/op_tests/utils/gen_correctness_vk.py +++ b/backends/vulkan/test/op_tests/utils/gen_correctness_vk.py @@ -129,6 +129,19 @@ def gen_parameterization(self) -> str: } } +ValueRef add_scalar_to_graph(ComputeGraph& graph, const at::Scalar& scalar) { + if (scalar.isBoolean()) { + return graph.add_scalar(scalar.toBool()); + } + if (scalar.isIntegral(/*includeBool=*/false)) { + return graph.add_scalar(scalar.toLong()); + } + if (scalar.isFloatingPoint()) { + return graph.add_scalar(scalar.toDouble()); + } + VK_THROW("Unsupported at::Scalar!"); +} + #ifdef USE_VULKAN_FP16_INFERENCE bool check_close(at::Tensor& t1, at::Tensor& t2, float rtol=1e-2, float atol=1e-2) { #else