Skip to content

[converter] Support SymInt repeats in aten.repeat lowering#15

Merged
gokulkrishna98 merged 4 commits into
apple:mainfrom
gokulkrishna98:dev/gokul/sam-repeat-symint
Jun 16, 2026
Merged

[converter] Support SymInt repeats in aten.repeat lowering#15
gokulkrishna98 merged 4 commits into
apple:mainfrom
gokulkrishna98:dev/gokul/sam-repeat-symint

Conversation

@gokulkrishna98

@gokulkrishna98 gokulkrishna98 commented Jun 15, 2026

Copy link
Copy Markdown
Contributor

Description:

  • replace_repeat previously called np.array(repeats, dtype=uint32), which raises TypeError when any entry is an fx.Node derived from a symbolic shape.
  • When any repeat entry is dynamic, build the tile-dims vector at runtime: emit constant chunks for plain ints and the resolved value (cast to uint32, lifted to rank-1) for SymInts, concat along axis 0, and pass the resulting value to tile. All-static repeats keep the existing fast path.

Stacking:

  1. [converter] Harden mixed-source SymInt lowerings under dynamic shapes #13
  2. [converter] Fix cat lowering when promoted shape still has dynamic axes #14
  3. [converter] Support SymInt repeats in aten.repeat lowering #15. <----

Testing

  • python unit test
  • enables conversion of dynamic config of Sam model

…amic shapes

Six related fixes that surface together when exporting models whose FX
graphs combine SymInt-derived shape arithmetic with mixed source
element types and ranks:

1. _aten_to_core_resolver / replace_binary_ops _op_map: register a
   bare 'pow' entry alongside the variant-suffixed ones. Some
   torch.export rewrites leave ``aten.pow`` as the OpOverloadPacket
   target with no overload suffix; without this entry the converter
   raises ``Unsupported ATen op: pow``.
2. Same registries for bare 'round': torch.export can leave
   ``aten.round`` without a ``.default`` overload, mirroring the pow
   case.
3. upsample_build_output_shape_dynamic: ensure each (out_h, out_w)
   operand is rank-1 with int32 element type before the concat that
   builds the output shape — the dialect verifier rejects mixed-rank
   / mixed-element-type concat inputs. Hits when out_h/out_w are
   SymInts derived from ``round(SymFloat)`` arithmetic.
4. get_operand mixed-list path (SymInt + plain int): normalise each
   resolved Value to the same canonical rank-1 si32 form and emit
   plain-int constants with explicit ``dtype=np.int32`` so the
   dim-vector concat sees uniform operands. Hits ops like ``view``,
   ``expand``, ``reshape``, ``repeat`` whenever a dim list mixes
   SymInts with ints.
5. replace_cat: when one input has a dynamic non-concat axis and a
   sibling has a known static size for that axis, reshape the
   dynamic side to that static size before the concat. Localised
   shape inference using the fact that all non-concat dims must be
   equal — multiple distinct static sizes is left for the dialect
   verifier to reject.
6. replace_arange_start_step: unify start/end/step element types to
   the FX node's output dtype before ``coreai.range_``. Mirrors
   aten.arange's internal type promotion since coreai.range_'s
   verifier requires uniform element types.

Adds a shared ``to_rank1_int32(v)`` helper in ``_utils.py`` so fixes
3 and 4 share one canonical normalization (rank-0 → rank-1, cast to
signed int32 if needed); both call sites collapse to one line per
operand.

One regression test per non-trivial fix, each verified to FAIL
without the fix and PASS with it (verified by reverting each fix
individually):

- TestRound: bare ``aten.round`` overload-packet target must lower.
- TestUpsampleNearest2d / TestUpsampleBilinear2d::test_round_symfloat_size:
  ``round((num / aspect) ** 0.5) * 14`` output_size produces SymInts
  whose Value type doesn't match the int32 constants used elsewhere.
  Pre-fix: ``coreai.concat`` raises ``Operation creation failed``.
- TestView::test_view_with_round_symfloat_dims: same trigger applied
  to a ``view([1, C, h, w])`` mixed list. Pre-fix: ``expected the
  same element type for all inputs to concat``.
- TestCat::test_dynamic_vs_static_non_concat_axis: ``Dim.AUTO`` on
  one side + static sibling forces non-concat-axis promotion.
- TestArange::test_symint_end_with_float_start_step: float ``arange``
  with SymInt-derived end exercises the element-type unify.
…ill-dynamic siblings

The previous replace_cat fix uses ``coreai.reshape(inp, new_shape)`` with
a Python list to promote a dynamic non-concat axis to its known static
size before the concat. The list form materializes the shape as an
``int32`` constant tensor, so every entry must be a real int — there is
no slot for the dynamic-size sentinel.

When a cat input has both a promotable axis (sibling has a known static
size) and an axis that is dynamic on every input (no static sibling),
post-promotion ``new_shape`` is a mix of concrete ints and the dynamic
sentinel; passing it to the list-form reshape raises
``OverflowError: Python integer -9223372036854775808 out of bounds for
int32``.

Split ``replace_cat`` into two reshape paths:

- All axes static post-promotion: keep the list-form reshape.
- Some axes still dynamic post-promotion: build the shape vector at
  runtime — ``coreai.get_shape(inp)`` for the still-dynamic axes,
  ``coreai.constant`` slices for the promoted axes, concat along axis 0
  to get a rank-1 ``int32`` Value, and pass that to Value-form
  ``coreai.ReshapeOp`` with a partially-static result type.

Adds a numerical regression that exercises the mixed path: a 4-D cat
where one input has a sibling-promotable axis and another axis that
stays dynamic on every input. Verified to fail with the list-form-only
fix and pass with this delta.
… dims

The previous ``replace_repeat`` lowering called
``np.array(node.args[1], dtype=np.uint32)`` on the repeats list. Under
dynamic shapes, ``tensor.repeat(...)`` can be called with at least one
entry derived from a symbolic shape; that entry appears in the FX graph
as a ``torch.fx.Node``, and ``np.array`` raises
``TypeError: int() argument must be a string ... not 'Node'`` trying to
coerce it.

Detect any ``fx.Node`` in the repeats list and take a dynamic path:
build a rank-1 ``uint32`` dim vector at runtime by emitting per-axis
``coreai.constant`` chunks for plain ints and the resolved Value (cast
to ``uint32``, lifted to rank-1 if scalar) for SymInts, concat along
axis 0, and pass the resulting Value to ``coreai.tile`` (which accepts
a runtime Value for its dims). The all-static fast path keeps the
existing ``np.array(...)`` form.

Adds an IR FileCheck test pinning the dynamic-tile shape and a
numerical-output test. Both fail without the fix (the IR test
reproduces the ``TypeError`` above).
@gokulkrishna98 gokulkrishna98 merged commit ced5268 into apple:main Jun 16, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants