[mlir][Python] Roll up of python API fixes.

* As discussed, fixes the ordering or (operands, results) -> (results, operands) in various `create` like methods.
* Fixes a syntax error in an ODS accessor method.
* Removes the linalg example in favor of a test case that exercises the same.
* Fixes FuncOp visibility to properly use None instead of the empty string and defaults it to None.
* Implements what was documented for requiring that trailing __init__ args `loc` and `ip` are keyword only.
* Adds a check to `InsertionPoint.insert` so that if attempting to insert past the terminator, an exception is raised telling you what to do instead. Previously, this would crash downstream (i.e. when trying to print the resultant module).
* Renames `_ods_build_default` -> `build_generic` and documents it.
* Removes `result` from the list of prohibited words and for single-result ops, defaults to naming the result `result`, thereby matching expectations and what is already implemented on the base class.
* This was intended to be a relatively small set of changes to be inlined with the broader support for ODS generating the most specific builder, but it spidered out once actually testing various combinations, so rolling up separately.

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

Author: stellaraccident

mlir/docs/Bindings/Python.md

@@ -439,8 +439,9 @@ defaults on `OpView`):
 #### Builders
 
 Presently, only a single, default builder is mapped to the `__init__` method.
-Generalizing this facility is under active development. It currently accepts
-arguments:
+The intent is that this `__init__` method represents the *most specific* of
+the builders typically generated for C++; however currently it is just the
+generic form below.
 
 * One argument for each declared result:
   * For single-valued results: Each will accept an `mlir.ir.Type`.
@@ -453,7 +454,11 @@ arguments:
   * `loc`: An explicit `mlir.ir.Location` to use. Defaults to the location
     bound to the thread (i.e. `with Location.unknown():`) or an error if none
     is bound nor specified.
-  * `context`: An explicit `mlir.ir.Context` to use. Default to the context
-    bound to the thread (i.e. `with Context():` or implicitly via `Location` or
-    `InsertionPoint` context managers) or an error if none is bound nor
-    specified.
+  * `ip`: An explicit `mlir.ir.InsertionPoint` to use. Default to the insertion
+    point bound to the thread (i.e. `with InsertionPoint(...):`).
+
+In addition, each `OpView` inherits a `build_generic` method which allows
+construction via a (nested in the case of variadic) sequence of `results` and
+`operands`. This can be used to get some default construction semantics for
+operations that are otherwise unsupported in Python, at the expense of having
+a very generic signature.

mlir/examples/python/.style.yapf

@@ -891,8 +891,8 @@ PyBlock PyOperation::getBlock() {
 }
 
 py::object PyOperation::create(
-    std::string name, llvm::Optional<std::vector<PyValue *>> operands,
-    llvm::Optional<std::vector<PyType *>> results,
+    std::string name, llvm::Optional<std::vector<PyType *>> results,
+    llvm::Optional<std::vector<PyValue *>> operands,
     llvm::Optional<py::dict> attributes,
     llvm::Optional<std::vector<PyBlock *>> successors, int regions,
     DefaultingPyLocation location, py::object maybeIp) {
@@ -1039,12 +1039,12 @@ py::object PyOperation::createOpView() {
 //------------------------------------------------------------------------------
 
 py::object
-PyOpView::odsBuildDefault(py::object cls, py::list operandList,
-                          py::list resultTypeList,
-                          llvm::Optional<py::dict> attributes,
-                          llvm::Optional<std::vector<PyBlock *>> successors,
-                          llvm::Optional<int> regions,
-                          DefaultingPyLocation location, py::object maybeIp) {
+PyOpView::buildGeneric(py::object cls, py::list resultTypeList,
+                       py::list operandList,
+                       llvm::Optional<py::dict> attributes,
+                       llvm::Optional<std::vector<PyBlock *>> successors,
+                       llvm::Optional<int> regions,
+                       DefaultingPyLocation location, py::object maybeIp) {
   PyMlirContextRef context = location->getContext();
   // Class level operation construction metadata.
   std::string name = py::cast<std::string>(cls.attr("OPERATION_NAME"));
@@ -1288,8 +1288,9 @@ PyOpView::odsBuildDefault(py::object cls, py::list operandList,
   }
 
   // Delegate to create.
-  return PyOperation::create(std::move(name), /*operands=*/std::move(operands),
+  return PyOperation::create(std::move(name),
                              /*results=*/std::move(resultTypes),
+                             /*operands=*/std::move(operands),
                              /*attributes=*/std::move(attributes),
                              /*successors=*/std::move(successors),
                              /*regions=*/*regions, location, maybeIp);
@@ -1357,6 +1358,16 @@ void PyInsertionPoint::insert(PyOperationBase &operationBase) {
     // Insert before operation.
     (*refOperation)->checkValid();
     beforeOp = (*refOperation)->get();
+  } else {
+    // Insert at end (before null) is only valid if the block does not
+    // already end in a known terminator (violating this will cause assertion
+    // failures later).
+    if (!mlirOperationIsNull(mlirBlockGetTerminator(block.get()))) {
+      throw py::index_error("Cannot insert operation at the end of a block "
+                            "that already has a terminator. Did you mean to "
+                            "use 'InsertionPoint.at_block_terminator(block)' "
+                            "versus 'InsertionPoint(block)'?");
+    }
   }
   mlirBlockInsertOwnedOperationBefore(block.get(), beforeOp, operation);
   operation.setAttached();
@@ -3646,8 +3657,8 @@ void mlir::python::populateIRSubmodule(py::module &m) {
 
   py::class_<PyOperation, PyOperationBase>(m, "Operation")
       .def_static("create", &PyOperation::create, py::arg("name"),
-                  py::arg("operands") = py::none(),
                   py::arg("results") = py::none(),
+                  py::arg("operands") = py::none(),
                   py::arg("attributes") = py::none(),
                   py::arg("successors") = py::none(), py::arg("regions") = 0,
                   py::arg("loc") = py::none(), py::arg("ip") = py::none(),
@@ -3681,12 +3692,11 @@ void mlir::python::populateIRSubmodule(py::module &m) {
   opViewClass.attr("_ODS_REGIONS") = py::make_tuple(0, true);
   opViewClass.attr("_ODS_OPERAND_SEGMENTS") = py::none();
   opViewClass.attr("_ODS_RESULT_SEGMENTS") = py::none();
-  opViewClass.attr("_ods_build_default") = classmethod(
-      &PyOpView::odsBuildDefault, py::arg("cls"),
-      py::arg("operands") = py::none(), py::arg("results") = py::none(),
-      py::arg("attributes") = py::none(), py::arg("successors") = py::none(),
-      py::arg("regions") = py::none(), py::arg("loc") = py::none(),
-      py::arg("ip") = py::none(),
+  opViewClass.attr("build_generic") = classmethod(
+      &PyOpView::buildGeneric, py::arg("cls"), py::arg("results") = py::none(),
+      py::arg("operands") = py::none(), py::arg("attributes") = py::none(),
+      py::arg("successors") = py::none(), py::arg("regions") = py::none(),
+      py::arg("loc") = py::none(), py::arg("ip") = py::none(),
       "Builds a specific, generated OpView based on class level attributes.");
 
   //----------------------------------------------------------------------------

mlir/examples/python/linalg_matmul.py

@@ -455,8 +455,8 @@ class PyOperation : public PyOperationBase, public BaseContextObject {
 
   /// Creates an operation. See corresponding python docstring.
   static pybind11::object
-  create(std::string name, llvm::Optional<std::vector<PyValue *>> operands,
-         llvm::Optional<std::vector<PyType *>> results,
+  create(std::string name, llvm::Optional<std::vector<PyType *>> results,
+         llvm::Optional<std::vector<PyValue *>> operands,
          llvm::Optional<pybind11::dict> attributes,
          llvm::Optional<std::vector<PyBlock *>> successors, int regions,
          DefaultingPyLocation location, pybind11::object ip);
@@ -498,12 +498,12 @@ class PyOpView : public PyOperationBase {
   pybind11::object getOperationObject() { return operationObject; }
 
   static pybind11::object
-  odsBuildDefault(pybind11::object cls, pybind11::list operandList,
-                  pybind11::list resultTypeList,
-                  llvm::Optional<pybind11::dict> attributes,
-                  llvm::Optional<std::vector<PyBlock *>> successors,
-                  llvm::Optional<int> regions, DefaultingPyLocation location,
-                  pybind11::object maybeIp);
+  buildGeneric(pybind11::object cls, pybind11::list resultTypeList,
+               pybind11::list operandList,
+               llvm::Optional<pybind11::dict> attributes,
+               llvm::Optional<std::vector<PyBlock *>> successors,
+               llvm::Optional<int> regions, DefaultingPyLocation location,
+               pybind11::object maybeIp);
 
 private:
   PyOperation &operation;           // For efficient, cast-free access from C++

mlir/lib/Bindings/Python/IRModules.cpp

@@ -7,9 +7,9 @@
 class ModuleOp:
   """Specialization for the module op class."""
 
-  def __init__(self, loc=None, ip=None):
+  def __init__(self, *, loc=None, ip=None):
     super().__init__(
-        self._ods_build_default(operands=[], results=[], loc=loc, ip=ip))
+        self.build_generic(results=[], operands=[], loc=loc, ip=ip))
     body = self.regions[0].blocks.append()
     with InsertionPoint(body):
       Operation.create("module_terminator")
@@ -25,7 +25,8 @@ class FuncOp:
   def __init__(self,
                name,
                type,
-               visibility,
+               *,
+               visibility=None,
                body_builder=None,
                loc=None,
                ip=None):
@@ -34,8 +35,8 @@ def __init__(self,
     - `name` is a string representing the function name.
     - `type` is either a FunctionType or a pair of list describing inputs and
       results.
-    - `visibility` is a string matching `public`, `private`, or `nested`. The
-      empty string implies a private visibility.
+    - `visibility` is a string matching `public`, `private`, or `nested`. None
+      implies private visibility.
     - `body_builder` is an optional callback, when provided a new entry block
       is created and the callback is invoked with the new op as argument within
       an InsertionPoint context already set for the block. The callback is
@@ -50,7 +51,7 @@ def __init__(self,
     type = TypeAttr.get(type)
     sym_visibility = StringAttr.get(
         str(visibility)) if visibility is not None else None
-    super().__init__(sym_name, type, sym_visibility, loc, ip)
+    super().__init__(sym_name, type, sym_visibility, loc=loc, ip=ip)
     if body_builder:
       entry_block = self.add_entry_block()
       with InsertionPoint(entry_block):

mlir/lib/Bindings/Python/IRModules.h

@@ -11,11 +11,10 @@ def __init__(self, inputs, outputs=(), results=(), loc=None, ip=None):
       raise ValueError(
           "Structured ops must have outputs or results, but not both.")
     super().__init__(
-        self._ods_build_default(operands=[list(inputs),
-                                          list(outputs)],
-                                results=list(results),
-                                loc=loc,
-                                ip=ip))
+        self.build_generic(results=list(results),
+                           operands=[list(inputs), list(outputs)],
+                           loc=loc,
+                           ip=ip))
 
 
 def select_opview_mixin(parent_opview_cls):

mlir/lib/Bindings/Python/mlir/dialects/_builtin.py

@@ -0,0 +1,57 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+from mlir.dialects import builtin
+from mlir.dialects import linalg
+from mlir.dialects import std
+
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+
+
+# CHECK-LABEL: TEST: testStructuredOpOnTensors
+def testStructuredOpOnTensors():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    f32 = F32Type.get()
+    tensor_type = RankedTensorType.get((2, 3, 4), f32)
+    with InsertionPoint.at_block_terminator(module.body):
+      func = builtin.FuncOp(name="matmul_test",
+                            type=FunctionType.get(
+                                inputs=[tensor_type, tensor_type],
+                                results=[tensor_type]))
+      with InsertionPoint(func.add_entry_block()):
+        lhs, rhs = func.entry_block.arguments
+        result = linalg.MatmulOp([lhs, rhs], results=[tensor_type]).result
+        std.ReturnOp([result])
+
+  # CHECK: %[[R:.*]] = linalg.matmul ins(%arg0, %arg1 : tensor<2x3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32>
+  print(module)
+
+
+run(testStructuredOpOnTensors)
+
+
+# CHECK-LABEL: TEST: testStructuredOpOnBuffers
+def testStructuredOpOnBuffers():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    f32 = F32Type.get()
+    memref_type = MemRefType.get((2, 3, 4), f32)
+    with InsertionPoint.at_block_terminator(module.body):
+      func = builtin.FuncOp(name="matmul_test",
+                            type=FunctionType.get(
+                                inputs=[memref_type, memref_type, memref_type],
+                                results=[]))
+      with InsertionPoint(func.add_entry_block()):
+        lhs, rhs, result = func.entry_block.arguments
+        linalg.MatmulOp([lhs, rhs], outputs=[result])
+        std.ReturnOp([])
+
+  # CHECK: linalg.matmul ins(%arg0, %arg1 : memref<2x3x4xf32>, memref<2x3x4xf32>) outs(%arg2 : memref<2x3x4xf32>)
+  print(module)
+
+
+run(testStructuredOpOnBuffers)

mlir/lib/Bindings/Python/mlir/dialects/_linalg.py

@@ -125,6 +125,21 @@ def test_insert_at_block_terminator_missing():
 run(test_insert_at_block_terminator_missing)
 
 
+# CHECK-LABEL: TEST: test_insert_at_end_with_terminator_errors
+def test_insert_at_end_with_terminator_errors():
+  with Context() as ctx, Location.unknown():
+    ctx.allow_unregistered_dialects = True
+    m = Module.create()  # Module is created with a terminator.
+    with InsertionPoint(m.body):
+      try:
+        Operation.create("custom.op1", results=[], operands=[])
+      except IndexError as e:
+        # CHECK: ERROR: Cannot insert operation at the end of a block that already has a terminator.
+        print(f"ERROR: {e}")
+
+run(test_insert_at_end_with_terminator_errors)
+
+
 # CHECK-LABEL: TEST: test_insertion_point_context
 def test_insertion_point_context():
   ctx = Context()