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
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions