Skip to content

[Frontend] Simplify _convert_sympy_to_mlir_expr and retire the deprecated old-tog negative-coefficient affine formatting #273

Description

@YWHyuk

Context

After #271 removed the dead convert_index floor/mod codegen path, the affine
load/store index reaching _convert_sympy_to_mlir_expr
(PyTorchSimFrontend/mlir/mlir_codegen_backend.py) is always a pure affine,
constant-stride
expression -- axis-split's affine-only contract
(docs/axis-split-scheduling.md) linearizes aligned FloorDiv/ModularIndexing
upstream, and any residual floor/mod now raises loudly.

With that narrower input, the function is more elaborate than it needs to be:

  • a custom recursive mlir_str(expr) builder (Add -> " + ", Mul -> " * ",
    Symbol/number -> str, else fallback str), and
  • an indices loop that re-extracts the per-term symbol from sorted_args,

which for pure-affine input is largely equivalent to str(expr) plus
sorted(expr.free_symbols, key=str). (The redundant assumption-stripping
Symbol(str(...)) + expr.replace round-trip was already dropped in #271.)

Why it isn't a one-line cleanup yet

The one behavioural difference is negative-coefficient formatting: mlir_str
emits e.g. i1 + -1 * i0 where sympy's str gives i1 - i0. That
negative-coefficient handling is deprecated old-tog behaviour -- it exists to
match what the old TOG path emitted, and should be retired gradually rather than
silently changed. Until it is gone, collapsing mlir_str into str(expr) would
change the emitted affine string for negative coefficients.

Proposed work

  1. Retire the deprecated old-tog negative-coefficient affine formatting (confirm
    nothing on the current trace path depends on the -1 * i0 form).
  2. Then simplify _convert_sympy_to_mlir_expr: collapse mlir_str toward
    str(expr) and the indices loop toward sorted(expr.free_symbols); the
    sorted_args parameter (and the matching term-sort in parse_indices) may
    then become unnecessary. Apply the same to the parallel logic in
    parse_index_list.
  3. Verify the view/affine suite is unchanged (tests/ops/view/* incl.
    test_floormod_axis_split, transpose2D/3D, view3D_2D, cat) plus the affine
    index paths in gemm/conv/reduce.

Notes

Low priority / cleanup. Tracking so the simplification opportunity surfaced
during #271 is not forgotten. The floor/mod guards must stay (they enforce the
affine-only contract).

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions