Add region for attention

nod-ai · Oct 22, 2024 · c74ba4f · c74ba4f
1 parent 0aa226d
commit c74ba4f
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/attentionbench/attention_utils.py b/attentionbench/attention_utils.py
@@ -139,8 +139,10 @@ def generate_mlir(config: AttentionConfig, tuning: Optional[TuningSpec] = None):
        {{ indexing_maps = [#Q, #K, #V, #S, #O]
          {",compilation_info = #tuning" if tuning and config.dtype == "f16" else ""}
        }}
-       ins(%Q, %K, %V, %scale : !Q, !K, !V, !dtype)
-       outs(%empty : !O) -> !O
+       ins(%Q, %K, %V, %scale : !Q, !K, !V, !dtype) outs(%empty : !O) {{
+          ^bb0(%score: f32):
+            iree_linalg_ext.yield %score : f32
+        }} -> !O
   return %O : !O
 }}
 """