[LinalgExt] Implement PartialReductionOpInterface for OnlineAttention…

…Op (#19684)
iree-org · Jan 14, 2025 · 6fd0fd0 · 6fd0fd0
1 parent 3c963dd
commit 6fd0fd0
Show file tree

Hide file tree

Showing 3 changed files with 462 additions and 15 deletions.
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td
@@ -649,7 +649,12 @@ def IREELinalgExt_OnlineAttentionOp : IREELinalgExt_PureOp<"online_attention",
       ["getIterationDomain",
        "getLoopIteratorTypes",
        "getResultTilePosition",
-       "getTiledImplementation"]>]> {
+       "getTiledImplementation"]>,
+     DeclareOpInterfaceMethods<PartialReductionOpInterface,
+      ["generateInitialTensorForPartialReduction",
+       "tileToPartialReduction",
+       "mergeReductions",
+       "getPartialResultTilePosition"]>]> {
   let summary = "Online Attention operator";
   let description = [{
     Traditional scaled dot product attention computes: