[converter] Harden mixed-source SymInt lowerings under dynamic shapes

coreai_torch/_aten_to_core.py

@@ -644,17 +644,20 @@ def replace_arange_start_step(
         else coreai.constant(1, dtype=start.type.element_type)
     )
 
-    # Squeeze rank-1 inputs to 0D scalars for coreai.range_.
+    # coreai.range_ requires scalar (rank-0) operands that share an element
+    # type. aten.arange promotes mixed-type scalars internally; we replicate
+    # that here by squeezing each operand to rank-0 and casting to the FX
+    # node's output dtype before the op.
+    target_type = get_output_element_type_from_node(node)
+
     def to_scalar(v: Value) -> Value:
         if v.type.rank > 0:
-            return coreai.shrink_dims(v, list(range(v.type.rank)))
+            v = coreai.shrink_dims(v, list(range(v.type.rank)))
+        if v.type.element_type != target_type:
+            v = coreai.cast(v, target_type)
         return v
 
-    result = coreai.range_(to_scalar(start), to_scalar(end), to_scalar(step))
-    target_type = get_output_element_type_from_node(node)
-    if result.type.element_type != target_type:
-        result = coreai.cast(result, target_type)
-    return result
+    return coreai.range_(to_scalar(start), to_scalar(end), to_scalar(step))
 
 
 def replace_batch_norm(
@@ -814,6 +817,7 @@ def replace_binary_ops(
         "pow.Scalar": coreai.broadcasting_pow,
         "pow.Tensor_Tensor": coreai.broadcasting_pow,
         "pow.Tensor_Scalar": coreai.broadcasting_pow,
+        "pow": coreai.broadcasting_pow,
         "sub.Tensor": coreai.broadcasting_sub,
         "sub.Scalar": coreai.broadcasting_sub,
         "sub": coreai.broadcasting_sub,
@@ -944,6 +948,36 @@ def replace_cat(values_map: dict[str, Value], node: fx.Node, loc: Location) -> V
 
     rank = inputs[0].type.rank
     dim = dim + rank if dim < 0 else dim
+
+    # coreai.concat requires all non-concat dims to be provably equal across
+    # inputs. Under dynamic shapes, one branch can carry a dynamic non-concat
+    # axis while a sibling has a static size for the same axis — the dynamic
+    # side must in fact equal that static size, but the type system doesn't
+    # know it. Reshape such inputs to the known static size before the concat.
+    # Multiple distinct static sizes on one axis is a real mismatch and is
+    # left for the dialect verifier to reject.
+    dyn = ShapedType.get_dynamic_size()
+
+    def known_static(axis: int) -> int | None:
+        if axis == dim:
+            return None
+        sizes = {inp.type.shape[axis] for inp in inputs if inp.type.shape[axis] != dyn}
+        return next(iter(sizes)) if len(sizes) == 1 else None
+
+    statics = [known_static(a) for a in range(rank)]
+    promoted: list[Value] = []
+    for inp in inputs:
+        new_shape = [
+            statics[a]
+            if statics[a] is not None and inp.type.shape[a] == dyn
+            else inp.type.shape[a]
+            for a in range(rank)
+        ]
+        if new_shape != list(inp.type.shape):
+            inp = coreai.reshape(inp, new_shape)
+        promoted.append(inp)
+    inputs = promoted
+
     return coreai.concat(dim, inputs)
 
 
@@ -2595,6 +2629,7 @@ def replace_unary_ops(
         "log.default": coreai.log,
         "relu.default": coreai.relu,
         "round.default": coreai.round_,
+        "round": coreai.round_,
         "rsqrt.default": coreai.rsqrt,
         "sigmoid.default": coreai.sigmoid,
         "silu.default": coreai.silu,
@@ -3434,13 +3469,15 @@ def sdpa_maskless(q: Value, k: Value, v: Value) -> Value:
     "pow.Scalar": replace_binary_ops,
     "pow.Tensor_Scalar": replace_binary_ops,
     "pow.Tensor_Tensor": replace_binary_ops,
+    "pow": replace_binary_ops,
     "prod.default": replace_prod_default,
     "prod.dim_int": replace_prod_dim_int,
     "reciprocal.default": replace_reciprocal,
     "relu.default": replace_unary_ops,
     "remainder.Tensor": replace_remainder,
     "round.default": replace_unary_ops,
     "round.decimals": replace_round_decimals,
+    "round": replace_unary_ops,
     "repeat.default": replace_repeat,
     "rsqrt.default": replace_unary_ops,
     "scaled_dot_product_attention.default": replace_sdpa,

coreai_torch/_utils.py

@@ -175,6 +175,23 @@ def __exit__(self, *exc: Any) -> None:
         self.close()
 
 
+def to_rank1_int32(v: Value) -> Value:
+    """Coerce a SymInt-derived Value to canonical rank-1 si32 form.
+
+    Dim-vector concats (used to build shape operands for ``coreai.reshape``,
+    ``coreai.interpolate``, etc.) require all inputs to share rank and
+    element type. SymInt values can arrive rank-0 (e.g. from
+    ``aten._local_scalar_dense``) or with a different int variant. This
+    helper produces the form that aligns with ``coreai.constant([i],
+    dtype=np.int32)`` and ``replace_sym_size_int``.
+    """
+    if v.type.rank == 0:
+        v = coreai.reshape(v, [1])
+    if v.type.element_type != IntegerType.get_signed(32):
+        v = coreai.cast(v, np.int32)
+    return v
+
+
 def upsample_build_output_shape_dynamic(
     x: Value, out_h: int | Value, out_w: int | Value
 ) -> Value:
@@ -192,8 +209,8 @@ def upsample_build_output_shape_dynamic(
     )
     shape = coreai.cast(coreai.get_shape(x), dtype=np.int32)
     non_spatial = coreai.slice_(shape, [0], [2], [1])
-    h = [out_h] if isinstance(out_h, int) else out_h
-    w = [out_w] if isinstance(out_w, int) else out_w
+    h = [out_h] if isinstance(out_h, int) else to_rank1_int32(out_h)
+    w = [out_w] if isinstance(out_w, int) else to_rank1_int32(out_w)
     return coreai.concat(0, [non_spatial, h, w])
 
 
@@ -986,9 +1003,13 @@ def get_operand(
     if isinstance(arg, fx.Node):
         return values_map[arg.name]
     if isinstance(arg, list) and any(isinstance(e, fx.Node) for e in arg):
-        # Mixed list: resolve fx.Node elements via values_map, keep ints as constants.
+        # Mixed list: SymInt fx.Nodes + plain ints. Concat the two sources
+        # into a single rank-1 si32 dim vector. Both branches must produce
+        # the same canonical form so the concat verifier accepts them.
         dim_vals = [
-            values_map[e.name] if isinstance(e, fx.Node) else coreai.constant([e])
+            to_rank1_int32(values_map[e.name])
+            if isinstance(e, fx.Node)
+            else coreai.constant([e], dtype=np.int32)
             for e in arg
         ]
         return coreai.concat(0, dim_vals) if len(dim_vals) > 1 else dim_vals[0]