[Mlir-commits] [mlir] c606fef - [MLIR][python bindings] Add more basic AttrBuilder for _ops_gen.py files

[Alex Zinenko]

Author: pengchao.hu
Date: 2023-05-22T18:38:25+02:00
New Revision: c606fefa8579f1d00b53d0a7a6b63319526003e4

URL: https://github.com/llvm/llvm-project/commit/c606fefa8579f1d00b53d0a7a6b63319526003e4
DIFF: https://github.com/llvm/llvm-project/commit/c606fefa8579f1d00b53d0a7a6b63319526003e4.diff

LOG: [MLIR][python bindings] Add more basic AttrBuilder for _ops_gen.py files

Add more attribute builders, such as "F32Attr", "F64Attr" and "F64ArrayAttr", which are useful to create operations by python bindings. For example, tosa.clamp in _tosa_ops_gen.py need 'F32Attr'.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D150757

Added: 
    

Modified: 
    mlir/python/mlir/ir.py
    mlir/test/python/dialects/python_test.py
    mlir/test/python/python_test_ops.td

Removed: 
    


################################################################################
diff  --git a/mlir/python/mlir/ir.py b/mlir/python/mlir/ir.py
index 714253426b025..be065d4639519 100644
--- a/mlir/python/mlir/ir.py
+++ b/mlir/python/mlir/ir.py
@@ -27,7 +27,7 @@ def _indexAttr(x, context):
 
 
 @register_attribute_builder("I16Attr")
-def _i32Attr(x, context):
+def _i16Attr(x, context):
   return IntegerAttr.get(IntegerType.get_signless(16, context=context), x)
 
 
@@ -41,6 +41,26 @@ def _i64Attr(x, context):
   return IntegerAttr.get(IntegerType.get_signless(64, context=context), x)
 
 
+ at register_attribute_builder("SI16Attr")
+def _si16Attr(x, context):
+  return IntegerAttr.get(IntegerType.get_signed(16, context=context), x)
+
+
+ at register_attribute_builder("SI32Attr")
+def _si32Attr(x, context):
+  return IntegerAttr.get(IntegerType.get_signed(32, context=context), x)
+
+
+ at register_attribute_builder("F32Attr")
+def _f32Attr(x, context):
+  return FloatAttr.get_f32(x, context=context)
+
+
+ at register_attribute_builder("F64Attr")
+def _f64Attr(x, context):
+  return FloatAttr.get_f64(x, context=context)
+
+
 @register_attribute_builder("StrAttr")
 def _stringAttr(x, context):
   return StringAttr.get(x, context=context)
@@ -61,11 +81,26 @@ def _arrayAttr(x, context):
   return ArrayAttr.get(x, context=context)
 
 
+ at register_attribute_builder("I32ArrayAttr")
+def _i32ArrayAttr(x, context):
+  return ArrayAttr.get([_i32Attr(v, context) for v in x])
+
+
 @register_attribute_builder("I64ArrayAttr")
 def _i64ArrayAttr(x, context):
   return ArrayAttr.get([_i64Attr(v, context) for v in x])
 
 
+ at register_attribute_builder("F32ArrayAttr")
+def _f32ArrayAttr(x, context):
+  return ArrayAttr.get([_f32Attr(v, context) for v in x])
+
+
+ at register_attribute_builder("F64ArrayAttr")
+def _f64ArrayAttr(x, context):
+  return ArrayAttr.get([_f64Attr(v, context) for v in x])
+
+
 @register_attribute_builder("DenseI64ArrayAttr")
 def _denseI64ArrayAttr(x, context):
   return DenseI64ArrayAttr.get(x, context=context)

diff  --git a/mlir/test/python/dialects/python_test.py b/mlir/test/python/dialects/python_test.py
index 8280e5ec73a76..6cde96e1da10d 100644
--- a/mlir/test/python/dialects/python_test.py
+++ b/mlir/test/python/dialects/python_test.py
@@ -131,6 +131,27 @@ def testAttributes():
     del op.unit
     print(f"Unit: {op.unit}")
 
+# CHECK-LABEL: TEST: attrBuilder
+ at run
+def attrBuilder():
+  with Context() as ctx, Location.unknown():
+    ctx.allow_unregistered_dialects = True
+    op = test.AttributesOp(x_bool=True,
+                           x_i16=1,
+                           x_i32=2,
+                           x_i64=3,
+                           x_si16=-1,
+                           x_si32=-2,
+                           x_f32=1.5,
+                           x_f64=2.5,
+                           x_str='x_str',
+                           x_i32_array=[1, 2, 3],
+                           x_i64_array=[4, 5, 6],
+                           x_f32_array=[1.5, -2.5, 3.5],
+                           x_f64_array=[4.5, 5.5, -6.5],
+                           x_i64_dense=[1, 2, 3, 4, 5, 6])
+    print(op)
+
 
 # CHECK-LABEL: TEST: inferReturnTypes
 @run

diff  --git a/mlir/test/python/python_test_ops.td b/mlir/test/python/python_test_ops.td
index e1a03c6ee217d..1be8415733ffb 100644
--- a/mlir/test/python/python_test_ops.td
+++ b/mlir/test/python/python_test_ops.td
@@ -57,6 +57,23 @@ def AttributedOp : TestOp<"attributed_op"> {
                    UnitAttr:$unit);
 }
 
+def AttributesOp : TestOp<"attributes_op"> {
+  let arguments = (ins BoolAttr:$x_bool,
+                   I16Attr: $x_i16,
+                   I32Attr: $x_i32,
+                   I64Attr: $x_i64,
+                   SI16Attr: $x_si16,
+                   SI32Attr: $x_si32,
+                   F32Attr: $x_f32,
+                   F64Attr: $x_f64,
+                   StrAttr: $x_str,
+                   I32ArrayAttr: $x_i32_array,
+                   I64ArrayAttr: $x_i64_array,
+                   F32ArrayAttr: $x_f32_array,
+                   F64ArrayAttr: $x_f64_array,
+                   DenseI64ArrayAttr: $x_i64_dense);
+}
+
 def PropertyOp : TestOp<"property_op"> {
   let arguments = (ins I32Attr:$property,
                    I32:$idx);