tree 54232463ae3917ee389a23cb09b0d1b0573a093d
parent 891f438c27f201ae0dde3e8408eefc6fca80bc5f
author rohan-tan-bhowmik <46410002+rohan-tan-bhowmik@users.noreply.github.com> 1726985238 -0700
committer GitHub <noreply@github.com> 1726985238 -0700
gpgsig -----BEGIN PGP SIGNATURE-----
 
 wsFcBAABCAAQBQJm77QWCRC1aQ7uu5UhlAAAwwIQAEskM3/Fw8uQgRymu3eOknwc
 8jt+bHywziNJspxlMyx8Vk/2ykDzaT/oREBWGah2x6DaSaMgI8xL2W2pJmF+TP6g
 DnqX+HMckDbYmtNCnFF9lOgPiipT3/I/yp6Ge4qlLgQiZ7G7iXoJYhgbuNDE/TKr
 ClC/iWZWm6DNwJnWABlqFis5wro9zrM7RYMe1ftSNddH0/WCs/O1f8inh3MZr2OD
 tiI4qxu82gag3unjVi14uzj7YgdPgDym/HPcqhKzOZrZn7vFphT8Aal/pju850Du
 ZhUatZDzbPvIgdm4h5dPQSM0RFiA3S8k0Jxtzl55+OJ+DwO6DZQFflgxjJntZhmq
 iG0eKMEbS8XaR4BP1gE7Pi+nFcrRtyZR8SnZqsis+lDy63eKlF6aQwiAFHdnUH7V
 C/7vSsz/JdY0jK0iWx7sf823osO63zMlN9dhPN23Sr/2njMgXYFImi+xbTbI8JQD
 FbJfbNN88RUU9/L0J3zKvH29IAKIavqCenfPcnaP35YmJi7dwA1Ty5pIft/uYFC+
 qu04Mmaw2TWDRWZdtn23azXw3F241FGs6T6rVB4HGZjvs0yG86bw+30OwKRYYUkq
 RR5W/703l7ALkq78h1RWLG55a+g4cNW8YCt0W+0g9wQ19yDTQd7/3ZA0rshwmsMu
 /Rh5R7dnQNmDy/3Ghdrt
 =dpwX
 -----END PGP SIGNATURE-----
 

[LinalgExt] Masked Attention Implementation (#18525)

Enables float/boolean mask as parameters and created linalg generic ops
to apply masking. This image (https://imgur.com/a/1MePgcy) elaborates on
the main files changed and how they enable masked attention:
- Blue boxes represent changed .cpp and .td files to
enable/pass/decompose the mask
  -  Yellow boxes represent the different op classes
- Red boxes represent test mlir files pertaining to certain .cpp/.td
implementations or ops
  
For quick reference, AggregateOpInterfaceImpl.cpp contains the bulk of
the actual mask decomposition (QK += mask)

And for clarification, TileAttention.cpp only holds the
convertToOnlineAttentionOp and getTileAttentionIndexingMaps functions;
TilingInterfaceImpl.cpp contains the main tiling capabilities in the
form of AttentionOp::getTiledImplementation and
OnlineAttentionOp::getTiledImplementation.
  

Updated version of https://github.com/iree-org/iree/pull/18461. This
version was created to include scale affine map and enable fused
attention (incorporated
https://github.com/IanWood1/iree/tree/raikonen/sdpa_mask).
- To that end, many modifications in tests are for adding the scale
affine map (without much functionality change)
- For tiling and decomposition tests, most functionality tests are
included in "tiling.mlir" and "decompose_online_attention.mlir". On the
other hand, the "tile_attention.mlir and "decompose_attention.mlir" are
old paths intended to be be retired and deprecate soon. Hence, no major
tests were added it there.

Test directory for numerical verification:
https://github.com/rohan-tan-bhowmik/iree-masked-attention-test

---------

Signed-off-by: Stanley Winata <stanley.winata@amd.com>
Co-authored-by: Stanley Winata <stanley.winata@amd.com>
Co-authored-by: Ian Wood <ianwood2024@u.northwestern.edu>