[MLIR][Python] enable precise registration

Author: makslevental

mlir/cmake/modules/AddMLIRPython.cmake

@@ -457,11 +457,14 @@ endfunction()
 #     This file is where the *Attrs are defined, not where the *Enums are defined.
 #     **WARNING**: This arg will shortly be removed when the TODO for
 #     declare_mlir_dialect_python_bindings is satisfied. Use at your risk.
+#   EMBED_CAPI_LINK_LIBS: Dependent CAPI libraries that this extension depends
+#     on. These will be collected for all extensions and put into an
+#     aggregate dylib that is linked against.
 function(declare_mlir_dialect_extension_python_bindings)
   cmake_parse_arguments(ARG
     "GEN_ENUM_BINDINGS"
     "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME;EXTENSION_NAME"
-    "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE"
+    "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE;EMBED_CAPI_LINK_LIBS"
     ${ARGN})
   # Source files.
   set(_extension_target "${ARG_ADD_TO_PARENT}.${ARG_EXTENSION_NAME}")
@@ -503,6 +506,7 @@ function(declare_mlir_dialect_extension_python_bindings)
       ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
       ADD_TO_PARENT "${_extension_target}"
       SOURCES ${_sources}
+      EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
     )
   endif()
 endfunction()

mlir/include/mlir-c/Bindings/Python/Interop.h

@@ -84,6 +84,8 @@
 #define MLIR_PYTHON_CAPSULE_VALUE MAKE_MLIR_PYTHON_QUALNAME("ir.Value._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_TYPEID                                             \
   MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID._CAPIPtr")
+#define MLIR_PYTHON_CAPSULE_DIALECT_HANDLE                                     \
+  MAKE_MLIR_PYTHON_QUALNAME("ir.DialectHandle._CAPIPtr")
 
 /** Attribute on MLIR Python objects that expose their C-API pointer.
  * This will be a type-specific capsule created as per one of the helpers
@@ -457,6 +459,13 @@ static inline MlirValue mlirPythonCapsuleToValue(PyObject *capsule) {
   return value;
 }
 
+static inline MlirDialectHandle
+mlirPythonCapsuleToDialectHandle(PyObject *capsule) {
+  void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_DIALECT_HANDLE);
+  MlirDialectHandle handle = {ptr};
+  return handle;
+}
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -0,0 +1,36 @@
+//===-- mlir-c/Dialect/Affine.h - C API for Affine dialect --------*- C -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// Affine dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_AFFINE_H
+#define MLIR_C_DIALECT_AFFINE_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Affine, affine);
+
+MLIR_CAPI_EXPORTED void
+mlirAffineRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_AFFINE_H

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

@@ -0,0 +1,36 @@
+//===-- mlir-c/Dialect/Bufferization.h - C API for Bufferization dialect --===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// Bufferization dialect. A dialect should be registered with a context to make
+// it available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_BUFFERIZATION_H
+#define MLIR_C_DIALECT_BUFFERIZATION_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Bufferization, bufferization);
+
+MLIR_CAPI_EXPORTED void mlirBufferizationRegisterTransformDialectExtension(
+    MlirDialectRegistry registry);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_BUFFERIZATION_H

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

@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/Builtin.h - C API for Builtin dialect ------*- C -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// Builtin dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_BUILTIN_H
+#define MLIR_C_DIALECT_BUILTIN_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Builtin, builtin);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_BUILTIN_H

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

@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/Complex.h - C API for Complex dialect ------*- C -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// Complex dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_COMPLEX_H
+#define MLIR_C_DIALECT_COMPLEX_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Complex, complex);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_COMPLEX_H

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

@@ -63,6 +63,9 @@ mlirGPUObjectAttrHasKernels(MlirAttribute mlirObjectAttr);
 MLIR_CAPI_EXPORTED MlirAttribute
 mlirGPUObjectAttrGetKernels(MlirAttribute mlirObjectAttr);
 
+MLIR_CAPI_EXPORTED void
+mlirGPURegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -55,6 +55,9 @@ mlirLinalgGetIndexingMapsAttribute(MlirOperation op);
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Linalg, linalg);
 
+MLIR_CAPI_EXPORTED void
+mlirLinalgRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -26,6 +26,9 @@ extern "C" {
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(MemRef, memref);
 
+MLIR_CAPI_EXPORTED void
+mlirMemRefRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -29,6 +29,9 @@ MLIR_CAPI_EXPORTED MlirType mlirNVGPUTensorMapDescriptorTypeGet(
     MlirContext ctx, MlirType tensorMemrefType, int swizzle, int l2promo,
     int oobFill, int interleave);
 
+MLIR_CAPI_EXPORTED void
+mlirNVGPURegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif