From ab4239a78f8123087986906f8aec5d176440db47 Mon Sep 17 00:00:00 2001
From: Amanda Tang <amandatang88888@gmail.com>
Date: Fri, 10 Nov 2023 17:29:42 -0500
Subject: [PATCH] Update linalg fill syntax

---
 mlir/dialects/linalg.py | 21 +++++++++++++--------
 tests/test_linalg.py    |  4 ++--
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/mlir/dialects/linalg.py b/mlir/dialects/linalg.py
index e348694..8b5c938 100644
--- a/mlir/dialects/linalg.py
+++ b/mlir/dialects/linalg.py
@@ -114,16 +114,21 @@ class LinalgDot(DialectOp):
 
 @dataclass
 class LinalgFill(DialectOp):
-    output_id: mast.SsaId
-    value_id: mast.SsaId
-    output_type: mast.Type
-    value_type: mast.Type
+    in_id: mast.SsaId
+    in_type: mast.Type
+    out_id: mast.SsaId
+    out_type: mast.Type
+    res_type: mast.Type
     attr: Optional[mast.Attribute] = None
 
-    _syntax_ = [("linalg.fill( {output_id.ssa_id} , {value_id.ssa_id} ) "
-                "{attr.attribute_value} : {output_type.type} , {value_type.type}"),
-                ("linalg.fill( {output_id.ssa_id} , {value_id.ssa_id} ) "
-                " : {output_type.type} , {value_type.type}")]
+    _syntax_ = [("linalg.fill"
+                 " ins( {in_id.ssa_id} : {in_type.type} )"
+                 " outs( {out_id.ssa_id} : {out_type.type} )"
+                 "  {attr.attribute_value} -> {res_type.type}"),
+                ("linalg.fill"
+                 " ins( {in_id.ssa_id} : {in_type.type} )"
+                 " outs( {out_id.ssa_id} : {out_type.type} )"
+                 " -> {res_type.type}")]
 
 
 @dataclass
diff --git a/tests/test_linalg.py b/tests/test_linalg.py
index 1b4f61b..ef5694f 100644
--- a/tests/test_linalg.py
+++ b/tests/test_linalg.py
@@ -68,8 +68,8 @@ def test_dot():
 
 def test_fill():
     assert_roundtrip_equivalence("""module {
-  func.func @fill_view(%arg0: memref<?xf32, strided<[1], offset: ?>>, %arg1: f32) {
-    linalg.fill( %arg0 , %arg1 )  : memref<?xf32, strided<[1], offset: ?>> , f32
+  func.func @fill_view(%arg0: f32, %arg1: tensor<?x?xf32>) {
+    linalg.fill ins( %arg0 : f32 ) outs( %arg1 : tensor<?x?xf32> ) -> tensor<?x?xf32>
     return
   }
 }""")