[mlir] Introduce Python bindings for the PDL dialect

This change adds full python bindings for PDL, including types and operations
with additional mixins to make operation construction more similar to the PDL
syntax.

Reviewed By: ftynse

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

Author: shabalind

mlir/include/mlir-c/Dialect/PDL.h

@@ -49,6 +49,8 @@ MLIR_CAPI_EXPORTED bool mlirTypeIsAPDLRangeType(MlirType type);
 
 MLIR_CAPI_EXPORTED MlirType mlirPDLRangeTypeGet(MlirType elementType);
 
+MLIR_CAPI_EXPORTED MlirType mlirPDLRangeTypeGetElementType(MlirType type);
+
 //===---------------------------------------------------------------------===//
 // TypeType
 //===---------------------------------------------------------------------===//

mlir/lib/Bindings/Python/DialectPDL.cpp

@@ -0,0 +1,102 @@
+//===- DialectPDL.cpp - 'pdl' dialect submodule ---------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/PDL.h"
+#include "mlir-c/IR.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
+
+namespace py = pybind11;
+using namespace llvm;
+using namespace mlir;
+using namespace mlir::python;
+using namespace mlir::python::adaptors;
+
+void populateDialectPDLSubmodule(const pybind11::module &m) {
+  //===-------------------------------------------------------------------===//
+  // PDLType
+  //===-------------------------------------------------------------------===//
+
+  auto pdlType = mlir_type_subclass(m, "PDLType", mlirTypeIsAPDLType);
+
+  //===-------------------------------------------------------------------===//
+  // AttributeType
+  //===-------------------------------------------------------------------===//
+
+  auto attributeType =
+      mlir_type_subclass(m, "AttributeType", mlirTypeIsAPDLAttributeType);
+  attributeType.def_classmethod(
+      "get",
+      [](py::object cls, MlirContext ctx) {
+        return cls(mlirPDLAttributeTypeGet(ctx));
+      },
+      "Get an instance of AttributeType in given context.", py::arg("cls"),
+      py::arg("context") = py::none());
+
+  //===-------------------------------------------------------------------===//
+  // OperationType
+  //===-------------------------------------------------------------------===//
+
+  auto operationType =
+      mlir_type_subclass(m, "OperationType", mlirTypeIsAPDLOperationType);
+  operationType.def_classmethod(
+      "get",
+      [](py::object cls, MlirContext ctx) {
+        return cls(mlirPDLOperationTypeGet(ctx));
+      },
+      "Get an instance of OperationType in given context.", py::arg("cls"),
+      py::arg("context") = py::none());
+
+  //===-------------------------------------------------------------------===//
+  // RangeType
+  //===-------------------------------------------------------------------===//
+
+  auto rangeType = mlir_type_subclass(m, "RangeType", mlirTypeIsAPDLRangeType);
+  rangeType.def_classmethod(
+      "get",
+      [](py::object cls, MlirType elementType) {
+        return cls(mlirPDLRangeTypeGet(elementType));
+      },
+      "Gets an instance of RangeType in the same context as the provided "
+      "element type.",
+      py::arg("cls"), py::arg("element_type"));
+  rangeType.def_property_readonly(
+      "element_type",
+      [](MlirType type) { return mlirPDLRangeTypeGetElementType(type); },
+      "Get the element type.");
+
+  //===-------------------------------------------------------------------===//
+  // TypeType
+  //===-------------------------------------------------------------------===//
+
+  auto typeType = mlir_type_subclass(m, "TypeType", mlirTypeIsAPDLTypeType);
+  typeType.def_classmethod(
+      "get",
+      [](py::object cls, MlirContext ctx) {
+        return cls(mlirPDLTypeTypeGet(ctx));
+      },
+      "Get an instance of TypeType in given context.", py::arg("cls"),
+      py::arg("context") = py::none());
+
+  //===-------------------------------------------------------------------===//
+  // ValueType
+  //===-------------------------------------------------------------------===//
+
+  auto valueType = mlir_type_subclass(m, "ValueType", mlirTypeIsAPDLValueType);
+  valueType.def_classmethod(
+      "get",
+      [](py::object cls, MlirContext ctx) {
+        return cls(mlirPDLValueTypeGet(ctx));
+      },
+      "Get an instance of TypeType in given context.", py::arg("cls"),
+      py::arg("context") = py::none());
+}
+
+PYBIND11_MODULE(_mlirDialectsPDL, m) {
+  m.doc() = "MLIR PDL dialect.";
+  populateDialectPDLSubmodule(m);
+}

mlir/lib/CAPI/Dialect/PDL.cpp

@@ -60,6 +60,10 @@ MlirType mlirPDLRangeTypeGet(MlirType elementType) {
   return wrap(pdl::RangeType::get(unwrap(elementType)));
 }
 
+MlirType mlirPDLRangeTypeGetElementType(MlirType type) {
+  return wrap(unwrap(type).cast<pdl::RangeType>().getElementType());
+}
+
 //===---------------------------------------------------------------------===//
 // TypeType
 //===---------------------------------------------------------------------===//

mlir/python/CMakeLists.txt

@@ -123,6 +123,15 @@ declare_mlir_python_sources(
     dialects/quant.py
     _mlir_libs/_mlir/dialects/quant.pyi)
 
+declare_mlir_python_sources(
+  MLIRPythonSources.Dialects.pdl
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  SOURCES
+    dialects/pdl.py
+    dialects/_pdl_ops_ext.py
+    _mlir_libs/_mlir/dialects/pdl.pyi)
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
@@ -243,6 +252,19 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
     MLIRCAPIQuant
 )
 
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
+  MODULE_NAME _mlirDialectsPDL
+  ADD_TO_PARENT MLIRPythonSources.Dialects.pdl
+  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  SOURCES
+    DialectPDL.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIIR
+    MLIRCAPIPDL
+)
+
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
   MODULE_NAME _mlirDialectsSparseTensor
   ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor

mlir/python/mlir/_mlir_libs/_mlir/dialects/pdl.pyi

@@ -0,0 +1,64 @@
+#  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 typing import Optional
+
+from mlir.ir import Type, Context
+
+__all__ = [
+    'PDLType',
+    'AttributeType',
+    'OperationType',
+    'RangeType',
+    'TypeType',
+    'ValueType',
+]
+
+
+class PDLType(Type):
+  @staticmethod
+  def isinstance(type: Type) -> bool: ...
+
+
+class AttributeType(Type):
+  @staticmethod
+  def isinstance(type: Type) -> bool: ...
+
+  @staticmethod
+  def get(context: Optional[Context] = None) -> AttributeType: ...
+
+
+class OperationType(Type):
+  @staticmethod
+  def isinstance(type: Type) -> bool: ...
+
+  @staticmethod
+  def get(context: Optional[Context] = None) -> OperationType: ...
+
+
+class RangeType(Type):
+  @staticmethod
+  def isinstance(type: Type) -> bool: ...
+
+  @staticmethod
+  def get(element_type: Type) -> RangeType: ...
+
+  @property
+  def element_type(self) -> Type: ...
+
+
+class TypeType(Type):
+  @staticmethod
+  def isinstance(type: Type) -> bool: ...
+
+  @staticmethod
+  def get(context: Optional[Context] = None) -> TypeType: ...
+
+
+class ValueType(Type):
+  @staticmethod
+  def isinstance(type: Type) -> bool: ...
+
+  @staticmethod
+  def get(context: Optional[Context] = None) -> ValueType: ...

mlir/python/mlir/dialects/PDLOps.td

@@ -0,0 +1,15 @@
+//===-- PDLOps.td - Entry point for PDLOps bind ------------*- 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_PDL_OPS
+#define PYTHON_BINDINGS_PDL_OPS
+
+include "mlir/Bindings/Python/Attributes.td"
+include "mlir/Dialect/PDL/IR/PDLOps.td"
+
+#endif

mlir/python/mlir/dialects/_ods_common.py

@@ -144,7 +144,8 @@ def get_op_result_or_value(
 
 
 def get_op_results_or_values(
-    arg: _Union[_cext.ir.OpView, _cext.ir.Operation, _Sequence[_cext.ir.Value]]
+    arg: _Union[_cext.ir.OpView, _cext.ir.Operation,
+                _Sequence[_Union[_cext.ir.OpView, _cext.ir.Operation, _cext.ir.Value]]]
 ) -> _Union[_Sequence[_cext.ir.Value], _cext.ir.OpResultList]:
   """Returns the given sequence of values or the results of the given op.
 
@@ -157,4 +158,4 @@ def get_op_results_or_values(
   elif isinstance(arg, _cext.ir.Operation):
     return arg.results
   else:
-    return arg
+    return [get_op_result_or_value(element) for element in arg]