[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

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);
 

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.");
 }

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()));
 

mlir/python/mlir/dialects/_linalg_ops_ext.py

@@ -34,8 +34,7 @@ class FillOp:
         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."""

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.