From 1fc710a1903da3e54c2db558943b02bf1ee7e6cc Mon Sep 17 00:00:00 2001
From: saienduri <77521230+saienduri@users.noreply.github.com>
Date: Thu, 24 Oct 2024 12:42:15 -0700
Subject: [PATCH] Add region for attention. (#27)

Related IREE commit:
https://github.com/iree-org/iree/commit/a488d3826eaf263a161eec27ae428e60ea717a5a
---
 attentionbench/attention_utils.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/attentionbench/attention_utils.py b/attentionbench/attention_utils.py
index 6397f01..d8d8d18 100644
--- 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
 }}
 """