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
- Retire the deprecated old-tog negative-coefficient affine formatting (confirm
nothing on the current trace path depends on the -1 * i0 form).
- 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.
- 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).
Context
After #271 removed the dead
convert_indexfloor/mod codegen path, the affineload/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/ModularIndexingupstream, and any residual floor/mod now raises loudly.
With that narrower input, the function is more elaborate than it needs to be:
mlir_str(expr)builder (Add -> " + ",Mul -> " * ",Symbol/number -> str, else fallbackstr), andindicesloop that re-extracts the per-term symbol fromsorted_args,which for pure-affine input is largely equivalent to
str(expr)plussorted(expr.free_symbols, key=str). (The redundant assumption-strippingSymbol(str(...))+expr.replaceround-trip was already dropped in #271.)Why it isn't a one-line cleanup yet
The one behavioural difference is negative-coefficient formatting:
mlir_stremits e.g.
i1 + -1 * i0where sympy'sstrgivesi1 - i0. Thatnegative-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_strintostr(expr)wouldchange the emitted affine string for negative coefficients.
Proposed work
nothing on the current trace path depends on the
-1 * i0form)._convert_sympy_to_mlir_expr: collapsemlir_strtowardstr(expr)and theindicesloop towardsorted(expr.free_symbols); thesorted_argsparameter (and the matching term-sort inparse_indices) maythen become unnecessary. Apply the same to the parallel logic in
parse_index_list.tests/ops/view/*incl.test_floormod_axis_split, transpose2D/3D, view3D_2D, cat) plus the affineindex 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).