[Mlir-commits] [mlir] e31c77b - [mlir][python] Reorganize MLIR python into namespace packages.

Stella Laurenzo llvmlistbot at llvm.org
Mon Mar 8 23:05:56 PST 2021 

Author: Stella Laurenzo
Date: 2021-03-08T23:01:34-08:00
New Revision: e31c77b1827fa4dd3511f21af11cfab18ecf6d38

URL: https://github.com/llvm/llvm-project/commit/e31c77b1827fa4dd3511f21af11cfab18ecf6d38
DIFF: https://github.com/llvm/llvm-project/commit/e31c77b1827fa4dd3511f21af11cfab18ecf6d38.diff

LOG: [mlir][python] Reorganize MLIR python into namespace packages.

* Only leaf packages are non-namespace packages. This allows most of the top levels to be split into different directories or deployment packages. In the previous state, the presence of __init__.py files at each level meant that the entire tree could only ever exist in one physical directory on the path.
* This changes the API usage slightly: `import mlir` will no longer do a deep import of `mlir.ir`, etc. This may necessitate some client code changes.
* Dialect gen code was restructured so that the user is responsible for providing the `my_dialect.py` file, which then must import its peer `_my_dialect_ops_gen`. This gives complete control of the dialect namespace to the user instead of to tablegen code, allowing further dialect-specific python APIs.
* Correspondingly, the previous extension modules `_my_dialect.py` are now `_my_dialect_ops_ext.py`.
* Now that the `linalg` namespace is open, moved the `linalg_opdsl` tool into it.
* This may require some corresponding downstream adjustments to npcomp, circt, et al:
  * Probably some shallow imports need to be converted to deep imports (i.e. not `import mlir` brings in the world).
  * Each tablegen generated dialect now needs an explicit `foo.py` which does a `from ._foo_ops_gen import *`. This is similar to the way that generated code operates in the C++ world.
  * If providing dialect op extensions, those need to be moved from `_foo.py` -> `_foo_ops_ext.py`.

Differential Revision: https://reviews.llvm.org/D98096

Added: 
    mlir/lib/Bindings/Python/mlir/_cext_loader.py
    mlir/lib/Bindings/Python/mlir/dialects/_builtin_ops_ext.py
    mlir/lib/Bindings/Python/mlir/dialects/_linalg_ops_ext.py
    mlir/lib/Bindings/Python/mlir/dialects/_ods_common.py
    mlir/lib/Bindings/Python/mlir/dialects/builtin.py
    mlir/lib/Bindings/Python/mlir/dialects/linalg/__init__.py
    mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/__init__.py
    mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/dump_oplib.py
    mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/__init__.py
    mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/affine.py
    mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/comprehension.py
    mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/config.py
    mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/dsl.py
    mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/scalar_expr.py
    mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/types.py
    mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/yaml_helper.py
    mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/ops/__init__.py
    mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
    mlir/lib/Bindings/Python/mlir/dialects/python_test.py
    mlir/lib/Bindings/Python/mlir/dialects/shape.py
    mlir/lib/Bindings/Python/mlir/dialects/std.py
    mlir/lib/Bindings/Python/mlir/dialects/tensor.py
    mlir/test/Bindings/Python/dialects/linalg/opdsl/assignments.py
    mlir/test/Bindings/Python/dialects/linalg/opdsl/doctests.py
    mlir/test/Bindings/Python/dialects/linalg/opdsl/interfaces.py
    mlir/test/Bindings/Python/dialects/linalg/opdsl/lit.local.cfg
    mlir/test/Bindings/Python/dialects/linalg/opdsl/shape_maps_iteration.py
    mlir/test/Bindings/Python/dialects/linalg/opdsl/test_core_named_ops.py
    mlir/test/Bindings/Python/dialects/linalg/ops.py

Modified: 
    mlir/cmake/modules/AddMLIRPythonExtension.cmake
    mlir/docs/Bindings/Python.md
    mlir/lib/Bindings/Python/mlir/conversions/__init__.py
    mlir/lib/Bindings/Python/mlir/execution_engine.py
    mlir/lib/Bindings/Python/mlir/ir.py
    mlir/lib/Bindings/Python/mlir/passmanager.py
    mlir/lib/Bindings/Python/mlir/transforms/__init__.py
    mlir/test/Bindings/Python/context_lifecycle.py
    mlir/test/Bindings/Python/dialects.py
    mlir/test/Bindings/Python/ir_operation.py
    mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

Removed: 
    mlir/lib/Bindings/Python/mlir/__init__.py
    mlir/lib/Bindings/Python/mlir/dialects/__init__.py
    mlir/lib/Bindings/Python/mlir/dialects/_builtin.py
    mlir/lib/Bindings/Python/mlir/dialects/_linalg.py
    mlir/lib/Bindings/Python/mlir/tools/__init__.py
    mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/__init__.py
    mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/dump_oplib.py
    mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/__init__.py
    mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/affine.py
    mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/comprehension.py
    mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/config.py
    mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/dsl.py
    mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/scalar_expr.py
    mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/types.py
    mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/yaml_helper.py
    mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/ops/__init__.py
    mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/ops/core_named_ops.py
    mlir/test/Bindings/Python/dialects/linalg.py
    mlir/test/Bindings/Python/tools/linalg_opdsl/assignments.py
    mlir/test/Bindings/Python/tools/linalg_opdsl/doctests.py
    mlir/test/Bindings/Python/tools/linalg_opdsl/interfaces.py
    mlir/test/Bindings/Python/tools/linalg_opdsl/lit.local.cfg
    mlir/test/Bindings/Python/tools/linalg_opdsl/shape_maps_iteration.py
    mlir/test/Bindings/Python/tools/linalg_opdsl/test_core_named_ops.py


################################################################################
diff  --git a/mlir/cmake/modules/AddMLIRPythonExtension.cmake b/mlir/cmake/modules/AddMLIRPythonExtension.cmake
index dbbe71e22b13..dce7d836f8d6 100644
--- a/mlir/cmake/modules/AddMLIRPythonExtension.cmake
+++ b/mlir/cmake/modules/AddMLIRPythonExtension.cmake
@@ -143,8 +143,9 @@ function(add_mlir_dialect_python_bindings tblgen_target)
     "DEPENDS"
     ${ARGN})
 
+  set(dialect_filename "_${ARG_DIALECT_NAME}_ops_gen.py")
   set(LLVM_TARGET_DEFINITIONS ${ARG_TD_FILE})
-  mlir_tablegen("${ARG_DIALECT_NAME}.py" -gen-python-op-bindings
+  mlir_tablegen("${dialect_filename}" -gen-python-op-bindings
                 -bind-dialect=${ARG_DIALECT_NAME})
   add_public_tablegen_target(
     ${tblgen_target})
@@ -154,9 +155,9 @@ function(add_mlir_dialect_python_bindings tblgen_target)
 
   add_custom_command(
     TARGET ${tblgen_target} POST_BUILD
-    COMMENT "Copying generated python source \"dialects/${ARG_DIALECT_NAME}.py\""
+    COMMENT "Copying generated python source \"dialects/${dialect_filename}\""
     COMMAND "${CMAKE_COMMAND}" -E copy_if_
diff erent
-      "${CMAKE_CURRENT_BINARY_DIR}/${ARG_DIALECT_NAME}.py"
-      "${PROJECT_BINARY_DIR}/python/mlir/dialects/${ARG_DIALECT_NAME}.py")
+      "${CMAKE_CURRENT_BINARY_DIR}/${dialect_filename}"
+      "${PROJECT_BINARY_DIR}/python/mlir/dialects/${dialect_filename}")
 endfunction()
 

diff  --git a/mlir/docs/Bindings/Python.md b/mlir/docs/Bindings/Python.md
index 7ea5d022b34f..e5e8e6d77c30 100644
--- a/mlir/docs/Bindings/Python.md
+++ b/mlir/docs/Bindings/Python.md
@@ -122,18 +122,21 @@ pass registration, etc.
 LLVM/MLIR is a non-trivial python-native project that is likely to co-exist with
 other non-trivial native extensions. As such, the native extension (i.e. the
 `.so`/`.pyd`/`.dylib`) is exported as a notionally private top-level symbol
-(`_mlir`), while a small set of Python code is provided in `mlir/__init__.py`
-and siblings which loads and re-exports it. This split provides a place to stage
-code that needs to prepare the environment *before* the shared library is loaded
-into the Python runtime, and also provides a place that one-time initialization
-code can be invoked apart from module constructors.
-
-To start with the `mlir/__init__.py` loader shim can be very simple and scale to
-future need:
-
-```python
-from _mlir import *
-```
+(`_mlir`), while a small set of Python code is provided in
+`mlir/_cext_loader.py` and siblings which loads and re-exports it. This
+split provides a place to stage code that needs to prepare the environment
+*before* the shared library is loaded into the Python runtime, and also
+provides a place that one-time initialization code can be invoked apart from
+module constructors.
+
+It is recommended to avoid using `__init__.py` files to the extent possible,
+until reaching a leaf package that represents a discrete component. The rule
+to keep in mind is that the presence of an `__init__.py` file prevents the
+ability to split anything at that level or below in the namespace into
+
diff erent directories, deployment packages, wheels, etc.
+
+See the documentation for more information and advice:
+https://packaging.python.org/guides/packaging-namespace-packages/
 
 ### Use the C-API
 
@@ -361,13 +364,16 @@ are multiple parts to this integration, outlined below. Most details have
 been elided: refer to the build rules and python sources under `mlir.dialects`
 for the canonical way to use this facility.
 
-### Generating `{DIALECT_NAMESPACE}.py` wrapper modules
+Users are responsible for providing a `{DIALECT_NAMESPACE}.py` (or an
+equivalent directory with `__init__.py` file) as the entrypoint.
+
+### Generating `_{DIALECT_NAMESPACE}_ops_gen.py` wrapper modules
 
 Each dialect with a mapping to python requires that an appropriate
-`{DIALECT_NAMESPACE}.py` wrapper module is created. This is done by invoking
-`mlir-tblgen` on a python-bindings specific tablegen wrapper that includes
-the boilerplate and actual dialect specific `td` file. An example, for the
-`StandardOps` (which is assigned the namespace `std` as a special case):
+`_{DIALECT_NAMESPACE}_ops_gen.py` wrapper module is created. This is done by
+invoking `mlir-tblgen` on a python-bindings specific tablegen wrapper that
+includes the boilerplate and actual dialect specific `td` file. An example, for
+the `StandardOps` (which is assigned the namespace `std` as a special case):
 
 ```tablegen
 #ifndef PYTHON_BINDINGS_STANDARD_OPS
@@ -387,6 +393,13 @@ mlir-tblgen -gen-python-op-bindings -bind-dialect={DIALECT_NAMESPACE} \
     {PYTHON_BINDING_TD_FILE}
 ```
 
+The generates op classes must be included in the `{DIALECT_NAMESPACE}.py` file
+in a similar way that generated headers are included for C++ generated code:
+
+```python
+from ._my_dialect_ops_gen import *
+```
+
 ### Extending the search path for wrapper modules
 
 When the python bindings need to locate a wrapper module, they consult the

diff  --git a/mlir/lib/Bindings/Python/mlir/__init__.py b/mlir/lib/Bindings/Python/mlir/_cext_loader.py
similarity index 78%
rename from mlir/lib/Bindings/Python/mlir/__init__.py
rename to mlir/lib/Bindings/Python/mlir/_cext_loader.py
index 3eb4eb797b4e..35847efa9e4a 100644
--- a/mlir/lib/Bindings/Python/mlir/__init__.py
+++ b/mlir/lib/Bindings/Python/mlir/_cext_loader.py
@@ -1,18 +1,7 @@
 #  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
-
-# Note that the only function of this module is currently to load the
-# native module and re-export its symbols. In the future, this file is
-# reserved as a trampoline to handle environment specific loading needs
-# and arbitrate any one-time initialization needed in various shared-library
-# scenarios.
-
-__all__ = [
-  "ir",
-  "execution_engine",
-  "passmanager",
-]
+"""Common module for looking up and manipulating C-Extensions."""
 
 # Packaged installs have a top-level _mlir_libs package with symbols:
 #   load_extension(name): Loads a named extension module
@@ -36,19 +25,21 @@ def _load_extension(name):
   _preload_dependency = _mlir_libs.preload_dependency
 
 _preload_dependency("MLIRPublicAPI")
+
 # Expose the corresponding C-Extension module with a well-known name at this
 # top-level module. This allows relative imports like the following to
 # function:
-#   from .. import _cext
+#   from .._cext_loader import _cext
 # This reduces coupling, allowing embedding of the python sources into another
 # project that can just vary based on this top-level loader module.
 _cext = _load_extension("_mlir")
 
+
 def _reexport_cext(cext_module_name, target_module_name):
   """Re-exports a named sub-module of the C-Extension into another module.
 
   Typically:
-    from . import _reexport_cext
+    from ._cext_loader import _reexport_cext
     _reexport_cext("ir", __name__)
     del _reexport_cext
   """
@@ -60,9 +51,5 @@ def _reexport_cext(cext_module_name, target_module_name):
       setattr(target_module, attr_name, getattr(source_module, attr_name))
 
 
-# Import sub-modules. Since these may import from here, this must come after
-# any exported definitions.
-from . import ir, execution_engine, passmanager
-
 # Add our 'dialects' parent module to the search path for implementations.
 _cext.globals.append_dialect_search_prefix("mlir.dialects")

diff  --git a/mlir/lib/Bindings/Python/mlir/conversions/__init__.py b/mlir/lib/Bindings/Python/mlir/conversions/__init__.py
index 21713432620d..0989449a447b 100644
--- a/mlir/lib/Bindings/Python/mlir/conversions/__init__.py
+++ b/mlir/lib/Bindings/Python/mlir/conversions/__init__.py
@@ -4,5 +4,5 @@
 
 # Expose the corresponding C-Extension module with a well-known name at this
 # level.
-from .. import _load_extension
+from .._cext_loader import _load_extension
 _cextConversions = _load_extension("_mlirConversions")

diff  --git a/mlir/lib/Bindings/Python/mlir/dialects/_builtin.py b/mlir/lib/Bindings/Python/mlir/dialects/_builtin_ops_ext.py
similarity index 93%
rename from mlir/lib/Bindings/Python/mlir/dialects/_builtin.py
rename to mlir/lib/Bindings/Python/mlir/dialects/_builtin_ops_ext.py
index 1cb84017215a..b0789299139d 100644
--- a/mlir/lib/Bindings/Python/mlir/dialects/_builtin.py
+++ b/mlir/lib/Bindings/Python/mlir/dialects/_builtin_ops_ext.py
@@ -1,15 +1,15 @@
 #  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
-from mlir.ir import *
+from ..ir import *
 
 
 class ModuleOp:
   """Specialization for the module op class."""
 
   def __init__(self, *, loc=None, ip=None):
-    super().__init__(
-        self.build_generic(results=[], operands=[], loc=loc, ip=ip))
+    super().__init__(self.build_generic(results=[], operands=[], loc=loc,
+                                        ip=ip))
     body = self.regions[0].blocks.append()
     with InsertionPoint(body):
       Operation.create("module_terminator")
@@ -84,10 +84,11 @@ def entry_block(self):
     return self.regions[0].blocks[0]
 
   def add_entry_block(self):
-    '''
-    Add an entry block to the function body using the function signature to infer block arguments
+    """
+    Add an entry block to the function body using the function signature to
+    infer block arguments.
     Returns the newly created block
-    '''
+    """
     if not self.is_external:
       raise IndexError('The function already has an entry block!')
     self.body.blocks.append(*self.type.inputs)

diff  --git a/mlir/lib/Bindings/Python/mlir/dialects/_linalg.py b/mlir/lib/Bindings/Python/mlir/dialects/_linalg_ops_ext.py
similarity index 100%
rename from mlir/lib/Bindings/Python/mlir/dialects/_linalg.py
rename to mlir/lib/Bindings/Python/mlir/dialects/_linalg_ops_ext.py

diff  --git a/mlir/lib/Bindings/Python/mlir/dialects/__init__.py b/mlir/lib/Bindings/Python/mlir/dialects/_ods_common.py
similarity index 99%
rename from mlir/lib/Bindings/Python/mlir/dialects/__init__.py
rename to mlir/lib/Bindings/Python/mlir/dialects/_ods_common.py
index f5a71bf88700..6d37700ecdc4 100644
--- a/mlir/lib/Bindings/Python/mlir/dialects/__init__.py
+++ b/mlir/lib/Bindings/Python/mlir/dialects/_ods_common.py
@@ -3,7 +3,7 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 # Re-export the parent _cext so that every level of the API can get it locally.
-from .. import _cext
+from .._cext_loader import _cext
 
 __all__ = [
     "equally_sized_accessor",

diff  --git a/mlir/lib/Bindings/Python/mlir/dialects/builtin.py b/mlir/lib/Bindings/Python/mlir/dialects/builtin.py
new file mode 100644
index 000000000000..30279e1611f9
--- /dev/null
+++ b/mlir/lib/Bindings/Python/mlir/dialects/builtin.py
@@ -0,0 +1,5 @@
+#  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
+
+from ._builtin_ops_gen import *

diff  --git a/mlir/lib/Bindings/Python/mlir/dialects/linalg/__init__.py b/mlir/lib/Bindings/Python/mlir/dialects/linalg/__init__.py
new file mode 100644
index 000000000000..81949b8f881c
--- /dev/null
+++ b/mlir/lib/Bindings/Python/mlir/dialects/linalg/__init__.py
@@ -0,0 +1,5 @@
+#  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
+
+from .._linalg_ops_gen import *

diff  --git a/mlir/lib/Bindings/Python/mlir/tools/__init__.py b/mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/__init__.py
similarity index 100%
rename from mlir/lib/Bindings/Python/mlir/tools/__init__.py
rename to mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/__init__.py

diff  --git a/mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/dump_oplib.py b/mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/dump_oplib.py
similarity index 95%
rename from mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/dump_oplib.py
rename to mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/dump_oplib.py
index 5e8c3bf6a330..98bf2e247ea1 100644
--- a/mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/dump_oplib.py
+++ b/mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/dump_oplib.py
@@ -64,7 +64,8 @@ def main(args):
   modules = []
   for module_name in args.modules:
     modules.append(
-        importlib.import_module(module_name, package="mlir.tools.linalg_opdsl"))
+        importlib.import_module(module_name,
+                                package="mlir.dialects.linalg.opdsl"))
   for i, file_path in enumerate(args.file or []):
     modules.append(load_module_from_file(f"_mlir_eval_oplib{i}", file_path))
   for m in modules:

diff  --git a/mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/__init__.py b/mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/__init__.py
similarity index 100%
rename from mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/__init__.py
rename to mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/__init__.py

diff  --git a/mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/affine.py b/mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/affine.py
similarity index 100%
rename from mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/affine.py
rename to mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/affine.py

diff  --git a/mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/comprehension.py b/mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/comprehension.py
similarity index 100%
rename from mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/comprehension.py
rename to mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/comprehension.py

diff  --git a/mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/config.py b/mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/config.py
similarity index 100%
rename from mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/config.py
rename to mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/config.py

diff  --git a/mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/dsl.py b/mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/dsl.py
similarity index 100%
rename from mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/dsl.py
rename to mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/dsl.py

diff  --git a/mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/scalar_expr.py b/mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/scalar_expr.py
similarity index 100%
rename from mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/scalar_expr.py
rename to mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/scalar_expr.py

diff  --git a/mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/types.py b/mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/types.py
similarity index 100%
rename from mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/types.py
rename to mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/types.py

diff  --git a/mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/yaml_helper.py b/mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/yaml_helper.py
similarity index 100%
rename from mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/yaml_helper.py
rename to mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/lang/yaml_helper.py

diff  --git a/mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/__init__.py b/mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/ops/__init__.py
similarity index 100%
rename from mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/__init__.py
rename to mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/ops/__init__.py

diff  --git a/mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/ops/core_named_ops.py b/mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
similarity index 100%
rename from mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/ops/core_named_ops.py
rename to mlir/lib/Bindings/Python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py

diff  --git a/mlir/lib/Bindings/Python/mlir/dialects/python_test.py b/mlir/lib/Bindings/Python/mlir/dialects/python_test.py
new file mode 100644
index 000000000000..524db4317678
--- /dev/null
+++ b/mlir/lib/Bindings/Python/mlir/dialects/python_test.py
@@ -0,0 +1,5 @@
+#  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
+
+from ._python_test_ops_gen import *

diff  --git a/mlir/lib/Bindings/Python/mlir/dialects/shape.py b/mlir/lib/Bindings/Python/mlir/dialects/shape.py
new file mode 100644
index 000000000000..cc987ac843e7
--- /dev/null
+++ b/mlir/lib/Bindings/Python/mlir/dialects/shape.py
@@ -0,0 +1,5 @@
+#  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
+
+from ._shape_ops_gen import *

diff  --git a/mlir/lib/Bindings/Python/mlir/dialects/std.py b/mlir/lib/Bindings/Python/mlir/dialects/std.py
new file mode 100644
index 000000000000..8e55807a0420
--- /dev/null
+++ b/mlir/lib/Bindings/Python/mlir/dialects/std.py
@@ -0,0 +1,5 @@
+#  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
+
+from ._std_ops_gen import *

diff  --git a/mlir/lib/Bindings/Python/mlir/dialects/tensor.py b/mlir/lib/Bindings/Python/mlir/dialects/tensor.py
new file mode 100644
index 000000000000..26edf6b6436d
--- /dev/null
+++ b/mlir/lib/Bindings/Python/mlir/dialects/tensor.py
@@ -0,0 +1,5 @@
+#  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
+
+from ._tensor_ops_gen import *

diff  --git a/mlir/lib/Bindings/Python/mlir/execution_engine.py b/mlir/lib/Bindings/Python/mlir/execution_engine.py
index 15a874a710aa..89bd4aad5658 100644
--- a/mlir/lib/Bindings/Python/mlir/execution_engine.py
+++ b/mlir/lib/Bindings/Python/mlir/execution_engine.py
@@ -3,7 +3,7 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 # Simply a wrapper around the extension module of the same name.
-from . import _cext
+from ._cext_loader import _cext
 import ctypes
 
 class ExecutionEngine(_cext.execution_engine.ExecutionEngine):

diff  --git a/mlir/lib/Bindings/Python/mlir/ir.py b/mlir/lib/Bindings/Python/mlir/ir.py
index 70d19737f5e6..e5ba1bdb0270 100644
--- a/mlir/lib/Bindings/Python/mlir/ir.py
+++ b/mlir/lib/Bindings/Python/mlir/ir.py
@@ -3,6 +3,6 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 # Simply a wrapper around the extension module of the same name.
-from . import _reexport_cext
+from ._cext_loader import _reexport_cext
 _reexport_cext("ir", __name__)
 del _reexport_cext

diff  --git a/mlir/lib/Bindings/Python/mlir/passmanager.py b/mlir/lib/Bindings/Python/mlir/passmanager.py
index 95119e52f971..6b267b76eb7d 100644
--- a/mlir/lib/Bindings/Python/mlir/passmanager.py
+++ b/mlir/lib/Bindings/Python/mlir/passmanager.py
@@ -3,6 +3,6 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 # Simply a wrapper around the extension module of the same name.
-from . import _reexport_cext
+from ._cext_loader import _reexport_cext
 _reexport_cext("passmanager", __name__)
 del _reexport_cext

diff  --git a/mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/ops/__init__.py b/mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/ops/__init__.py
deleted file mode 100644
index e69de29bb2d1..000000000000

diff  --git a/mlir/lib/Bindings/Python/mlir/transforms/__init__.py b/mlir/lib/Bindings/Python/mlir/transforms/__init__.py
index 1155c6392eb8..2149933d0848 100644
--- a/mlir/lib/Bindings/Python/mlir/transforms/__init__.py
+++ b/mlir/lib/Bindings/Python/mlir/transforms/__init__.py
@@ -4,5 +4,5 @@
 
 # Expose the corresponding C-Extension module with a well-known name at this
 # level.
-from .. import _load_extension
+from .._cext_loader import _load_extension
 _cextTransforms = _load_extension("_mlirTransforms")

diff  --git a/mlir/test/Bindings/Python/context_lifecycle.py b/mlir/test/Bindings/Python/context_lifecycle.py
index 460f41ccd4a7..c20270999425 100644
--- a/mlir/test/Bindings/Python/context_lifecycle.py
+++ b/mlir/test/Bindings/Python/context_lifecycle.py
@@ -1,7 +1,7 @@
 # RUN: %PYTHON %s
 # Standalone sanity check of context life-cycle.
 import gc
-import mlir
+import mlir.ir
 
 assert mlir.ir.Context._get_live_count() == 0
 

diff  --git a/mlir/test/Bindings/Python/dialects.py b/mlir/test/Bindings/Python/dialects.py
index 82e25b9e54fa..128f64cab199 100644
--- a/mlir/test/Bindings/Python/dialects.py
+++ b/mlir/test/Bindings/Python/dialects.py
@@ -35,7 +35,7 @@ def testUserDialectClass():
   d = ctx.dialects.std
   # Note that the standard dialect namespace prints as ''. Others will print
   # as "<Dialect %namespace (..."
-  # CHECK: <Dialect (class mlir.dialects.std._Dialect)>
+  # CHECK: <Dialect (class mlir.dialects._std_ops_gen._Dialect)>
   print(d)
   try:
     _ = ctx.dialects.not_existing
@@ -46,7 +46,7 @@ def testUserDialectClass():
 
   # Access using index.
   d = ctx.dialects["std"]
-  # CHECK: <Dialect (class mlir.dialects.std._Dialect)>
+  # CHECK: <Dialect (class mlir.dialects._std_ops_gen._Dialect)>
   print(d)
   try:
     _ = ctx.dialects["not_existing"]
@@ -57,7 +57,7 @@ def testUserDialectClass():
 
   # Using the 'd' alias.
   d = ctx.d["std"]
-  # CHECK: <Dialect (class mlir.dialects.std._Dialect)>
+  # CHECK: <Dialect (class mlir.dialects._std_ops_gen._Dialect)>
   print(d)
 
 run(testUserDialectClass)

diff  --git a/mlir/test/Bindings/Python/tools/linalg_opdsl/assignments.py b/mlir/test/Bindings/Python/dialects/linalg/opdsl/assignments.py
similarity index 86%
rename from mlir/test/Bindings/Python/tools/linalg_opdsl/assignments.py
rename to mlir/test/Bindings/Python/dialects/linalg/opdsl/assignments.py
index 225b38c8c529..e96bc0de2204 100644
--- a/mlir/test/Bindings/Python/tools/linalg_opdsl/assignments.py
+++ b/mlir/test/Bindings/Python/dialects/linalg/opdsl/assignments.py
@@ -1,6 +1,6 @@
-# RUN: %PYTHON -m mlir.tools.linalg_opdsl.dump_oplib --file %s | FileCheck %s
+# RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s | FileCheck %s
 
-from mlir.tools.linalg_opdsl.lang import *
+from mlir.dialects.linalg.opdsl.lang import *
 
 # CHECK: ---
 # CHECK-LABEL: matmul

diff  --git a/mlir/test/Bindings/Python/tools/linalg_opdsl/doctests.py b/mlir/test/Bindings/Python/dialects/linalg/opdsl/doctests.py
similarity index 69%
rename from mlir/test/Bindings/Python/tools/linalg_opdsl/doctests.py
rename to mlir/test/Bindings/Python/dialects/linalg/opdsl/doctests.py
index e9cddcc73519..4aae76884881 100644
--- a/mlir/test/Bindings/Python/tools/linalg_opdsl/doctests.py
+++ b/mlir/test/Bindings/Python/dialects/linalg/opdsl/doctests.py
@@ -9,5 +9,5 @@ def test_module(module_name):
   doctest.testmod(m, verbose=True, raise_on_error=True, report=True)
 
 
-test_module("mlir.tools.linalg_opdsl.lang.affine")
-test_module("mlir.tools.linalg_opdsl.lang.types")
+test_module("mlir.dialects.linalg.opdsl.lang.affine")
+test_module("mlir.dialects.linalg.opdsl.lang.types")

diff  --git a/mlir/test/Bindings/Python/tools/linalg_opdsl/interfaces.py b/mlir/test/Bindings/Python/dialects/linalg/opdsl/interfaces.py
similarity index 73%
rename from mlir/test/Bindings/Python/tools/linalg_opdsl/interfaces.py
rename to mlir/test/Bindings/Python/dialects/linalg/opdsl/interfaces.py
index d70820b83c8a..46689a07bbbb 100644
--- a/mlir/test/Bindings/Python/tools/linalg_opdsl/interfaces.py
+++ b/mlir/test/Bindings/Python/dialects/linalg/opdsl/interfaces.py
@@ -1,6 +1,6 @@
-# RUN: %PYTHON -m mlir.tools.linalg_opdsl.dump_oplib --file %s | FileCheck %s
+# RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s | FileCheck %s
 
-from mlir.tools.linalg_opdsl.lang import *
+from mlir.dialects.linalg.opdsl.lang import *
 
 # CHECK: ---
 # CHECK-LABEL: matmul

diff  --git a/mlir/test/Bindings/Python/tools/linalg_opdsl/lit.local.cfg b/mlir/test/Bindings/Python/dialects/linalg/opdsl/lit.local.cfg
similarity index 100%
rename from mlir/test/Bindings/Python/tools/linalg_opdsl/lit.local.cfg
rename to mlir/test/Bindings/Python/dialects/linalg/opdsl/lit.local.cfg

diff  --git a/mlir/test/Bindings/Python/tools/linalg_opdsl/shape_maps_iteration.py b/mlir/test/Bindings/Python/dialects/linalg/opdsl/shape_maps_iteration.py
similarity index 92%
rename from mlir/test/Bindings/Python/tools/linalg_opdsl/shape_maps_iteration.py
rename to mlir/test/Bindings/Python/dialects/linalg/opdsl/shape_maps_iteration.py
index 2627f2c5fcb0..61453da13f49 100644
--- a/mlir/test/Bindings/Python/tools/linalg_opdsl/shape_maps_iteration.py
+++ b/mlir/test/Bindings/Python/dialects/linalg/opdsl/shape_maps_iteration.py
@@ -1,6 +1,6 @@
-# RUN: %PYTHON -m mlir.tools.linalg_opdsl.dump_oplib --file %s | FileCheck %s
+# RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s | FileCheck %s
 
-from mlir.tools.linalg_opdsl.lang import *
+from mlir.dialects.linalg.opdsl.lang import *
 
 
 # Verify that simple case with iteration order defined lexically and reduction

diff  --git a/mlir/test/Bindings/Python/dialects/linalg/opdsl/test_core_named_ops.py b/mlir/test/Bindings/Python/dialects/linalg/opdsl/test_core_named_ops.py
new file mode 100644
index 000000000000..0dceba9c13a4
--- /dev/null
+++ b/mlir/test/Bindings/Python/dialects/linalg/opdsl/test_core_named_ops.py
@@ -0,0 +1,4 @@
+# RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib .ops.core_named_ops | FileCheck %s
+
+# Just verify that at least one known op is generated.
+# CHECK: name: matmul

diff  --git a/mlir/test/Bindings/Python/dialects/linalg.py b/mlir/test/Bindings/Python/dialects/linalg/ops.py
similarity index 100%
rename from mlir/test/Bindings/Python/dialects/linalg.py
rename to mlir/test/Bindings/Python/dialects/linalg/ops.py

diff  --git a/mlir/test/Bindings/Python/ir_operation.py b/mlir/test/Bindings/Python/ir_operation.py
index 034b28ec25bf..a71194b354cb 100644
--- a/mlir/test/Bindings/Python/ir_operation.py
+++ b/mlir/test/Bindings/Python/ir_operation.py
@@ -492,7 +492,7 @@ def testKnownOpView():
     # addf should map to a known OpView class in the std dialect.
     # We know the OpView for it defines an 'lhs' attribute.
     addf = module.body.operations[2]
-    # CHECK: <mlir.dialects.std._AddFOp object
+    # CHECK: <mlir.dialects._std_ops_gen._AddFOp object
     print(repr(addf))
     # CHECK: "custom.f32"()
     print(addf.lhs)

diff  --git a/mlir/test/Bindings/Python/tools/linalg_opdsl/test_core_named_ops.py b/mlir/test/Bindings/Python/tools/linalg_opdsl/test_core_named_ops.py
deleted file mode 100644
index d8c827468bf7..000000000000
--- a/mlir/test/Bindings/Python/tools/linalg_opdsl/test_core_named_ops.py
+++ /dev/null
@@ -1,4 +0,0 @@
-# RUN: %PYTHON -m mlir.tools.linalg_opdsl.dump_oplib .ops.core_named_ops | FileCheck %s
-
-# Just verify that at least one known op is generated.
-# CHECK: name: matmul

diff  --git a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
index c45ad495e4e8..d00555dc1c09 100644
--- a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
@@ -27,12 +27,12 @@ using namespace mlir::tblgen;
 constexpr const char *fileHeader = R"Py(
 # Autogenerated by mlir-tblgen; don't manually edit.
 
-from . import _cext as _ods_cext
-from . import extend_opview_class as _ods_extend_opview_class, segmented_accessor as _ods_segmented_accessor, equally_sized_accessor as _ods_equally_sized_accessor, get_default_loc_context as _ods_get_default_loc_context
+from ._ods_common import _cext as _ods_cext
+from ._ods_common import extend_opview_class as _ods_extend_opview_class, segmented_accessor as _ods_segmented_accessor, equally_sized_accessor as _ods_equally_sized_accessor, get_default_loc_context as _ods_get_default_loc_context
 _ods_ir = _ods_cext.ir
 
 try:
-  from . import _{0} as _ods_ext_module
+  from . import _{0}_ops_ext as _ods_ext_module
 except ImportError:
   _ods_ext_module = None

More information about the Mlir-commits mailing list