Vulkan: do not partition ops whose symbolic arguments the runtime reads at build time

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<float>() 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,13 +1773,29 @@ 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(
         inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
         inputs_dtypes=utils.FP_T,
         supports_prepacking=True,
         supports_resize=True,
+        are_node_inputs_supported_fn=_check_batch_norm_is_4d,
     )