search

iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.
http://iree.dev/
Apache License 2.0
2.47k stars 548 forks source link

[LLVMGPU] Generalize AMDGPUChainedMatmul pass to multiple dimensions #17684

Closed Groverkss closed 1 week ago

Groverkss commented 1 week ago

This patch generalizes the AMDGPUChainedMatmul pass to use VectorContractOpInfo to query and transpose dims, instead of hardcoding indexing maps.

Groverkss commented 1 week ago

Since this introduces a n, m matmul, I need to check if attention still works with this patch.

Groverkss commented 1 week ago

Actually, CI seems to be passing on this, so I'm going to land this