iree-org · ingomueller-net · Jan 17, 2023 · Jan 17, 2023 · Jan 17, 2023 · Jan 17, 2023
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/lib/Conversion/IteratorsToLLVM/IteratorsToLLVM.cpp b/experimental/iterators/lib/Conversion/IteratorsToLLVM/IteratorsToLLVM.cpp
@@ -19,7 +19,7 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Func/Transforms/FuncConversions.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
-#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"

diff --git a/experimental/iterators/lib/Conversion/PassDetail.h b/experimental/iterators/lib/Conversion/PassDetail.h
@@ -11,7 +11,7 @@
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
-#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {

diff --git a/experimental/iterators/lib/Conversion/StatesToLLVM/StatesToLLVM.cpp b/experimental/iterators/lib/Conversion/StatesToLLVM/StatesToLLVM.cpp
@@ -14,7 +14,7 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Func/Transforms/FuncConversions.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
-#include "mlir/Dialect/SCF/Transforms.h"
+#include "mlir/Dialect/SCF/Transforms/Transforms.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/Transforms/DialectConversion.h"
 

diff --git a/experimental/iterators/lib/Utils/MLIRSupport.cpp b/experimental/iterators/lib/Utils/MLIRSupport.cpp
@@ -9,7 +9,7 @@
 #include "iterators/Utils/MLIRSupport.h"
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
-#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 
 using namespace mlir;

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
 
 

diff --git a/third_party/iree b/third_party/iree