diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index eed287b3a08..bd61c0ecb4c 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -208,6 +208,22 @@ def register_copy_op(): # ============================================================================= +def _check_clamp_bounds_are_static(node: torch.fx.Node) -> bool: + """Only support clamp when both bounds are static. + + get_val_or_inf() in UnaryOp.cpp reads each bound with + extract_scalar() when the node is BUILT and bakes the result into the + dispatch. A bound derived from a dynamic dimension is therefore never + refreshed, and the op silently computes against a stale limit -- no error is + raised. This is especially damaging when clamp is applied to index tensors, + where a wrong limit silently reorders or drops data downstream. + """ + for bound in node.args[1:3]: + if isinstance(bound, torch.fx.Node): + return False + return True + + @update_features( [ exir_ops.edge.aten.abs.default, @@ -234,6 +250,9 @@ def register_unaryop_cpp_ops(): inputs_storage=utils.ANY_STORAGE, inputs_dtypes=utils.FP_T, supports_resize=True, + # hardtanh/hardshrink/leaky_relu read their bounds through the same + # build-time get_val_or_inf() path as clamp. + are_node_inputs_supported_fn=_check_clamp_bounds_are_static, ) @@ -243,6 +262,7 @@ def register_clamp(): inputs_storage=utils.ANY_STORAGE, inputs_dtypes=utils.FP_INT_T, supports_resize=True, + are_node_inputs_supported_fn=_check_clamp_bounds_are_static, ) @@ -1508,12 +1528,30 @@ def register_arange(): # ============================================================================= +def _check_pad_is_static(node: torch.fx.Node) -> bool: + """Only support constant_pad_nd when the pad amounts are static. + + A symbolic pad list is serialized as a VALUELIST rather than an INTLIST, and + Pad.cpp reads it with get_int_list(), which throws "Expected value to have + type IntList, got VALUELIST instead". Separately, + add_constant_pad_nd_node() bakes the amounts into a params buffer at build + time, so a pad derived from a dynamic dim would use stale values even if the + list were read symbolically. Decline the node until the padding is plumbed + through as a symint. + """ + pad = node.args[1] + if not isinstance(pad, (list, tuple)): + return False + return all(isinstance(p, int) for p in pad) + + @update_features(exir_ops.edge.aten.constant_pad_nd.default) def register_constant_pad_nd(): return OpFeatures( inputs_storage=utils.ANY_STORAGE, inputs_dtypes=utils.FP_INT_BOOL_T, supports_resize=True, + are_node_inputs_supported_fn=_check_pad_is_static, ) @@ -1735,6 +1773,21 @@ def register_embedding_q4gsw(): # ============================================================================= +def _check_batch_norm_is_4d(node: torch.fx.Node) -> bool: + """Only support batch norm on a 4d input. + + add_native_batch_norm_node() asserts + VK_CHECK_COND(in_sizes.size() == 4, "BatchNorm only support 4d tensor"), so + partitioning a batch norm whose input is not 4d (any conv1d model, where + activations are rank 3) yields a .pte that aborts at execute time. + """ + input_node = node.args[0] + if not isinstance(input_node, torch.fx.Node): + return False + val = input_node.meta.get("val") + return val is not None and val.dim() == 4 + + @update_features(exir_ops.edge.aten._native_batch_norm_legit_no_training.default) def register_native_batch_norm_legit_no_training(): return OpFeatures( @@ -1742,6 +1795,7 @@ def register_native_batch_norm_legit_no_training(): inputs_dtypes=utils.FP_T, supports_prepacking=True, supports_resize=True, + are_node_inputs_supported_fn=_check_batch_norm_is_4d, )