tree a5df137e2c4a0427107e8f76adcabb8dc28f26d8
parent 68a9309688a6605151d97d369a68cb5b620802bc
author Zhewen Yu <zhewenyu@amd.com> 1764847392 +0000
committer GitHub <noreply@github.com> 1764847392 +0000
gpgsig -----BEGIN PGP SIGNATURE-----
 
 wsFcBAABCAAQBQJpMW8gCRC1aQ7uu5UhlAAAcPoQAG9TXI53SiEKPSkDWLqfwvr8
 QWRTaII6j1OSzOiNb7j3qDNyH1/fNggohvXVrWPybJ2xjhZRJX3HaWZaHXQLFSMm
 asUO1tBHcTxCFQ44xSnD+lc5xrqQCpZQh8T32tLjlvyr1Qhwgx5d6Z6GNqLybqyB
 zrVpS9eymn2U7ax9wFz0Zo+6cIo2W3/aSv/VMwQpXqRoDcG1l5drAN4faElcvfnP
 Amx8M/M9IX2t+ZussiB1TxmJ3B02+1dJ3kzbobOCfm+QVT0l6LeyINnygjjtrzeM
 36HOIFZGTUXCUQaYY9QFHHCDa2bToeynqL18htrOlmIk7aBzEHxzKJv6aFtEW31a
 hhw6evlImsHGuAFXI+w9YC8F8FWxtjSwZkPZzlp90zPf0EevhxbWSqUItuMmx9gX
 gwwyJivexHmKQo9cscVEEIgH07vUZe8dVhQaSyaTrfMyqoS6uw/fWmWGhojxExNq
 re+/S7qC0CYtvtuA5NfqkK30raAP5HQIBIcYlhyYrvJqyvnt/5rmjZAIbFanBh/7
 WzAzGNXoDifKf7M4vlbmZ/0G4SeRsCzxqMC+KQ52XorRs85XwCJGNum64rgexYt4
 HFFn+SbSfjBNoTLhDrlIGZoLomTGRfzvtYnyEF7NKJaRwQezyRvb67Q6jsCVDaoS
 bao8rkKYX5tQjHazK12+
 =m7ww
 -----END PGP SIGNATURE-----
 

[GPU] Fix alignment check for scaled matmul (#22737)

## Problem

The current alignment check in `GPUHeuristics.cpp` is incorrect for any
intrinsic that has multiple M, N, and K dimensions. The root cause is
that the product of intrinsic sizes is passed to `GPUMMASchedule`
instead of passing the individual dimension sizes as a vector.

## Example


https://github.com/iree-org/iree/blob/b98c1b92cb630bd696992f47df591bb2f247a8d7/compiler/src/iree/compiler/Codegen/Common/GPU/GPUHeuristics.cpp#L516-L525

Consider the scaled MFMA where `intrinsic.kSizes = [K, KB] = [4, 32]`.
Instead of passing the vector `[4, 32]`, the value `128` (product: 4 ×
32) is passed to `GPUMMASchedule`.


https://github.com/iree-org/iree/blob/b98c1b92cb630bd696992f47df591bb2f247a8d7/compiler/src/iree/compiler/Codegen/Common/GPU/GPUHeuristics.cpp#L98-L108

Assume tile size = `[4, 1]`. The returned schedule sizes become `[4,
128]` instead of the correct `[16, 32]`. As a result, the last dimension
`128` always makes the alignment check fail, since the problem size of
KB is `32` and `32 % 128 != 0`.

When the alignment check fails, no intrinsic is selected and the
operation falls back to complete serialization. This leads to extremely
slow execution for workloads like Llama 405B FP4 prefill with direct
codegen.

## Solution

This PR passes all intrinsic sizes as vectors to `GPUMMASchedule`.

## Performance

**Llama 405B FP4 prefill direct codegen with shark-ai:**
- Before: 11 minutes
- After: 234 ms

Closes: #22559

ci-extra: test_torch

---------

Signed-off-by: Yu-Zhewen <zhewenyu@amd.com>