[mlir] Use public PybindAdaptors in Linalg dialect bindings

[ftynse]

Summary

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.

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

 #include "mlir-c/Registration.h"
 #include "mlir-c/Support.h"
 
 #ifdef __cplusplus
 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);
 
 #ifdef __cplusplus
 }
 #endif
 
 #include "mlir/Dialect/Linalg/Passes.capi.h.inc"
 
 #endif // MLIR_C_DIALECT_LINALG_H

mlir/lib/Bindings/Python/DialectLinalg.cpp

 //===- DialectLinalg.cpp - Pybind module for Linalg dialect API support --===//
 //
 // 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 "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

 #include "mlir/CAPI/Registration.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 
 using namespace mlir;
 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()));
 
   ImplicitLocOpBuilder b(op->getLoc(), op->getContext());
   Region &region = op->getRegion(0);
   Block *body = b.createBlock(&region, /*insertPt=*/{}, argTypes);
   b.setInsertionPointToStart(body);
   fun(b, *body);
 }
 
 MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Linalg, linalg, LinalgDialect)

mlir/python/mlir/dialects/_linalg_ops_ext.py

if isa(RankedTensorType, output.type):
       results = [output.type]
     op = self.build_generic(
         results=results,
         operands=[_get_op_result_or_value(o) for o in [value, output]],
         attributes=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."""
 
   def __init__(self,
                sizes: Union[Sequence[int], Sequence[Value]],
                element_type: Type,
                *,

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

def emit_named_structured_op(op_config: LinalgStructuredOpConfig, op_name: str,
   ctx = Context.current
   fully_qualified_name = "linalg." + op_name
   if (not ctx.is_registered_operation(fully_qualified_name) or
       not op_class_name in linalg.__dict__.keys()):
     raise NotImplementedError(
         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.
   named_op.operation.attributes[
       "linalg.memoized_indexing_maps"] = indexing_maps_attr
   # iterator_types are hardcoded in C++ both in the yaml and non-yaml path.
 
   # Additionally set all named attributes.