Update CMake dependencies of the MLIR Python extension.

This follows the changes of 5e83a5b of the LLVM mono repo.
iree-org · Jan 17, 2023 · d39fafa · d39fafa
1 parent 44f1182
commit d39fafa
Show file tree

Hide file tree

Showing 7 changed files with 23 additions and 15 deletions.
diff --git a/experimental/iterators/benchmarks/inner_product/run.py b/experimental/iterators/benchmarks/inner_product/run.py
@@ -12,7 +12,6 @@
 import pandas as pd
 
 from iree.compiler.runtime.np_to_memref import get_ranked_memref_descriptor
-from mlir.sandbox.utils import realign
 from mlir_iterators.dialects import iterators as it
 from mlir_iterators.dialects import tabular as tab
 from mlir_iterators.dialects import arith, func, memref, scf
@@ -31,7 +30,6 @@
 )
 from mlir_iterators.passmanager import PassManager
 from mlir_iterators.runtime.pandas_to_iterators import to_tabular_view_descriptor
-import mlir_iterators.all_passes_registration
 
 _MLIR_RUNNER_UTILS_LIB_ENV = "MLIR_RUNNER_UTILS_LIB"
 _MLIR_RUNNER_UTILS_LIB_DEFAULT = "libmlir_runner_utils.so"
@@ -82,6 +80,23 @@ def emit_benchmarking_function(name: str, bench: func.FuncOp) -> func.FuncOp:
   return wrapper
 
 
+# Copied from mlir.sandbox.utils. That package uses the vanilla `mlir` package
+# instead of `mlir_iterators` as the rest of this file, so they are
+# incompatible.
+def realign(allocated_unaligned: np.ndarray, byte_alignment: int = 64):
+  shape = allocated_unaligned.shape
+  dt = allocated_unaligned.dtype
+  effective_size_in_bytes = np.prod(shape) * np.dtype(dt).itemsize
+  total_size_in_bytes = effective_size_in_bytes + byte_alignment
+  buf = np.empty(total_size_in_bytes, dtype=np.byte)
+  off = (-buf.ctypes.data % byte_alignment)
+  allocated_aligned = buf[off:off +
+                          effective_size_in_bytes].view(dt).reshape(shape)
+  np.copyto(allocated_aligned, allocated_unaligned)
+  assert allocated_aligned.ctypes.data % byte_alignment == 0
+  return allocated_aligned
+
+
 def setup_data(num_elements, dtype):
   """Sets up the input data: two `numpy.arrays` with `num_elements` each of
   type `dtype`."""

diff --git a/experimental/iterators/include/iterators-c/Dialects.h b/experimental/iterators/include/iterators-c/Dialects.h
@@ -10,7 +10,6 @@
 #define ITERATORS_C_DIALECTS_H
 
 #include "mlir-c/IR.h"
-#include "mlir-c/Registration.h"
 #include "mlir-c/Support.h"
 
 #ifdef __cplusplus

diff --git a/experimental/iterators/python/CMakeLists.txt b/experimental/iterators/python/CMakeLists.txt
@@ -64,11 +64,10 @@ add_mlir_python_common_capi_library(IteratorsPythonCAPI
   RELATIVE_INSTALL_ROOT "../../../.."
   DECLARED_SOURCES
   IteratorsPythonSources
-  MLIRPythonExtension.AllPassesRegistration
+  MLIRPythonExtension.RegisterEverything
   MLIRPythonSources.Core
   MLIRPythonSources.Dialects
   MLIRPythonSources.ExecutionEngine
-  MLIRPythonSources.Passes
 )
 
 # ###############################################################################
@@ -79,11 +78,10 @@ add_mlir_python_modules(IteratorsPythonModules
   INSTALL_PREFIX "python_packages/mlir_iterators"
   DECLARED_SOURCES
   IteratorsPythonSources
-  MLIRPythonExtension.AllPassesRegistration
+  MLIRPythonExtension.RegisterEverything
   MLIRPythonSources.Core
   MLIRPythonSources.Dialects
   MLIRPythonSources.ExecutionEngine
-  MLIRPythonSources.Passes
   COMMON_CAPI_LINK_LIBS
   IteratorsPythonCAPI
 )
diff --git a/experimental/iterators/python/IteratorsDialects.cpp b/experimental/iterators/python/IteratorsDialects.cpp
@@ -13,7 +13,6 @@
 
 #include "iterators-c/Dialects.h"
 #include "mlir-c/IR.h"
-#include "mlir-c/Registration.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
 
 namespace py = pybind11;

diff --git a/experimental/iterators/test/python/dialects/iterators/dialect.py b/experimental/iterators/test/python/dialects/iterators/dialect.py
@@ -11,7 +11,6 @@
 from mlir_iterators.passmanager import PassManager
 from mlir_iterators.execution_engine import ExecutionEngine
 from mlir_iterators.ir import Context, Module, IntegerType
-import mlir_iterators.all_passes_registration
 
 
 def run(f):

diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
@@ -54,7 +54,7 @@ add_mlir_python_common_capi_library(SandboxSandboxCAPI
   RELATIVE_INSTALL_ROOT "../../../.."
   DECLARED_SOURCES
     MLIRPythonSources
-    MLIRPythonExtension.AllPassesRegistration
+    MLIRPythonExtension.RegisterEverything
     SandboxSources
 )
 
@@ -69,7 +69,7 @@ add_mlir_python_modules(SandboxSandboxPythonModules
   INSTALL_PREFIX "python_packages/iree_llvm_sandbox/mlir"
   DECLARED_SOURCES
     MLIRPythonSources
-    MLIRPythonExtension.AllPassesRegistration
+    MLIRPythonExtension.RegisterEverything
     SandboxSources
   COMMON_CAPI_LINK_LIBS
     SandboxSandboxCAPI
@@ -103,7 +103,7 @@ add_mlir_python_common_capi_library(IREECAPI
   RELATIVE_INSTALL_ROOT "../../../../.."
   DECLARED_SOURCES
     MLIRPythonSources
-    MLIRPythonExtension.AllPassesRegistration
+    MLIRPythonExtension.RegisterEverything
     ${_addl_sources}
 )
 
@@ -118,7 +118,7 @@ add_mlir_python_modules(IREEPythonModules
   INSTALL_PREFIX "python_packages/iree_llvm_sandboxiree/compiler"
   DECLARED_SOURCES
     MLIRPythonSources
-    MLIRPythonExtension.AllPassesRegistration
+    MLIRPythonExtension.RegisterEverything
     ${_addl_sources}
   COMMON_CAPI_LINK_LIBS
     IREECAPI

diff --git a/python/mlir/sandbox/transforms.py b/python/mlir/sandbox/transforms.py
@@ -6,8 +6,6 @@
 from mlir.sandbox.variables import *
 from mlir.sandbox.transform import Transform
 
-# import mlir.all_passes_registration
-
 import typing as tp