DRAFT: Adapt IREE to MLIR Python API registration overhaul.

compiler/src/iree/compiler/API/python/CMakeLists.txt

@@ -39,6 +39,7 @@ set(CMAKE_PLATFORM_NO_VERSIONED_SONAME 1)
 declare_mlir_python_sources(IREECompilerAPIPythonSources
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler"
   SOURCES
+    _mlir_libs/_site_initialize_0.py
     transforms/ireec.py
 )
 declare_mlir_python_sources(IREECompilerAPIPythonExtensions)

compiler/src/iree/compiler/API/python/iree/compiler/_mlir_libs/_site_initialize_0.py

@@ -0,0 +1,19 @@
+# Copyright 2022 The IREE Authors
+#
+# Licensed 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
+
+# Note that this does some load time, once-only init.
+from . import _ireecTransforms
+
+
+def register_dialects(registry):
+  # TODO: Convert the below context_init_hook to be DialectRegistry
+  # based and use it here.
+  pass
+
+
+def context_init_hook(context):
+  _ireecTransforms.register_all_dialects(context)
+

compiler/src/iree/compiler/API/python/test/transforms/ireec/compile_sample_module.py

@@ -29,8 +29,6 @@
 
 # Test the compiler API.
 with ir.Context() as ctx:
-  ireec.register_all_dialects(ctx)
-
   input_module = ir.Module.parse(r"""
     builtin.module  {
       func.func @fabs(%arg0: tensor<1x4xf32>, %arg1: tensor<4x1xf32>) -> tensor<4x4xf32> {

llvm-external-projects/iree-dialects/include/iree-dialects-c/Dialects.h

@@ -9,7 +9,6 @@
 
 #include "mlir-c/IR.h"
 #include "mlir-c/Pass.h"
-#include "mlir-c/Registration.h"
 
 #ifdef __cplusplus
 extern "C" {

llvm-external-projects/iree-dialects/python/CMakeLists.txt

@@ -85,15 +85,12 @@ declare_mlir_python_extension(IREEDialectsPythonExtensions.Main
 ################################################################################
 
 set(_source_components
-  # TODO: Core is now implicitly building/registering all dialects, increasing
-  # build burden by ~5x. Make it stop.
   MLIRPythonSources.Core
   MLIRPythonSources.Dialects.builtin
   MLIRPythonSources.Dialects.cf
   MLIRPythonSources.Dialects.func
   MLIRPythonSources.Dialects.pdl
   MLIRPythonSources.Dialects.transform
-  MLIRPythonSources.Passes
   IREEDialectsPythonSources
   IREEDialectsPythonExtensions
 )

llvm-external-projects/iree-dialects/python/IREEDialectsModule.cpp

@@ -10,7 +10,7 @@
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
 #include "mlir-c/Diagnostics.h"
-#include "mlir-c/Registration.h"
+#include "mlir-c/IR.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
 
 namespace py = pybind11;