[mlir][python] Provide more convenient constructors for std.CallOp

The new constructor relies on type-based dynamic dispatch and allows one to
construct call operations given an object representing a FuncOp or its name as
a string, as opposed to requiring an explicitly constructed attribute.

Depends On D110947

Reviewed By: stellaraccident

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

2 files changed, 80 insertions, 6 deletions

diff --git a/mlir/python/mlir/dialects/_builtin_ops_ext.py b/mlir/python/mlir/dialects/_builtin_ops_ext.py
index d464819f2dc8..462850d63156 100644
--- a/mlir/python/mlir/dialects/_builtin_ops_ext.py
+++ b/mlir/python/mlir/dialects/_builtin_ops_ext.py
@@ -3,7 +3,7 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 try:
-  from typing import Optional, Sequence
+  from typing import Optional, Sequence, Union
 
   import inspect
 
@@ -82,8 +82,8 @@ class FuncOp:
     return self.attributes["sym_visibility"]
 
   @property
-  def name(self):
-    return self.attributes["sym_name"]
+  def name(self) -> StringAttr:
+    return StringAttr(self.attributes["sym_name"])
 
   @property
   def entry_block(self):
@@ -104,11 +104,15 @@ class FuncOp:
 
   @property
   def arg_attrs(self):
-    return self.attributes[ARGUMENT_ATTRIBUTE_NAME]
+    return ArrayAttr(self.attributes[ARGUMENT_ATTRIBUTE_NAME])
 
   @arg_attrs.setter
-  def arg_attrs(self, attribute: ArrayAttr):
-    self.attributes[ARGUMENT_ATTRIBUTE_NAME] = attribute
+  def arg_attrs(self, attribute: Union[ArrayAttr, list]):
+    if isinstance(attribute, ArrayAttr):
+      self.attributes[ARGUMENT_ATTRIBUTE_NAME] = attribute
+    else:
+      self.attributes[ARGUMENT_ATTRIBUTE_NAME] = ArrayAttr.get(
+          attribute, context=self.context)
 
   @property
   def arguments(self):
diff --git a/mlir/python/mlir/dialects/_std_ops_ext.py b/mlir/python/mlir/dialects/_std_ops_ext.py
index bb67fe44d3b5..39da1c83f449 100644
--- a/mlir/python/mlir/dialects/_std_ops_ext.py
+++ b/mlir/python/mlir/dialects/_std_ops_ext.py
@@ -69,3 +69,73 @@ class ConstantOp:
       return FloatAttr(self.value).value
     else:
       raise ValueError("only integer and float constants have literal values")
+
+
+class CallOp:
+  """Specialization for the call op class."""
+
+  def __init__(self,
+               calleeOrResults: Union[FuncOp, List[Type]],
+               argumentsOrCallee: Union[List, FlatSymbolRefAttr, str],
+               arguments: Optional[List] = None,
+               *,
+               loc=None,
+               ip=None):
+    """Creates an call operation.
+
+    The constructor accepts three different forms:
+
+      1. A function op to be called followed by a list of arguments.
+      2. A list of result types, followed by the name of the function to be
+         called as string, following by a list of arguments.
+      3. A list of result types, followed by the name of the function to be
+         called as symbol reference attribute, followed by a list of arguments.
+
+    For example
+
+        f = builtin.FuncOp("foo", ...)
+        std.CallOp(f, [args])
+        std.CallOp([result_types], "foo", [args])
+
+    In all cases, the location and insertion point may be specified as keyword
+    arguments if not provided by the surrounding context managers.
+    """
+
+    # TODO: consider supporting constructor "overloads", e.g., through a custom
+    # or pybind-provided metaclass.
+    if isinstance(calleeOrResults, FuncOp):
+      if not isinstance(argumentsOrCallee, list):
+        raise ValueError(
+            "when constructing a call to a function, expected " +
+            "the second argument to be a list of call arguments, " +
+            f"got {type(argumentsOrCallee)}")
+      if arguments is not None:
+        raise ValueError("unexpected third argument when constructing a call" +
+                         "to a function")
+
+      super().__init__(
+          calleeOrResults.type.results,
+          FlatSymbolRefAttr.get(
+              calleeOrResults.name.value,
+              context=_get_default_loc_context(loc)),
+          argumentsOrCallee,
+          loc=loc,
+          ip=ip)
+      return
+
+    if isinstance(argumentsOrCallee, list):
+      raise ValueError("when constructing a call to a function by name, " +
+                       "expected the second argument to be a string or a " +
+                       f"FlatSymbolRefAttr, got {type(argumentsOrCallee)}")
+
+    if isinstance(argumentsOrCallee, FlatSymbolRefAttr):
+      super().__init__(
+          calleeOrResults, argumentsOrCallee, arguments, loc=loc, ip=ip)
+    elif isinstance(argumentsOrCallee, str):
+      super().__init__(
+          calleeOrResults,
+          FlatSymbolRefAttr.get(
+              argumentsOrCallee, context=_get_default_loc_context(loc)),
+          arguments,
+          loc=loc,
+          ip=ip)