[MLIR][Python] enable precise registration

Author: makslevental

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

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

mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp

@@ -32,7 +32,7 @@ NB_MODULE(_standaloneDialectsNanobind, m) {
         mlirDialectHandleRegisterDialect(standaloneHandle, context);
         if (load) {
           mlirDialectHandleLoadDialect(arithHandle, context);
-          mlirDialectHandleRegisterDialect(standaloneHandle, context);
+          mlirDialectHandleLoadDialect(standaloneHandle, context);
         }
       },
       nb::arg("context").none() = nb::none(), nb::arg("load") = true,

mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp

@@ -31,7 +31,7 @@ PYBIND11_MODULE(_standaloneDialectsPybind11, m) {
         mlirDialectHandleRegisterDialect(standaloneHandle, context);
         if (load) {
           mlirDialectHandleLoadDialect(arithHandle, context);
-          mlirDialectHandleRegisterDialect(standaloneHandle, context);
+          mlirDialectHandleLoadDialect(standaloneHandle, context);
         }
       },
       py::arg("context") = py::none(), py::arg("load") = true);

mlir/examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py

@@ -4,3 +4,10 @@
 
 from ._standalone_ops_gen import *
 from .._mlir_libs._standaloneDialectsNanobind.standalone import *
+
+from .._mlir_libs import get_dialect_registry as _get_dialect_registry
+from .._mlir_libs._capi import register_dialect as _register_dialect
+
+_dialect_registry = _get_dialect_registry()
+if "mlirGetDialectHandle__quant__" not in _dialect_registry.dialect_names:
+    _register_dialect("mlirGetDialectHandle__quant__", _dialect_registry)

mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py

@@ -4,3 +4,10 @@
 
 from ._standalone_ops_gen import *
 from .._mlir_libs._standaloneDialectsPybind11.standalone import *
+
+from .._mlir_libs import get_dialect_registry as _get_dialect_registry
+from .._mlir_libs._capi import register_dialect as _register_dialect
+
+_dialect_registry = _get_dialect_registry()
+if "mlirGetDialectHandle__quant__" not in _dialect_registry.dialect_names:
+    _register_dialect("mlirGetDialectHandle__quant__", _dialect_registry)

mlir/examples/standalone/test/python/smoketest.py

@@ -11,15 +11,26 @@
 else:
     raise ValueError("Expected either pybind11 or nanobind as arguments")
 
+from mlir_standalone.dialects import arith, quant
 
-with Context():
+
+with Context(), Location.unknown():
     standalone_d.register_dialects()
-    module = Module.parse(
-        """
-    %0 = arith.constant 2 : i32
-    %1 = standalone.foo %0 : i32
-    """
+    f32 = F32Type.get()
+    i8 = IntegerType.get_signless(8)
+    i32 = IntegerType.get_signless(32)
+    uniform = quant.UniformQuantizedType.get(
+        quant.UniformQuantizedType.FLAG_SIGNED, i8, f32, 0.99872, 127, -8, 7
     )
-    # CHECK: %[[C:.*]] = arith.constant 2 : i32
-    # CHECK: standalone.foo %[[C]] : i32
+
+    module = Module.create()
+    with InsertionPoint(module.body):
+        two_i32 = arith.constant(i32, 2)
+        standalone_d.foo(two_i32)
+        two_f32 = arith.constant(f32, 2.0)
+        quant.qcast(uniform, two_f32)
+    # CHECK: %[[TWOI32:.*]] = arith.constant 2 : i32
+    # CHECK: standalone.foo %[[TWOI32]] : i32
+    # CHECK: %[[TWOF32:.*]] = arith.constant 2.000000e+00 : f32
+    # CHECK: quant.qcast %[[TWOF32]] : f32 to !quant.uniform<i8<-8:7>:f32, 9.987200e-01:127>
     print(str(module))

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