Add Python bindings for the builtin dialect

[mehdi_amini]

Summary

This includes some minor customization for FuncOp and ModuleOp.

[rriddle]

mlir/lib/Bindings/Python/BuiltinOps.td

2

Linalg?

[stellaraccident]

Thanks. Lg modulo comments

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

30

Prefer the ternary form:

sym_visibility = ... if visibility is not None else None

35

Doesn't this return the region on the c++ side?

38

See if it works with this removed. It should bind by class name by default.

[stellaraccident]

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

35

Also, functions without a body are legal

[ftynse]

mlir/test/Bindings/Python/dialects/builtin.py

25

Maybe also define FuncOp.add_body ?

[mehdi_amini]

Address comment, add more helpers/properties on FuncOp, extend the test

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

35

Good point! PTAL

38

Indeed, I had to fix a bug in the class_decorator magic but now it works :)

mlir/test/Bindings/Python/dialects/builtin.py

25

add_entry_block to align with C++

[mehdi_amini]

Minor update to add_entry_block + test fix

[mehdi_amini]

Add a body_builder callback

[mehdi_amini]

Landing now, feel free to add more comments as needed!

mlir/lib/Bindings/Python/BuiltinOps.td

+//===-- BuiltinOps.td - Entry point for builtin bindings ---*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_BUILTIN_OPS
+#define PYTHON_BINDINGS_BUILTIN_OPS
+
+include "mlir/Bindings/Python/Attributes.td"
+include "mlir/IR/BuiltinOps.td"
+
+#endif

mlir/lib/Bindings/Python/CMakeLists.txt

 include(AddMLIRPythonExtension)
 add_custom_target(MLIRBindingsPythonExtension)
 
 ################################################################################
 # Copy python source tree.
 ################################################################################
 
 set(PY_SRC_FILES
   mlir/__init__.py
   mlir/_dlloader.py
   mlir/ir.py
   mlir/dialects/__init__.py
   mlir/dialects/_linalg.py
+  mlir/dialects/_builtin.py
   mlir/ir.py
   mlir/passmanager.py
   mlir/transforms/__init__.py
 )
 
 add_custom_target(MLIRBindingsPythonSources ALL
   DEPENDS ${PY_SRC_FILES}
 )
COMMAND "${CMAKE_COMMAND}" -E copy_if_different
         "${CMAKE_CURRENT_SOURCE_DIR}/${PY_SRC_FILE}" "${PY_DEST_FILE}"
   )
 endforeach()
 
 ################################################################################
 # Generate dialect-specific bindings.
 ################################################################################
 
+add_mlir_dialect_python_bindings(MLIRBindingsPythonBuiltinOps
+  TD_FILE BuiltinOps.td
+  DIALECT_NAME builtin)
+add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonBuiltinOps)
+
 add_mlir_dialect_python_bindings(MLIRBindingsPythonLinalgOps
   TD_FILE LinalgOps.td
   DIALECT_NAME linalg
   DEPENDS LinalgOdsGen)
 add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonLinalgOps)
 
 add_mlir_dialect_python_bindings(MLIRBindingsPythonShapeOps
   TD_FILE ShapeOps.td

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

def class_decorator(parent_opview_cls: type):
     if ext_module is None:
       return parent_opview_cls
     mixin_cls = NotImplemented
     try:
       select_mixin = getattr(ext_module, "select_opview_mixin")
     except AttributeError:
       # Try to default resolve it.
       try:
-        select_mixin = getattr(ext_module, parent_opview_cls.__name__)
+        mixin_cls = getattr(ext_module, parent_opview_cls.__name__)
       except AttributeError:
         pass
     else:
       mixin_cls = select_mixin(parent_opview_cls)
     if mixin_cls is NotImplemented or mixin_cls is None:
       return parent_opview_cls
 
     # Have a mixin_cls. Create an appropriate subclass.

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

+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+from mlir.ir import *
+
+
+class ModuleOp:
+  """Specialization for the module op class."""
+
+  def __init__(self, loc=None, ip=None):
+    super().__init__(
+        self._ods_build_default(operands=[], results=[], loc=loc, ip=ip))
+    body = self.regions[0].blocks.append()
+    with InsertionPoint(body):
+      Operation.create("module_terminator")
+
+  @property
+  def body(self):
+    return self.regions[0].blocks[0]
+
+
+class FuncOp:
+  """Specialization for the func op class."""
+
+  def __init__(self,
+               name,
+               type,
+               visibility,
+               body_builder=None,
+               loc=None,
+               ip=None):
+    """
+    Create a FuncOp with the provided `name`, `type`, and `visibility`.
+    - `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.
+    - `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
+      expected to insert a terminator in the block.
+    """
+    sym_name = StringAttr.get(str(name))
+
+    # If the type is passed as a tuple, build a FunctionType on the fly.
+    if isinstance(type, tuple):
+      type = FunctionType.get(inputs=type[0], results=type[1])
+
+    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)
+    if body_builder:
+      entry_block = self.add_entry_block()
+      with InsertionPoint(entry_block):
+        body_builder(self)
+
+  @property
+  def is_external(self):
+    return len(self.regions[0].blocks) == 0
+
+  @property
+  def body(self):
+    return self.regions[0]
+
+  @property
+  def type(self):
+    return FunctionType(TypeAttr(self.attributes["type"]).value)
+
+  @property
+  def visibility(self):
+    return self.attributes["sym_visibility"]
+
+  @property
+  def name(self):
+    return self.attributes["sym_name"]
+
+  @property
+  def entry_block(self):
+    if self.is_external:
+      raise IndexError('External function does not have a body')
+    return self.regions[0].blocks[0]
+
+  def add_entry_block(self):
+    '''
+    Add an entry block to the function body using the function signature to infer block arguments
+    Returns the newly created block
+    '''
+    if not self.is_external:
+      raise IndexError('The function already has an entry block!')
+    self.body.blocks.append(*self.type.inputs)
+    return self.body.blocks[0]

mlir/test/Bindings/Python/.style.yapf

+[style]
+  based_on_style = google
+  column_limit = 80
+  indent_width = 2

mlir/test/Bindings/Python/dialects/builtin.py

+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+import mlir.dialects.builtin as builtin
+import mlir.dialects.std as std
+
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+
+
+# CHECK-LABEL: TEST: testBuildFuncOp
+def testBuildFuncOp():
+  ctx = Context()
+  with Location.unknown(ctx) as loc:
+    m = builtin.ModuleOp()
+
+    f32 = F32Type.get()
+    tensor_type = RankedTensorType.get((2, 3, 4), f32)
+    with InsertionPoint.at_block_begin(m.body):
+      func = builtin.FuncOp(name="some_func",
+                            type=FunctionType.get(
+                                inputs=[tensor_type, tensor_type],
+                                results=[tensor_type]),
+                            visibility="nested")
+      # CHECK: Name is: "some_func"
+      print("Name is: ", func.name)
+
+      # CHECK: Type is: (tensor<2x3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32>
+      print("Type is: ", func.type)
+
+      # CHECK: Visibility is: "nested"
+      print("Visibility is: ", func.visibility)
+
+      try:
+        entry_block = func.entry_block
+      except IndexError as e:
+        # CHECK: External function does not have a body
+        print(e)
+
+      with InsertionPoint(func.add_entry_block()):
+        std.ReturnOp([func.entry_block.arguments[0]])
+        pass
+
+      try:
+        func.add_entry_block()
+      except IndexError as e:
+        # CHECK: The function already has an entry block!
+        print(e)
+
+      # Try the callback builder and passing type as tuple.
+      func = builtin.FuncOp(name="some_other_func",
+                            type=([tensor_type, tensor_type], [tensor_type]),
+                            visibility="nested",
+                            body_builder=lambda func: std.ReturnOp(
+                                [func.entry_block.arguments[0]]))
+
+  # CHECK: module  {
+  # CHECK:  func nested @some_func(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
+  # CHECK:   return %arg0 : tensor<2x3x4xf32>
+  # CHECK:  }
+  # CHECK:  func nested @some_other_func(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
+  # CHECK:   return %arg0 : tensor<2x3x4xf32>
+  # CHECK:  }
+  print(m)
+
+
+run(testBuildFuncOp)

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

static bool emitAllOps(const llvm::RecordKeeper &records, raw_ostream &os) {
   if (clDialectName.empty())
     llvm::PrintFatalError("dialect name not provided");
 
   AttributeClasses attributeClasses;
   constructAttributeMapping(records, attributeClasses);
 
   os << llvm::formatv(fileHeader, clDialectName.getValue());
   os << llvm::formatv(dialectClassTemplate, clDialectName.getValue());
+
+  if (clDialectName == "builtin")
+    clDialectName = "";
+
   for (const llvm::Record *rec : records.getAllDerivedDefinitions("Op")) {
     Operator op(rec);
     if (op.getDialectName() == clDialectName.getValue())
       emitOpBindings(op, attributeClasses, os);
   }
   return false;
 }
 
 static GenRegistration
     genPythonBindings("gen-python-op-bindings",
                       "Generate Python bindings for MLIR Ops", &emitAllOps);