[Mlir-commits] [mlir] d716cfc - [mlir] Use public PybindAdaptors in Linalg dialect bindings

[Alex Zinenko]

Author: Alex Zinenko
Date: 2022-01-05T16:18:30+01:00
New Revision: d716cfc4fa3e2152df73dcb9eda693303c8a33a9

URL: https://github.com/llvm/llvm-project/commit/d716cfc4fa3e2152df73dcb9eda693303c8a33a9
DIFF: https://github.com/llvm/llvm-project/commit/d716cfc4fa3e2152df73dcb9eda693303c8a33a9.diff

LOG: [mlir] Use public PybindAdaptors in Linalg dialect bindings

Previously, the Python bindings for the Linalg dialect relied on the internal
implementation of core bindings. Most of that functionality was moved, and the
remaining one does not need access to the implementation: it used to accept a
dialect pointer as argument, but it can always be extracted from the operation
that it also accepts; operations are available through PybindAdaptors in an
opaque way. Change the bindings in that direction.

This enables the decoupling of the Linalg dialect Python extension from the
core IR Python extension.

Reviewed By: nicolasvasilache

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

Added: 
    

Modified: 
    mlir/include/mlir-c/Dialect/Linalg.h
    mlir/lib/Bindings/Python/DialectLinalg.cpp
    mlir/lib/CAPI/Dialect/Linalg.cpp
    mlir/python/mlir/dialects/_linalg_ops_ext.py
    mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir-c/Dialect/Linalg.h b/mlir/include/mlir-c/Dialect/Linalg.h
index 27f2f7bc897f7..2fe1872bed0cb 100644
--- a/mlir/include/mlir-c/Dialect/Linalg.h
+++ b/mlir/include/mlir-c/Dialect/Linalg.h
@@ -18,9 +18,9 @@ extern "C" {
 #endif
 
 /// Apply the special region builder for the builtin named Linalg op.
-/// Assert that `op` is a builtin named Linalg op.
+/// Assert that `mlirOp` is a builtin named Linalg op.
 MLIR_CAPI_EXPORTED void
-mlirLinalgFillBuiltinNamedOpRegion(MlirDialect linalgDialect, MlirOperation op);
+mlirLinalgFillBuiltinNamedOpRegion(MlirOperation mlirOp);
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Linalg, linalg);
 

diff  --git a/mlir/lib/Bindings/Python/DialectLinalg.cpp b/mlir/lib/Bindings/Python/DialectLinalg.cpp
index a2a54249e6d68..a168256159ed6 100644
--- a/mlir/lib/Bindings/Python/DialectLinalg.cpp
+++ b/mlir/lib/Bindings/Python/DialectLinalg.cpp
@@ -7,24 +7,17 @@
 //===----------------------------------------------------------------------===//
 
 #include "Dialects.h"
-#include "IRModule.h"
 #include "mlir-c/Dialect/Linalg.h"
 #include "mlir-c/IR.h"
-
-// TODO: Port this to operate only on the public PybindAdaptors.h
-#include "PybindUtils.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
 
 namespace py = pybind11;
-using namespace mlir;
-using namespace mlir::python;
 
 void mlir::python::populateDialectLinalgSubmodule(py::module m) {
   m.def(
       "fill_builtin_region",
-      [](PyDialectDescriptor &dialect, PyOperation &op) {
-        mlirLinalgFillBuiltinNamedOpRegion(dialect.get(), op.get());
-      },
-      py::arg("dialect"), py::arg("op"),
+      [](MlirOperation op) { mlirLinalgFillBuiltinNamedOpRegion(op); },
+      py::arg("op"),
       "Fill the region for `op`, which is assumed to be a builtin named Linalg "
       "op.");
 }

diff  --git a/mlir/lib/CAPI/Dialect/Linalg.cpp b/mlir/lib/CAPI/Dialect/Linalg.cpp
index dd796167b6335..6c5ba9a883811 100644
--- a/mlir/lib/CAPI/Dialect/Linalg.cpp
+++ b/mlir/lib/CAPI/Dialect/Linalg.cpp
@@ -15,20 +15,19 @@ using namespace mlir::linalg;
 
 /// Apply the special region builder for the builtin named Linalg op.
 /// Assert that `op` is a builtin named Linalg op.
-void mlirLinalgFillBuiltinNamedOpRegion(MlirDialect linalgDialect,
-                                        MlirOperation mlirOp) {
+void mlirLinalgFillBuiltinNamedOpRegion(MlirOperation mlirOp) {
   Operation *op = unwrap(mlirOp);
-
+  auto linalgOp = cast<LinalgOp>(op);
+  auto *dialect = static_cast<LinalgDialect *>(linalgOp->getDialect());
   LinalgDialect::RegionBuilderFunType fun =
-      static_cast<LinalgDialect *>(unwrap(linalgDialect))
-          ->getRegionBuilder(op->getName().getStringRef());
+      dialect->getRegionBuilder(op->getName().getStringRef());
+
   assert(fun && "Expected a builtin named Linalg op.");
   assert(op->getNumRegions() == 1 && "Expected Linalg op with 1 region");
   assert(op->getRegion(0).getBlocks().empty() &&
          "Expected Linalg op with 0 blocks");
 
   SmallVector<Type, 8> argTypes;
-  auto linalgOp = cast<LinalgOp>(op);
   for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands())
     argTypes.push_back(getElementTypeOrSelf(opOperand->get().getType()));
 

diff  --git a/mlir/python/mlir/dialects/_linalg_ops_ext.py b/mlir/python/mlir/dialects/_linalg_ops_ext.py
index d6c57547ee163..90f9227248994 100644
--- a/mlir/python/mlir/dialects/_linalg_ops_ext.py
+++ b/mlir/python/mlir/dialects/_linalg_ops_ext.py
@@ -34,8 +34,7 @@ def __init__(self, output: Value, value: Value, *, loc=None, ip=None):
         loc=loc,
         ip=ip)
     OpView.__init__(self, op)
-    linalgDialect = Context.current.get_dialect_descriptor("linalg")
-    fill_builtin_region(linalgDialect, self.operation)
+    fill_builtin_region(self.operation)
 
 class InitTensorOp:
   """Extends the linalg.init_tensor op."""

diff  --git a/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py b/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py
index 933c26ad9a7be..c3cfdfac95dcf 100644
--- a/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py
+++ b/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py
@@ -173,8 +173,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)
-  linalgDialect = ctx.get_dialect_descriptor("linalg")
-  fill_builtin_region(linalgDialect, named_op.operation)
+  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.