[mlir] Split out Python bindings for dialects into separate libs

Historically, the bindings for the Linalg dialect were included into the
"core" bindings library because they depended on the C++ implementation
of the "core" bindings. The other dialects followed the pattern. Now
that this dependency is gone, split out each dialect into a separate
Python extension library.

Depends On D116649, D116605

Reviewed By: stellaraccident

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

Author: ftynse

mlir/lib/Bindings/Python/DialectLinalg.cpp

@@ -6,18 +6,23 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "Dialects.h"
 #include "mlir-c/Dialect/Linalg.h"
 #include "mlir-c/IR.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
 
 namespace py = pybind11;
 
-void mlir::python::populateDialectLinalgSubmodule(py::module m) {
+static void populateDialectLinalgSubmodule(py::module m) {
   m.def(
       "fill_builtin_region",
       [](MlirOperation op) { mlirLinalgFillBuiltinNamedOpRegion(op); },
       py::arg("op"),
       "Fill the region for `op`, which is assumed to be a builtin named Linalg "
       "op.");
 }
+
+PYBIND11_MODULE(_mlirDialectsLinalg, m) {
+  m.doc() = "MLIR Linalg dialect.";
+
+  populateDialectLinalgSubmodule(m);
+}

mlir/lib/Bindings/Python/DialectQuant.cpp

@@ -6,7 +6,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "Dialects.h"
 #include "mlir-c/Dialect/Quant.h"
 #include "mlir-c/IR.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
@@ -16,16 +15,13 @@ using namespace llvm;
 using namespace mlir;
 using namespace mlir::python::adaptors;
 
-void mlir::python::populateDialectQuantSubmodule(const py::module &m,
-                                                 const py::module &irModule) {
-  auto typeClass = irModule.attr("Type");
-
+static void populateDialectQuantSubmodule(const py::module &m) {
   //===-------------------------------------------------------------------===//
   // QuantizedType
   //===-------------------------------------------------------------------===//
 
-  auto quantizedType = mlir_type_subclass(m, "QuantizedType",
-                                          mlirTypeIsAQuantizedType, typeClass);
+  auto quantizedType =
+      mlir_type_subclass(m, "QuantizedType", mlirTypeIsAQuantizedType);
   quantizedType.def_staticmethod(
       "default_minimum_for_integer",
       [](bool isSigned, unsigned integralWidth) {
@@ -305,3 +301,9 @@ void mlir::python::populateDialectQuantSubmodule(const py::module &m,
     return mlirCalibratedQuantizedTypeGetMax(type);
   });
 }
+
+PYBIND11_MODULE(_mlirDialectsQuant, m) {
+  m.doc() = "MLIR Quantization dialect";
+
+  populateDialectQuantSubmodule(m);
+}

mlir/lib/Bindings/Python/DialectSparseTensor.cpp

@@ -6,7 +6,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "Dialects.h"
 #include "mlir-c/Dialect/SparseTensor.h"
 #include "mlir-c/IR.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
@@ -16,18 +15,14 @@ using namespace llvm;
 using namespace mlir;
 using namespace mlir::python::adaptors;
 
-void mlir::python::populateDialectSparseTensorSubmodule(
-    const py::module &m, const py::module &irModule) {
-  auto attributeClass = irModule.attr("Attribute");
-
+static void populateDialectSparseTensorSubmodule(const py::module &m) {
   py::enum_<MlirSparseTensorDimLevelType>(m, "DimLevelType", py::module_local())
       .value("dense", MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE)
       .value("compressed", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED)
       .value("singleton", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON);
 
   mlir_attribute_subclass(m, "EncodingAttr",
-                          mlirAttributeIsASparseTensorEncodingAttr,
-                          attributeClass)
+                          mlirAttributeIsASparseTensorEncodingAttr)
       .def_classmethod(
           "get",
           [](py::object cls,
@@ -72,3 +67,8 @@ void mlir::python::populateDialectSparseTensorSubmodule(
         return mlirSparseTensorEncodingAttrGetIndexBitWidth(self);
       });
 }
+
+PYBIND11_MODULE(_mlirDialectsSparseTensor, m) {
+  m.doc() = "MLIR SparseTensor dialect.";
+  populateDialectSparseTensorSubmodule(m);
+}

mlir/lib/Bindings/Python/Dialects.h

@@ -10,7 +10,6 @@
 
 #include "PybindUtils.h"
 
-#include "Dialects.h"
 #include "Globals.h"
 #include "IRModule.h"
 #include "Pass.h"
@@ -100,13 +99,4 @@ PYBIND11_MODULE(_mlir, m) {
   auto passModule =
       m.def_submodule("passmanager", "MLIR Pass Management Bindings");
   populatePassManagerSubmodule(passModule);
-
-  // Define and populate dialect submodules.
-  auto dialectsModule = m.def_submodule("dialects");
-  auto linalgModule = dialectsModule.def_submodule("linalg");
-  populateDialectLinalgSubmodule(linalgModule);
-  populateDialectSparseTensorSubmodule(
-      dialectsModule.def_submodule("sparse_tensor"), irModule);
-  populateDialectQuantSubmodule(dialectsModule.def_submodule("quant"),
-                                irModule);
 }

mlir/lib/Bindings/Python/MainModule.cpp

@@ -25,8 +25,6 @@ declare_mlir_python_sources(MLIRPythonSources.Core
     _mlir_libs/_mlir/__init__.pyi
     _mlir_libs/_mlir/ir.pyi
     _mlir_libs/_mlir/passmanager.pyi
-    # TODO: this should be split out into a separate library.
-    _mlir_libs/_mlir/dialects/quant.pyi
 )
 
 declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine
@@ -122,7 +120,8 @@ declare_mlir_python_sources(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   SOURCES
-    dialects/quant.py)
+    dialects/quant.py
+    _mlir_libs/_mlir/dialects/quant.pyi)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -191,9 +190,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
   ADD_TO_PARENT MLIRPythonSources.Core
   ROOT_DIR "${PYTHON_SOURCE_DIR}"
   SOURCES
-    DialectLinalg.cpp  # TODO: Break this out.
-    DialectSparseTensor.cpp  # TODO: Break this out.
-    DialectQuant.cpp # TODO: Break this out.
     MainModule.cpp
     IRAffine.cpp
     IRAttributes.cpp
@@ -205,7 +201,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
     Pass.cpp
 
     # Headers must be included explicitly so they are installed.
-    Dialects.h
     Globals.h
     IRModule.h
     Pass.h
@@ -219,10 +214,46 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
     MLIRCAPIRegistration  # TODO: See about dis-aggregating
 
     # Dialects
-    MLIRCAPILinalg  # TODO: Remove when above is removed.
-    MLIRCAPISparseTensor  # TODO: Remove when above is removed.
     MLIRCAPIStandard
-    MLIRCAPIQuant  # TODO: Remove when above is removed.
+)
+
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
+  MODULE_NAME _mlirDialectsLinalg
+  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
+  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  SOURCES
+    DialectLinalg.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIIR
+    MLIRCAPILinalg
+)
+
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
+  MODULE_NAME _mlirDialectsQuant
+  ADD_TO_PARENT MLIRPythonSources.Dialects.quant
+  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  SOURCES
+    DialectQuant.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIIR
+    MLIRCAPIQuant
+)
+
+declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
+  MODULE_NAME _mlirDialectsSparseTensor
+  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
+  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  SOURCES
+    DialectSparseTensor.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIIR
+    MLIRCAPISparseTensor
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.AllPassesRegistration

mlir/python/CMakeLists.txt

@@ -6,7 +6,7 @@
   from typing import Optional, Sequence, Union
   from ..ir import *
   from ._ods_common import get_default_loc_context
-  from .._mlir_libs._mlir.dialects.linalg import fill_builtin_region
+  from .._mlir_libs._mlirDialectsLinalg import fill_builtin_region
 except ImportError as e:
   raise RuntimeError("Error loading imports from extension module") from e
 

mlir/python/mlir/dialects/_linalg_ops_ext.py

@@ -2,6 +2,9 @@
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+# Re-export the objects provided by pybind.
+from ..._mlir_libs._mlirDialectsLinalg import *
+
 # These are the backing OpView classes generated from the linalg tablegen
 # definitions following these steps:
 #   DSL -> YAML -> tblgen -> pytblgen -> build/.../_linalg_ops_gen.py.
@@ -15,39 +18,39 @@
 #               C=TensorDef(U, S.M, S.N, output=True)):
 # ```
 # using the linalg-py eDSL.
-# The linalg-py eDSL builds a python representation (PyRepr) that is 
+# The linalg-py eDSL builds a python representation (PyRepr) that is
 # used in following ways:
 #  1. PyRepr -> YAML to generate the C++ and Python .td files. These
 #     then turn into the core C++ Op classes and Python OpView classes
-#     respectively (made available in _linalg_ops_gen). The generic OpView class 
+#     respectively (made available in _linalg_ops_gen). The generic OpView class
 #     mechanism makes the C++ classes available to python through the CAPI.
 #     PyRepr -> YAML currently occurs before compiler compile time.
 #     The other steps in this category occur at compiler compile time.
-#  2. PyRepr -> linalg.core_named_ops calls: piggybacks on the 
+#  2. PyRepr -> linalg.core_named_ops calls: piggybacks on the
 #     _linalg_ops_gen classes and the OpView mechanism to build IR at
 #     runtime in python:
 #       a. by default, the Named Op Form is emitted, e.g.:
 #          `linalg.matmul(lhs, rhs, outs=[out])` creates the following IR:
 #          ```
-#             %1 = linalg.matmul ins(%arg0, %arg1 : tensor<4x16xf32>, tensor<16x8xf32>) 
+#             %1 = linalg.matmul ins(%arg0, %arg1 : tensor<4x16xf32>, tensor<16x8xf32>)
 #                               outs(%0 : tensor<4x8xf32>)
-#                  -> tensor<4x8xf32>   
+#                  -> tensor<4x8xf32>
 #          ```
 #       b. by setting emit_generic=True, the Generic Op Form is emitted, e.g.:
 #           `linalg.matmul(lhs, rhs, outs=[out], emit_generic=True)` creates the following IR:
 #          ```
-#             %1 = linalg.generic {indexing_maps = [...], iterator_types = [...]} 
-#               ins(%arg0, %arg1 : tensor<4x16xf32>, tensor<16x8xf32>) 
+#             %1 = linalg.generic {indexing_maps = [...], iterator_types = [...]}
+#               ins(%arg0, %arg1 : tensor<4x16xf32>, tensor<16x8xf32>)
 #              outs(%0 : tensor<4x8xf32>) {
-#               ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):  
+#               ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
 #                  ...
 #                  linalg.yield %3 : f32
-#             } -> tensor<4x8xf32>  
+#             } -> tensor<4x8xf32>
 #          ```
 #  3. PyRepr -> Runtime Custom Op definitions: directly generates a
 #     linalg.generic form like in 2.b.
-#     !!!WARNING!!!: if one creates a runtime custom op with the same name 
+#     !!!WARNING!!!: if one creates a runtime custom op with the same name
 #     as an existing core named op, step 2. will likely take precedence.
-#     TODO: guard against surprises and fail create Runtime Custom Ops with 
+#     TODO: guard against surprises and fail create Runtime Custom Ops with
 #     the same name as existing Core Named Ops.
 from .opdsl.ops.core_named_ops import *

mlir/python/mlir/dialects/linalg/__init__.py

@@ -5,7 +5,6 @@
 from typing import Dict, List, Sequence, Tuple, Union
 
 from .....ir import *
-from ....._mlir_libs._mlir.dialects.linalg import fill_builtin_region
 
 from .... import linalg
 from .... import std
@@ -173,7 +172,7 @@ def emit_named_structured_op(op_config: LinalgStructuredOpConfig, op_name: str,
         f"Unknown named op_name / op_class_name: {op_name} / {op_class_name}")
 
   named_op = getattr(linalg, op_class_name)(ins, outs, result_types)
-  fill_builtin_region(named_op.operation)
+  linalg.fill_builtin_region(named_op.operation)
   # Note: mlir-linalg-ods-yaml-gen.cpp uses a special linalg.memoized_indexing_maps
   # attribute that the non-yaml path does not. The non-yaml path hardcodes the
   # indexing_maps in C++ directly.

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

@@ -2,4 +2,4 @@
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from .._mlir_libs._mlir.dialects.quant import *
+from .._mlir_libs._mlirDialectsQuant import *

mlir/python/mlir/dialects/quant.py

@@ -3,5 +3,5 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ._sparse_tensor_ops_gen import *
-from .._mlir_libs._mlir.dialects.sparse_tensor import *
+from .._mlir_libs._mlirDialectsSparseTensor import *
 from .._mlir_libs import _mlirSparseTensorPasses as _cextSparseTensorPasses