[MLIR][Python] enable precise registration

mlir/cmake/modules/AddMLIRPython.cmake

@@ -23,11 +23,14 @@
 #     grouping. Source groupings form a DAG.
 #   SOURCES: List of specific source files relative to ROOT_DIR to include.
 #   SOURCES_GLOB: List of glob patterns relative to ROOT_DIR to include.
+#   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_python_sources name)
   cmake_parse_arguments(ARG
     ""
     "ROOT_DIR;ADD_TO_PARENT"
-    "SOURCES;SOURCES_GLOB"
+    "SOURCES;SOURCES_GLOB;EMBED_CAPI_LINK_LIBS"
     ${ARGN})
 
   if(NOT ARG_ROOT_DIR)
@@ -53,9 +56,10 @@ function(declare_mlir_python_sources name)
   set_target_properties(${name} PROPERTIES
     # Yes: Leading-lowercase property names are load bearing and the recommended
     # way to do this: https://gitlab.kitware.com/cmake/cmake/-/issues/19261
-    EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_DEPENDS"
+    EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_DEPENDS;mlir_python_EMBED_CAPI_LINK_LIBS"
     mlir_python_SOURCES_TYPE pure
     mlir_python_DEPENDS ""
+    mlir_python_EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
   )
 
   # Use the interface include directories and sources on the target to carry the
@@ -374,6 +378,9 @@ endfunction()
 #     This file is where the *EnumAttrs are defined, not where the *Enums are defined.
 #     **WARNING**: This arg will shortly be removed when the just-below TODO 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.
 #
 # TODO: Right now `TD_FILE` can't be the actual dialect tablegen file, since we
 #       use its path to determine where to place the generated python file. If
@@ -383,7 +390,7 @@ function(declare_mlir_dialect_python_bindings)
   cmake_parse_arguments(ARG
     "GEN_ENUM_BINDINGS"
     "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME"
-    "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE"
+    "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE;EMBED_CAPI_LINK_LIBS"
     ${ARGN})
   # Sources.
   set(_dialect_target "${ARG_ADD_TO_PARENT}.${ARG_DIALECT_NAME}")
@@ -424,6 +431,7 @@ function(declare_mlir_dialect_python_bindings)
       ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
       ADD_TO_PARENT "${_dialect_target}"
       SOURCES ${_sources}
+      EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
     )
   endif()
 endfunction()

mlir/examples/standalone/pyproject.toml

@@ -37,6 +37,8 @@ cmake.source-dir = "."
 # This is for installing/distributing the python bindings target and only the python bindings target.
 build.targets = ["StandalonePythonModules"]
 install.components = ["StandalonePythonModules"]
+# The default is true but make it explicit to highlight that this option exists (turn off for debug symbols).
+install.strip = true
 
 [tool.scikit-build.cmake.define]
 # Optional
@@ -51,6 +53,9 @@ LLVM_USE_LINKER = { env = "LLVM_USE_LINKER", default = "" }
 CMAKE_VISIBILITY_INLINES_HIDDEN = "ON"
 CMAKE_C_VISIBILITY_PRESET = "hidden"
 CMAKE_CXX_VISIBILITY_PRESET = "hidden"
+# Disables generation of "version soname" (i.e. libFoo.so.<version>),
+# which causes pure duplication of various shlibs for Python wheels.
+CMAKE_PLATFORM_NO_VERSIONED_SONAME = "ON"
 
 # Non-optional (alternatively you could use CMAKE_PREFIX_PATH here).
 MLIR_DIR = { env = "MLIR_DIR", default = "" }

mlir/examples/standalone/python/CMakeLists.txt

@@ -138,6 +138,7 @@ set(_declared_sources
   StandalonePythonSources
   MLIRPythonSources.Core
   MLIRPythonSources.Dialects.builtin
+  MLIRPythonSources.Dialects.arith
 )
 # For an external projects build, the MLIRPythonExtension.Core.type_stub_gen
 # target already exists and can just be added to DECLARED_SOURCES.

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

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

@@ -66,6 +66,9 @@ MLIR_CAPI_EXPORTED bool mlirTypeIsAPDLValueType(MlirType type);
 
 MLIR_CAPI_EXPORTED MlirType mlirPDLValueTypeGet(MlirContext ctx);
 
+MLIR_CAPI_EXPORTED void
+mlirPDLRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -104,6 +104,9 @@ mlirSMTAttrGetBVCmpPredicate(MlirContext ctx, MlirStringRef str);
 MLIR_CAPI_EXPORTED MlirAttribute mlirSMTAttrGetIntPredicate(MlirContext ctx,
                                                             MlirStringRef str);
 
+MLIR_CAPI_EXPORTED void
+mlirSMTRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -107,6 +107,9 @@ mlirSparseTensorEncodingAttrBuildLvlType(
     const enum MlirSparseTensorLevelPropertyNondefault *properties,
     unsigned propSize, unsigned n, unsigned m);
 
+MLIR_CAPI_EXPORTED void
+mlirSparseTensorRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -18,6 +18,9 @@ extern "C" {
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Tensor, tensor);
 
+MLIR_CAPI_EXPORTED void
+mlirTensorRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif

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

@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/Tosa.h - C API for Tosa 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
+// Tosa 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_TOSA_H
+#define MLIR_C_DIALECT_TOSA_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Tosa, tosa);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_TOSA_H

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

@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/UB.h - C API for UB 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
+// UB 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_UB_H
+#define MLIR_C_DIALECT_UB_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(UB, ub);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_UB_H