iree-org / iree-turbine

IREE's PyTorch Frontend, based on Torch Dynamo.
Apache License 2.0
55 stars 25 forks source link

Add support for varying vector shapes #247

Closed harsh-nod closed 2 weeks ago

harsh-nod commented 2 weeks ago

This PR adds support for expanding operators with varying vector shapes, specifically for the MMA case where either the same dimension has different vector shapes in different mmas or if different instructions are being used.

The idea is to insert a reshape operator whenever such a shape mismatch is discovered. The reshape operator lowers to an extract or concatenate operation, depending on the context.