tree c2f353ef7abdefef6a960f1bd501061e9826a957
parent e83ce9e225eefd9245cd38d37d843a915fbaa034
author harsh-nod <harsh@nod-labs.com> 1707945664 -0800
committer GitHub <noreply@github.com> 1707945664 -0800
gpgsig -----BEGIN PGP SIGNATURE-----
 
 wsFcBAABCAAQBQJlzS7ACRC1aQ7uu5UhlAAAqPUQAFAfvr8ppSaAnbwk84CJEjm4
 zf0ZsjM+T6WoBr7kfc6QfWFTpeGfEI37LW2i4A7VUJCihCzkupMzSIGELbUZedah
 GXIwZBHUFKHg/1Tfqlpz10X8RZzzgbHfqnsAsxKscLNVGSssU35EXh1ugHPvFJlG
 izwvfo8iEwFbkffLQG44A9M2O0DlTHEbyUxZgOeLjwXmomciYC1EUTwF3J5s2E32
 VivPkX/aNXOuR+AvFAF8h/Il1zyeyE6Aepza/jgdMJF0iqh5Jj9nFUDxGupkeCTb
 I868T88pAsnSn+N/2SMemLYpkZzEliKYToYqKouRM6o0DHwxmJwyr8Ns1tWkJudo
 kwhrem5GEfLYlO1LShhHQc0CKXWbLMwou0rIss1KbTL/Ugf/UaavOPSruG0MiE1U
 RPxX32SWNh3qfTBc8lLBBnzzG62fpWwviKix6nZWhegBEO8xsDdTjhspqBu5Mf3+
 JZqgMDUgg+qWlXC1sgyua8kMyxM6XxadH1Aexdnoj2Y4p9KSJmuhxz1OGKWXtbmD
 4OlQu5k5Zq8OfGXOVzQ6JefsbYW+10K89EsRKeHyK0Ge5vn3H3A5bXE4jihTwmXf
 4FDcC5ys3o/3KB3L5bXqJ+P2zXsM+4vKIYdpa8P1Q/D0vnvUE0szaCuuqgdzwaQ4
 lO4I7HMLqVHD3wXKpF1z
 =WkRf
 -----END PGP SIGNATURE-----
 

Add AMDGPU pattern for chained matmuls (#16398)

This PR adds a pattern for chained matmuls operators like Flash
Attention where by swapping the operands of the sequential matmuls, one
can keep the data in registers and avoid a trip to shared memory and/or
additional shuffle instructions.

The pattern swaps the operands of the contract op, and inserts
transposes for the accumulator and result, if the contracts are chained
and satisfy the MMT
indexing maps.