Skip to content

[Torch] add aten._scaled_mm op support and FX import plumbing#4558

Open
catcor01 wants to merge 6 commits intollvm:mainfrom
catcor01:cc/scaled-mm-op-import
Open

[Torch] add aten._scaled_mm op support and FX import plumbing#4558
catcor01 wants to merge 6 commits intollvm:mainfrom
catcor01:cc/scaled-mm-op-import

Conversation

@catcor01
Copy link
Copy Markdown
Contributor

@catcor01 catcor01 commented May 6, 2026

Adds Torch dialect support and FX importer plumbing for aten._scaled_mm.

This is the first PR in a three-part stack:

  1. Torch op + FX import plumbing
  2. Static FP8 TorchToTosa lowering
  3. MXFP8/block-scaled TorchToTosa lowering

Add Torch op definitions, abstract interpretation support, and FX
importer handling for aten._scaled_mm, including float8_e8m0fnu
blocked-scale tensors.

Keep this change independent of TOSA legalization.

Include frontend export/import coverage for per-tensor and
blocked-scale FP8 shapes, plus out_dtype=None dtype inference.

Signed-off-by: Cathal Corbett <cathal.corbett@arm.com>
Change-Id: I102732ad725f89477b7f8fb2339d4fe920fa647b
@catcor01 catcor01 force-pushed the cc/scaled-mm-op-import branch from 6d12002 to c4b7aef Compare May 6, 2026 08:16
@catcor01
Copy link
Copy Markdown
Contributor Author

catcor01 commented May 6, 2026

Comment thread test/python/fx_importer/basic_test.py
Change-Id: I555367f4c8a3725fba75a1389e170bf402ad5eb9
Copy link
Copy Markdown
Member

@Lallapallooza Lallapallooza left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the patch, few comments.

Comment thread test/python/fx_importer/basic_test.py
catcor01 added 4 commits May 6, 2026 15:47
Change-Id: I3f62b8d60de4155295ec6f552f68c3d5f21f2c08
Change-Id: I1c1a386a6395b6416c6ada5d06cc82f87271bd5d
Change-Id: I7d4343af324fafe830ff514d79d8c6f03ac200e3
Change-Id: I52f0543dc1120e3a8e3d358486c2cafc8a63d1eb
@catcor01 catcor01 requested a review from Lallapallooza May 7, 2026 11:04
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