Move MLIR python sources to mlir/python.

[stellaraccident]

Summary

• NFC but has some fixes for CMake glitches discovered along the way (things not cleaning properly, co-mingled depends).
• Includes previously unsubmitted fix in D98681 and a TODO to fix it more appropriately in a smaller followup.

[stellaraccident]

Fix install.

[stellaraccident]

Rebase.

[mehdi_amini]

Seems mechanical, LGTM on the principle.

[stellaraccident]

In D101493#2727002, @mehdi_amini wrote:

Seems mechanical, LGTM on the principle.

Yes, mostly mechanical with CMake fixups that became apparent when things were separated.

[ftynse]

Tests seem to fail for me locally

: 'RUN: at line 1';   /usr/bin/python3.9 /home/zinenko/llvm-project/mlir/test/python/ir/value.py | /home/zinenko/llvm-project/build2/bin/FileCheck /home/zinenko/llvm-project/mlir/test/python/ir/value.py
--
Exit Code: 2

Command Output (stderr):
--
Traceback (most recent call last):
  File "/home/zinenko/llvm-project/build2/python/mlir/__init__.py", line 23, in <module>
    import _mlir_libs
ModuleNotFoundError: No module named '_mlir_libs'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/zinenko/llvm-project/build2/python/mlir/_cext_loader.py", line 12, in <module>
    import _mlir_libs
ModuleNotFoundError: No module named '_mlir_libs'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/zinenko/llvm-project/mlir/test/python/ir/value.py", line 4, in <module>
    from mlir.ir import *
  File "/home/zinenko/llvm-project/build2/python/mlir/__init__.py", line 28, in <module>
    from ._dlloader import preload_dependency as _preload_dependency
  File "/home/zinenko/llvm-project/build2/python/mlir/_dlloader.py", line 6, in <module>
    import platform
  File "/usr/lib/python3.9/platform.py", line 117, in <module>
    import re
  File "/usr/lib/python3.9/re.py", line 124, in <module>
    import enum
  File "/usr/lib/python3.9/enum.py", line 2, in <module>
    from types import MappingProxyType, DynamicClassAttribute
  File "/home/zinenko/llvm-project/mlir/test/python/ir/types.py", line 4, in <module>
    from mlir.ir import *
  File "/home/zinenko/llvm-project/build2/python/mlir/ir.py", line 6, in <module>
    from ._cext_loader import _reexport_cext
  File "/home/zinenko/llvm-project/build2/python/mlir/_cext_loader.py", line 17, in <module>
    from ._dlloader import preload_dependency as _preload_dependency
ImportError: cannot import name 'preload_dependency' from partially initialized module 'mlir._dlloader' (most likely due to a circular import) (/home/zinenko/llvm-project/build2/python/mlir/_dlloader.py)
FileCheck error: '<stdin>' is empty.
FileCheck command line:  /home/zinenko/llvm-project/build2/bin/FileCheck /home/zinenko/llvm-project/mlir/test/python/ir/value.py

--

********************
********************
Failed Tests (15):
  MLIR :: python/ir/affine_expr.py
  MLIR :: python/ir/affine_map.py
  MLIR :: python/ir/array_attributes.py
  MLIR :: python/ir/attributes.py
  MLIR :: python/ir/context_lifecycle.py
  MLIR :: python/ir/context_managers.py
  MLIR :: python/ir/debug.py
  MLIR :: python/ir/dialects.py
  MLIR :: python/ir/insertion_point.py
  MLIR :: python/ir/integer_set.py
  MLIR :: python/ir/location.py
  MLIR :: python/ir/module.py
  MLIR :: python/ir/operation.py
  MLIR :: python/ir/types.py
  MLIR :: python/ir/value.py

[stellaraccident]

In D101493#2727037, @ftynse wrote:

Tests seem to fail for me locally

: 'RUN: at line 1';   /usr/bin/python3.9 /home/zinenko/llvm-project/mlir/test/python/ir/value.py | /home/zinenko/llvm-project/build2/bin/FileCheck /home/zinenko/llvm-project/mlir/test/python/ir/value.py
--
Exit Code: 2

Command Output (stderr):
--
Traceback (most recent call last):
  File "/home/zinenko/llvm-project/build2/python/mlir/__init__.py", line 23, in <module>
    import _mlir_libs
ModuleNotFoundError: No module named '_mlir_libs'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/zinenko/llvm-project/build2/python/mlir/_cext_loader.py", line 12, in <module>
    import _mlir_libs
ModuleNotFoundError: No module named '_mlir_libs'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/zinenko/llvm-project/mlir/test/python/ir/value.py", line 4, in <module>
    from mlir.ir import *
  File "/home/zinenko/llvm-project/build2/python/mlir/__init__.py", line 28, in <module>
    from ._dlloader import preload_dependency as _preload_dependency
  File "/home/zinenko/llvm-project/build2/python/mlir/_dlloader.py", line 6, in <module>
    import platform
  File "/usr/lib/python3.9/platform.py", line 117, in <module>
    import re
  File "/usr/lib/python3.9/re.py", line 124, in <module>
    import enum
  File "/usr/lib/python3.9/enum.py", line 2, in <module>
    from types import MappingProxyType, DynamicClassAttribute
  File "/home/zinenko/llvm-project/mlir/test/python/ir/types.py", line 4, in <module>
    from mlir.ir import *
  File "/home/zinenko/llvm-project/build2/python/mlir/ir.py", line 6, in <module>
    from ._cext_loader import _reexport_cext
  File "/home/zinenko/llvm-project/build2/python/mlir/_cext_loader.py", line 17, in <module>
    from ._dlloader import preload_dependency as _preload_dependency
ImportError: cannot import name 'preload_dependency' from partially initialized module 'mlir._dlloader' (most likely due to a circular import) (/home/zinenko/llvm-project/build2/python/mlir/_dlloader.py)
FileCheck error: '<stdin>' is empty.
FileCheck command line:  /home/zinenko/llvm-project/build2/bin/FileCheck /home/zinenko/llvm-project/mlir/test/python/ir/value.py

--

********************
********************
Failed Tests (15):
  MLIR :: python/ir/affine_expr.py
  MLIR :: python/ir/affine_map.py
  MLIR :: python/ir/array_attributes.py
  MLIR :: python/ir/attributes.py
  MLIR :: python/ir/context_lifecycle.py
  MLIR :: python/ir/context_managers.py
  MLIR :: python/ir/debug.py
  MLIR :: python/ir/dialects.py
  MLIR :: python/ir/insertion_point.py
  MLIR :: python/ir/integer_set.py
  MLIR :: python/ir/location.py
  MLIR :: python/ir/module.py
  MLIR :: python/ir/operation.py
  MLIR :: python/ir/types.py
  MLIR :: python/ir/value.py

Tried to ping you on discord. Included here as well:

Regarding https://reviews.llvm.org/D101493, were you starting from a pre-existing build directory? (I had blown mine away and ninja check-mlir from scratch works)
Just verified again from a fresh build and tests pass.
Notably: this file should not exist (it was deleted quite some time ago): /home/zinenko/llvm-project/build2/python/mlir/__init__.py

[ftynse]

Tried to ping you on discord. Included here as well:

Sorry, I don't follow discord. Any synchronous discussion media basically shuts me off because of nine-hour time zone difference with most of the people. By the time I see the message, the discussion has usually ended or diverged and any reaction to previous messages would be largely off topic.

Regarding https://reviews.llvm.org/D101493, were you starting from a pre-existing build directory?

Yes, I was starting from a pre-existing build directory, but there seems to be another problem even with a clean directory.

FAIL: MLIR :: python/ir/value.py (766 of 954)
******************** TEST 'MLIR :: python/ir/value.py' FAILED ********************
Script:
--
: 'RUN: at line 1';   /usr/bin/python3.9 /home/zinenko/llvm-project/mlir/test/python/ir/value.py | /home/zinenko/llvm-project/build/bin/FileCheck /home/zinenko/llvm-project/mlir/test/python/ir/value.py
--
Exit Code: 1

Command Output (stderr):
--
Traceback (most recent call last):
  File "/home/zinenko/llvm-project/build/python/mlir/_cext_loader.py", line 12, in <module>
    import _mlir_libs
ModuleNotFoundError: No module named '_mlir_libs'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/zinenko/llvm-project/mlir/test/python/ir/value.py", line 4, in <module>
    from mlir.ir import *
  File "/home/zinenko/llvm-project/build/python/mlir/ir.py", line 6, in <module>
    from ._cext_loader import _reexport_cext
  File "/home/zinenko/llvm-project/build/python/mlir/_cext_loader.py", line 17, in <module>
    from ._dlloader import preload_dependency as _preload_dependency
  File "/home/zinenko/llvm-project/build/python/mlir/_dlloader.py", line 6, in <module>
    import platform
  File "/usr/lib/python3.9/platform.py", line 117, in <module>
    import re
  File "/usr/lib/python3.9/re.py", line 124, in <module>
    import enum
  File "/usr/lib/python3.9/enum.py", line 2, in <module>
    from types import MappingProxyType, DynamicClassAttribute
  File "/home/zinenko/llvm-project/mlir/test/python/ir/types.py", line 25, in <module>
    run(testParsePrint)
  File "/home/zinenko/llvm-project/mlir/test/python/ir/types.py", line 8, in run
    f()
  File "/home/zinenko/llvm-project/mlir/test/python/ir/types.py", line 15, in testParsePrint
    ctx = Context()
NameError: name 'Context' is not defined

[ftynse]

This particular problem seems to come from the name clash between the types.py test and the types.py distributed with Python. It's strange that you don't hit the problem.

[bondhugula]

I think you are missing an update to docs/Bindings/Python.md? (The paths there have to change.)

[stellaraccident]

Address comments and rebase.

[stellaraccident]

In D101493#2731131, @bondhugula wrote:

I think you are missing an update to docs/Bindings/Python.md? (The paths there have to change.)

Updated.

[stellaraccident]

In D101493#2728298, @ftynse wrote:

This particular problem seems to come from the name clash between the types.py test and the types.py distributed with Python. It's strange that you don't hit the problem.

Strange. In lieu of debugging, I just renamed that to builtin_types.py.

[stellaraccident]

Alex: Made a change that might appease your machine. Let me know.

[ftynse]

I tried that myself before with renaming (back) to ir_types and it worked, so it's okay.

[bondhugula]

Thanks for updating the docs!

mlir/CMakeLists.txt

else()
     message(WARNING "gtest not found, unittests will not be available")
   endif()
   add_subdirectory(test)
 endif()
 # Tools needs to come late to ensure that MLIR_ALL_LIBS is populated.
 # Generally things after this point may depend on MLIR_ALL_LIBS or libMLIR.so.
 add_subdirectory(tools)
 
+if(MLIR_BINDINGS_PYTHON_ENABLED)
+  # Python sources: built extensions come in via lib/Bindings/Python
+  add_subdirectory(python)
+endif()
+
 if( LLVM_INCLUDE_EXAMPLES )
   add_subdirectory(examples)
 endif()
 
 option(MLIR_INCLUDE_DOCS "Generate build targets for the MLIR docs."
   ${LLVM_INCLUDE_DOCS})
 if (MLIR_INCLUDE_DOCS)
   add_subdirectory(docs)

mlir/cmake/modules/AddMLIRPython.cmake

add_public_tablegen_target(
     ${tblgen_target})
   if(ARG_DEPENDS)
     add_dependencies(${tblgen_target} ${ARG_DEPENDS})
   endif()
 
   add_custom_command(
     TARGET ${tblgen_target} POST_BUILD
     COMMENT "Copying generated python source \"dialects/${dialect_filename}\""
+    BYPRODUCTS "${PROJECT_BINARY_DIR}/python/mlir/dialects/${dialect_filename}"
     COMMAND "${CMAKE_COMMAND}" -E copy_if_different
       "${CMAKE_CURRENT_BINARY_DIR}/${dialect_filename}"
       "${PROJECT_BINARY_DIR}/python/mlir/dialects/${dialect_filename}")
 endfunction()
 

mlir/cmake/modules/CMakeLists.txt

if (NOT LLVM_INSTALL_TOOLCHAIN_ONLY)
   if(mlir_has_exports)
     install(EXPORT MLIRTargets DESTINATION ${MLIR_INSTALL_PACKAGE_DIR}
             COMPONENT mlir-cmake-exports)
   endif()
 
   install(FILES
     ${CMAKE_CURRENT_BINARY_DIR}/CMakeFiles/MLIRConfig.cmake
     ${CMAKE_CURRENT_SOURCE_DIR}/AddMLIR.cmake
-    ${CMAKE_CURRENT_SOURCE_DIR}/AddMLIRPythonExtension.cmake
+    ${CMAKE_CURRENT_SOURCE_DIR}/AddMLIRPython.cmake
     ${CMAKE_CURRENT_SOURCE_DIR}/MLIRDetectPythonEnv.cmake
     DESTINATION ${MLIR_INSTALL_PACKAGE_DIR}
     COMPONENT mlir-cmake-exports)
 
   if(NOT LLVM_ENABLE_IDE)
     # Add a dummy target so this can be used with LLVM_DISTRIBUTION_COMPONENTS
     add_custom_target(mlir-cmake-exports)
     add_llvm_install_targets(install-mlir-cmake-exports
                              COMPONENT mlir-cmake-exports)
   endif()
 endif()

mlir/lib/Bindings/Python/CMakeLists.txt

-include(AddMLIRPythonExtension)
+include(AddMLIRPython)
 add_custom_target(MLIRBindingsPythonExtension)
 
 ################################################################################
-# Copy python source tree.
-################################################################################
-
-file(GLOB_RECURSE PY_SRC_FILES
-  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
-  "${CMAKE_CURRENT_SOURCE_DIR}/mlir/*.py")
-
-add_custom_target(MLIRBindingsPythonSources ALL
-  DEPENDS ${PY_SRC_FILES}
-)
-add_dependencies(MLIRBindingsPythonExtension MLIRBindingsPythonSources)
-
-foreach(PY_SRC_FILE ${PY_SRC_FILES})
-  set(PY_DEST_FILE "${PROJECT_BINARY_DIR}/python/${PY_SRC_FILE}")
-  get_filename_component(PY_DEST_DIR "${PY_DEST_FILE}" DIRECTORY)
-  file(MAKE_DIRECTORY "${PY_DEST_DIR}")
-  add_custom_command(
-    TARGET MLIRBindingsPythonSources PRE_BUILD
-    COMMENT "Copying python source ${PY_SRC_FILE} -> ${PY_DEST_FILE}"
-    DEPENDS "${PY_SRC_FILE}"
-    COMMAND "${CMAKE_COMMAND}" -E create_symlink
-        "${CMAKE_CURRENT_SOURCE_DIR}/${PY_SRC_FILE}" "${PY_DEST_FILE}"
-  )
-endforeach()
-
-################################################################################
-# Generate dialect-specific bindings.
-################################################################################
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonAsyncOps
-  TD_FILE AsyncOps.td
-  DIALECT_NAME async_dialect)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonAsyncOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonBuiltinOps
-  TD_FILE BuiltinOps.td
-  DIALECT_NAME builtin)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonBuiltinOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonGPUOps
-  TD_FILE GPUOps.td
-  DIALECT_NAME gpu)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonGPUOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonLinalgOps
-  TD_FILE LinalgOps.td
-  DIALECT_NAME linalg
-  DEPENDS LinalgOdsGen)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonLinalgOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonMemRefOps
-  TD_FILE MemRefOps.td
-  DIALECT_NAME memref)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonMemRefOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonShapeOps
-  TD_FILE ShapeOps.td
-  DIALECT_NAME shape)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonShapeOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonStandardOps
-  TD_FILE StandardOps.td
-  DIALECT_NAME std)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonStandardOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonTensorOps
-  TD_FILE TensorOps.td
-  DIALECT_NAME tensor)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonTensorOps)
-
-################################################################################
 # Build core python extension
 ################################################################################
 add_mlir_python_extension(MLIRCoreBindingsPythonExtension _mlir
   INSTALL_DIR
     python
   SOURCES
     DialectLinalg.cpp
     MainModule.cpp
     IRAffine.cpp
     IRAttributes.cpp
     IRCore.cpp
     IRTypes.cpp
     PybindUtils.cpp
     Pass.cpp
     ExecutionEngine.cpp
 )
 add_dependencies(MLIRBindingsPythonExtension MLIRCoreBindingsPythonExtension)
 
-# Note that we copy from the source tree just like for headers because
-# it will not be polluted with py_cache runtime artifacts (from testing and
-# such).
-install(
-  DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/mlir
-  DESTINATION python
-  COMPONENT MLIRBindingsPythonSources
-  FILES_MATCHING PATTERN "*.py"
-)
-
-if (NOT LLVM_ENABLE_IDE)
-  add_llvm_install_targets(
-    install-MLIRBindingsPythonSources
-    DEPENDS MLIRBindingsPythonSources
-    COMPONENT MLIRBindingsPythonSources)
-endif()
-
-# Dialect sources are generated. Install separately.
-# Note that __pycache__ directories may have been left by tests and other
-# executions. And __init__.py is handled as a regular source file.
-install(
-  DIRECTORY ${PROJECT_BINARY_DIR}/python/mlir/dialects
-  DESTINATION python/mlir
-  COMPONENT MLIRBindingsPythonDialects
-  FILES_MATCHING PATTERN "*.py"
-  PATTERN "__pycache__" EXCLUDE
-  PATTERN "__init__.py" EXCLUDE
-)
-
-if (NOT LLVM_ENABLE_IDE)
-  add_llvm_install_targets(
-    install-MLIRBindingsPythonDialects
-    DEPENDS MLIRBindingsPythonSources
-    COMPONENT MLIRBindingsPythonDialects)
-endif()
-
 add_subdirectory(Transforms)
 add_subdirectory(Conversions)
 
 add_mlir_python_extension(MLIRAsyncPassesBindingsPythonExtension _mlirAsyncPasses
   INSTALL_DIR
     python
   SOURCES
     AsyncPasses.cpp

mlir/python/.style.yapf

 [style]
   based_on_style = google
   column_limit = 80
   indent_width = 2

mlir/python/CMakeLists.txt

+################################################################################
+# Copy python source tree.
+################################################################################
+
+file(GLOB_RECURSE PY_SRC_FILES
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "${CMAKE_CURRENT_SOURCE_DIR}/mlir/*.py")
+
+add_custom_target(MLIRBindingsPythonSources ALL
+  DEPENDS
+    ${PY_SRC_FILES}
+)
+
+foreach(PY_SRC_FILE ${PY_SRC_FILES})
+  set(PY_DEST_FILE "${PROJECT_BINARY_DIR}/python/${PY_SRC_FILE}")
+  get_filename_component(PY_DEST_DIR "${PY_DEST_FILE}" DIRECTORY)
+  file(MAKE_DIRECTORY "${PY_DEST_DIR}")
+  add_custom_command(
+    TARGET MLIRBindingsPythonSources PRE_BUILD
+    COMMENT "Copying python source ${PY_SRC_FILE} -> ${PY_DEST_FILE}"
+    DEPENDS "${PY_SRC_FILE}"
+    BYPRODUCTS "${PY_DEST_FILE}"
+    COMMAND "${CMAKE_COMMAND}" -E create_symlink
+        "${CMAKE_CURRENT_SOURCE_DIR}/${PY_SRC_FILE}" "${PY_DEST_FILE}"
+  )
+endforeach()
+
+# Note that we copy from the source tree just like for headers because
+# it will not be polluted with py_cache runtime artifacts (from testing and
+# such).
+install(
+  DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/mlir
+  DESTINATION python
+  COMPONENT MLIRBindingsPythonSources
+  FILES_MATCHING PATTERN "*.py"
+)
+
+if (NOT LLVM_ENABLE_IDE)
+  add_llvm_install_targets(
+    install-MLIRBindingsPythonSources
+    DEPENDS MLIRBindingsPythonSources
+    COMPONENT MLIRBindingsPythonSources)
+endif()
+
+################################################################################
+# Generated sources.
+################################################################################
+
+add_subdirectory(mlir/dialects)

mlir/python/mlir/_cext_loader.py

 #  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
 """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
 #   preload_dependency(public_name): Loads a shared-library/DLL into the
 #     namespace. TODO: Remove this in favor of a more robust mechanism.
 # Conditionally switch based on whether we are in a package context.
 try:
   import _mlir_libs
 except ModuleNotFoundError:
   # Assume that we are in-tree.
   # The _dlloader takes care of platform specific setup before we try to
   # load a shared library.
   from ._dlloader import preload_dependency as _preload_dependency
 
   def _load_extension(name):
     import importlib
     return importlib.import_module(name)  # i.e. '_mlir' at the top level
 else:
   # Packaged distribution.
   _load_extension = _mlir_libs.load_extension
   _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 .._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 ._cext_loader import _reexport_cext
     _reexport_cext("ir", __name__)
     del _reexport_cext
   """
   import sys
   target_module = sys.modules[target_module_name]
   source_module = getattr(_cext, cext_module_name)
   for attr_name in dir(source_module):
     if not attr_name.startswith("__"):
       setattr(target_module, attr_name, getattr(source_module, attr_name))
 
 
 # Add our 'dialects' parent module to the search path for implementations.
 _cext.globals.append_dialect_search_prefix("mlir.dialects")

mlir/python/mlir/_dlloader.py

 #  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
 
 import os
 import platform
 
 _is_windows = platform.system() == "Windows"
 _this_directory = os.path.dirname(__file__)
 
 # The standard LLVM build/install tree for Windows is laid out as:
 #   bin/
 #     MLIRPublicAPI.dll
 #   python/
 #     _mlir.*.pyd (dll extension)
 #     mlir/
 #       _dlloader.py (this file)
 # First check the python/ directory level for DLLs co-located with the pyd
 # file, and then fall back to searching the bin/ directory.
 # TODO: This should be configurable at some point.
 _dll_search_path = [
   os.path.join(_this_directory, ".."),
   os.path.join(_this_directory, "..", "..", "bin"),
 ]
 
 # Stash loaded DLLs to keep them alive.
 _loaded_dlls = []
 
 def preload_dependency(public_name):
   """Preloads a dylib by its soname or DLL name.
 
   On Windows and Linux, doing this prior to loading a dependency will populate
   the library in the flat namespace so that a subsequent library that depend
   on it will resolve to this preloaded version.
 
   On OSX, resolution is completely path based so this facility no-ops. On
   Linux, as long as RPATHs are setup properly, resolution is path based but
   this facility can still act as an escape hatch for relocatable distributions.
   """
   if _is_windows:
     _preload_dependency_windows(public_name)
 
 
 def _preload_dependency_windows(public_name):
   dll_basename = public_name + ".dll"
   found_path = None
   for search_dir in _dll_search_path:
     candidate_path = os.path.join(search_dir, dll_basename)
     if os.path.exists(candidate_path):
       found_path = candidate_path
       break
 
   if found_path is None:
     raise RuntimeError(
       f"Unable to find dependency DLL {dll_basename} in search "
       f"path {_dll_search_path}")
 
   import ctypes
   _loaded_dlls.append(ctypes.CDLL(found_path))

mlir/python/mlir/conversions/__init__.py

 #  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
 
 # Expose the corresponding C-Extension module with a well-known name at this
 # level.
 from .._cext_loader import _load_extension
 _cextConversions = _load_extension("_mlirConversions")

mlir/python/mlir/dialects/AsyncOps.td

 //===-- AsyncOps.td - Entry point async_dialect bindings --*- tablegen -*-===//
 //
 // 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
 //
 //===---------------------------------------------------------------------===//
 
 #ifndef PYTHON_BINDINGS_ASYNC_OPS
 #define PYTHON_BINDINGS_ASYNC_OPS
 
 include "mlir/Bindings/Python/Attributes.td"
 include "mlir/Dialect/Async/IR/AsyncOps.td"
 
 #endif

mlir/python/mlir/dialects/BuiltinOps.td

 //===-- BuiltinOps.td - Entry point for builtin bindings ---*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef PYTHON_BINDINGS_BUILTIN_OPS
 #define PYTHON_BINDINGS_BUILTIN_OPS
 
 include "mlir/Bindings/Python/Attributes.td"
 include "mlir/IR/BuiltinOps.td"
 
 #endif

mlir/python/mlir/dialects/CMakeLists.txt

+include(AddMLIRPython)
+
+################################################################################
+# Generate dialect-specific bindings.
+################################################################################
+
+add_mlir_dialect_python_bindings(MLIRBindingsPythonAsyncOps
+  TD_FILE AsyncOps.td
+  DIALECT_NAME async_dialect)
+add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonAsyncOps)
+
+add_mlir_dialect_python_bindings(MLIRBindingsPythonBuiltinOps
+  TD_FILE BuiltinOps.td
+  DIALECT_NAME builtin)
+add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonBuiltinOps)
+
+add_mlir_dialect_python_bindings(MLIRBindingsPythonGPUOps
+  TD_FILE GPUOps.td
+  DIALECT_NAME gpu)
+add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonGPUOps)
+
+add_mlir_dialect_python_bindings(MLIRBindingsPythonLinalgOps
+  TD_FILE LinalgOps.td
+  DIALECT_NAME linalg
+  DEPENDS LinalgOdsGen)
+add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonLinalgOps)
+
+add_mlir_dialect_python_bindings(MLIRBindingsPythonMemRefOps
+  TD_FILE MemRefOps.td
+  DIALECT_NAME memref)
+add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonMemRefOps)
+
+add_mlir_dialect_python_bindings(MLIRBindingsPythonShapeOps
+  TD_FILE ShapeOps.td
+  DIALECT_NAME shape)
+add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonShapeOps)
+
+add_mlir_dialect_python_bindings(MLIRBindingsPythonStandardOps
+  TD_FILE StandardOps.td
+  DIALECT_NAME std)
+add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonStandardOps)
+
+add_mlir_dialect_python_bindings(MLIRBindingsPythonTensorOps
+  TD_FILE TensorOps.td
+  DIALECT_NAME tensor)
+add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonTensorOps)
+
+################################################################################
+# Installation.
+################################################################################
+
+# Dialect sources are generated. Install separately.
+# Note that __pycache__ directories may have been left by tests and other
+# executions. And __init__.py is handled as a regular source file.
+# TODO: Eliminate this glob install, instead adding INSTALL_COMPONENT to
+# add_mlir_dialect_python_bindings and installing the precise file there.
+install(
+  DIRECTORY ${PROJECT_BINARY_DIR}/python/mlir/dialects
+  DESTINATION python/mlir
+  COMPONENT MLIRBindingsPythonDialects
+  FILES_MATCHING PATTERN "_*_gen.py"
+  PATTERN "__pycache__" EXCLUDE
+  PATTERN "__init__.py" EXCLUDE
+)
+
+if (NOT LLVM_ENABLE_IDE)
+  add_llvm_install_targets(
+    install-MLIRBindingsPythonDialects
+    DEPENDS MLIRBindingsPythonSources
+    COMPONENT MLIRBindingsPythonDialects)
+endif()

mlir/python/mlir/dialects/GPUOps.td

 //===-- GPUOps.td - Entry point GPU_dialect bindings ------*- tablegen -*-===//
 //
 // 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
 //
 //===---------------------------------------------------------------------===//
 
 #ifndef PYTHON_BINDINGS_GPU_OPS
 #define PYTHON_BINDINGS_GPU_OPS
 
 include "mlir/Bindings/Python/Attributes.td"
 include "mlir/Dialect/GPU/GPUOps.td"
 
 #endif

mlir/python/mlir/dialects/LinalgOps.td

 //===-- LinalgOps.td - Entry point for linalg bind ---------*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef PYTHON_BINDINGS_LINALG_OPS
 #define PYTHON_BINDINGS_LINALG_OPS
 
 include "mlir/Bindings/Python/Attributes.td"
 include "mlir/Dialect/Linalg/IR/LinalgOps.td"
 include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.td"
 
 #endif

mlir/python/mlir/dialects/MemRefOps.td

 //===-- MemRefOps.td - Entry point for MemRefOps bind ------*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef PYTHON_BINDINGS_MEMREF_OPS
 #define PYTHON_BINDINGS_MEMREF_OPS
 
 include "mlir/Bindings/Python/Attributes.td"
 include "mlir/Dialect/MemRef/IR/MemRefOps.td"
 
 #endif

mlir/python/mlir/dialects/ShapeOps.td

 //===-- ShapeOps.td - Entry point for TensorOps bind -------*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef PYTHON_BINDINGS_SHAPE_OPS
 #define PYTHON_BINDINGS_SHAPE_OPS
 
 include "mlir/Bindings/Python/Attributes.td"
 include "mlir/Dialect/Shape/IR/ShapeOps.td"
 
 #endif

mlir/python/mlir/dialects/StandardOps.td

 //===-- StandardOps.td - Entry point for StandardOps bind --*- tablegen -*-===//
 //
 // 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 is the main file from which the Python bindings for the Standard
 // dialect are generated.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef PYTHON_BINDINGS_STANDARD_OPS
 #define PYTHON_BINDINGS_STANDARD_OPS
 
 include "mlir/Bindings/Python/Attributes.td"
 include "mlir/Dialect/StandardOps/IR/Ops.td"
 
 #endif

mlir/python/mlir/dialects/TensorOps.td

 //===-- TensorOps.td - Entry point for TensorOps bind ------*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef PYTHON_BINDINGS_TENSOR_OPS
 #define PYTHON_BINDINGS_TENSOR_OPS
 
 include "mlir/Bindings/Python/Attributes.td"
 include "mlir/Dialect/Tensor/IR/TensorOps.td"
 
 #endif

mlir/python/mlir/dialects/_builtin_ops_ext.py

 #  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 typing import Optional, Sequence
 
 import inspect
 
 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))
     body = self.regions[0].blocks.append()
 
   @property
   def body(self):
     return self.regions[0].blocks[0]
 
 
 class FuncOp:
   """Specialization for the func op class."""
 
   def __init__(self,
                name,
                type,
                *,
                visibility=None,
                body_builder=None,
                loc=None,
                ip=None):
     """
     Create a FuncOp with the provided `name`, `type`, and `visibility`.
     - `name` is a string representing the function name.
     - `type` is either a FunctionType or a pair of list describing inputs and
       results.
     - `visibility` is a string matching `public`, `private`, or `nested`. None
       implies private visibility.
     - `body_builder` is an optional callback, when provided a new entry block
       is created and the callback is invoked with the new op as argument within
       an InsertionPoint context already set for the block. The callback is
       expected to insert a terminator in the block.
     """
     sym_name = StringAttr.get(str(name))
 
     # If the type is passed as a tuple, build a FunctionType on the fly.
     if isinstance(type, tuple):
       type = FunctionType.get(inputs=type[0], results=type[1])
 
     type = TypeAttr.get(type)
     sym_visibility = StringAttr.get(
         str(visibility)) if visibility is not None else None
     super().__init__(sym_name, type, sym_visibility, loc=loc, ip=ip)
     if body_builder:
       entry_block = self.add_entry_block()
       with InsertionPoint(entry_block):
         body_builder(self)
 
   @property
   def is_external(self):
     return len(self.regions[0].blocks) == 0
 
   @property
   def body(self):
     return self.regions[0]
 
   @property
   def type(self):
     return FunctionType(TypeAttr(self.attributes["type"]).value)
 
   @property
   def visibility(self):
     return self.attributes["sym_visibility"]
 
   @property
   def name(self):
     return self.attributes["sym_name"]
 
   @property
   def entry_block(self):
     if self.is_external:
       raise IndexError('External function does not have a body')
     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.
     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)
     return self.body.blocks[0]
 
   @classmethod
   def from_py_func(FuncOp,
                    *inputs: Type,
                    results: Optional[Sequence[Type]] = None,
                    name: Optional[str] = None):
     """Decorator to define an MLIR FuncOp specified as a python function.
 
     Requires that an `mlir.ir.InsertionPoint` and `mlir.ir.Location` are
     active for the current thread (i.e. established in a `with` block).
 
     When applied as a decorator to a Python function, an entry block will
     be constructed for the FuncOp with types as specified in `*inputs`. The
     block arguments will be passed positionally to the Python function. In
     addition, if the Python function accepts keyword arguments generally or
     has a corresponding keyword argument, the following will be passed:
       * `func_op`: The `func` op being defined.
 
     By default, the function name will be the Python function `__name__`. This
     can be overriden by passing the `name` argument to the decorator.
 
     If `results` is not specified, then the decorator will implicitly
     insert a `ReturnOp` with the `Value`'s returned from the decorated
     function. It will also set the `FuncOp` type with the actual return
     value types. If `results` is specified, then the decorated function
     must return `None` and no implicit `ReturnOp` is added (nor are the result
     types updated). The implicit behavior is intended for simple, single-block
     cases, and users should specify result types explicitly for any complicated
     cases.
 
     The decorated function can further be called from Python and will insert
     a `CallOp` at the then-current insertion point, returning either None (
     if no return values), a unary Value (for one result), or a list of Values).
     This mechanism cannot be used to emit recursive calls (by construction).
     """
 
     def decorator(f):
       from . import std
       # Introspect the callable for optional features.
       sig = inspect.signature(f)
       has_arg_func_op = False
       for param in sig.parameters.values():
         if param.kind == param.VAR_KEYWORD:
           has_arg_func_op = True
         if param.name == "func_op" and (param.kind
                                         == param.POSITIONAL_OR_KEYWORD or
                                         param.kind == param.KEYWORD_ONLY):
           has_arg_func_op = True
 
       # Emit the FuncOp.
       implicit_return = results is None
       symbol_name = name or f.__name__
       function_type = FunctionType.get(
           inputs=inputs, results=[] if implicit_return else results)
       func_op = FuncOp(name=symbol_name, type=function_type)
       with InsertionPoint(func_op.add_entry_block()):
         func_args = func_op.entry_block.arguments
         func_kwargs = {}
         if has_arg_func_op:
           func_kwargs["func_op"] = func_op
         return_values = f(*func_args, **func_kwargs)
         if not implicit_return:
           return_types = list(results)
           assert return_values is None, (
               "Capturing a python function with explicit `results=` "
               "requires that the wrapped function returns None.")
         else:
           # Coerce return values, add ReturnOp and rewrite func type.
           if return_values is None:
             return_values = []
           elif isinstance(return_values, Value):
             return_values = [return_values]
           else:
             return_values = list(return_values)
           std.ReturnOp(return_values)
           # Recompute the function type.
           return_types = [v.type for v in return_values]
           function_type = FunctionType.get(inputs=inputs, results=return_types)
           func_op.attributes["type"] = TypeAttr.get(function_type)
 
       def emit_call_op(*call_args):
         call_op = std.CallOp(return_types, FlatSymbolRefAttr.get(symbol_name),
                              call_args)
         if return_types is None:
           return None
         elif len(return_types) == 1:
           return call_op.result
         else:
           return call_op.results
 
       wrapped = emit_call_op
       wrapped.__name__ = f.__name__
       wrapped.func_op = func_op
       return wrapped
 
     return decorator

mlir/python/mlir/dialects/_linalg_ops_ext.py

 #  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 typing import Optional, Sequence, Union
 from ..ir import *
 from ._ods_common import get_default_loc_context
 # TODO: resolve name collision for Linalg functionality that is injected inside
 # the _mlir.dialects.linalg directly via pybind.
 from _mlir.dialects.linalg import fill_builtin_region
 
 
 def isa(cls : Type, ty : Type):
   try:
     cls(ty)
     return True
   except ValueError:
     return False
 
 
 class FillOp:
   """Extends the linalg.fill op."""
 
   def __init__(self,
                output: Value,
                value: Value,
                *,
                loc=None,
                ip=None):
     results = []
     if isa(RankedTensorType, output.type):
       results = [output.type]
     op = self.build_generic(results=results,
                             operands=[output, value],
                             attributes=None,
                             loc=loc,
                             ip=ip)
     OpView.__init__(self, op)
     linalgDialect = Context.current.get_dialect_descriptor("linalg")
     fill_builtin_region(linalgDialect, self.operation, [value])
     # TODO: self.result is None. When len(results) == 1 we expect it to be
     # results[0] as per _linalg_ops_gen.py. This seems like an orthogonal bug
     # in the generator of _linalg_ops_gen.py where we have:
     # ```
     # def result(self):
     #   return self.operation.results[0] \
     #     if len(self.operation.results) > 1 else None
     # ```
 
 
 class InitTensorOp:
   """Extends the linalg.init_tensor op."""
 
   def __init__(self,
                sizes: Union[Sequence[int], Sequence[Value]],
                element_type: Type,
                *,
                loc=None,
                ip=None):
     """Constructs an `init_tensor` with either static or dynamic sizes."""
     context = get_default_loc_context(loc)
     operands = []
     attributes = {}
     # TODO: Refactor the InitTensorOp to take an element type attribute and
     # then use normal result type inference, unifying the Python and C++ side
     # with a standard mechanism (versus stashing that in builders).
     if sizes and isinstance(sizes[0], Value):
       # Dynamic sizes.
       operands.extend(sizes)
       static_size_ints = [-1] * len(sizes)
       result_type = RankedTensorType.get(static_size_ints, element_type)
     else:
       # Static sizes.
       result_type = RankedTensorType.get(sizes, element_type)
       static_size_ints = sizes
 
     index_type = IndexType.get(context)
     attributes["static_sizes"] = ArrayAttr.get(
         [IntegerAttr.get(index_type, s) for s in static_size_ints],
         context=context)
     op = self.build_generic(results=[result_type],
                             operands=operands,
                             attributes=attributes,
                             loc=loc,
                             ip=ip)
     OpView.__init__(self, op)
 
 
 class StructuredOpMixin:
   """All structured ops use the same mixin class."""
 
   def __init__(self, inputs, outputs=(), results=(), loc=None, ip=None):
     super().__init__(
         self.build_generic(results=list(results),
                            operands=[list(inputs), list(outputs)],
                            loc=loc,
                            ip=ip))
 
 
 def select_opview_mixin(parent_opview_cls):
   # TODO: This shouldn't be a heuristic: we should have a way to annotate
   # the OpView to note that it is a structured op.
   if ("__init__" not in parent_opview_cls.__dict__ and
       hasattr(parent_opview_cls, "inputs") and
       hasattr(parent_opview_cls, "outputs") and
       hasattr(parent_opview_cls, "result_tensors")):
     return StructuredOpMixin

mlir/python/mlir/dialects/_ods_common.py

 #  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
 
 # Re-export the parent _cext so that every level of the API can get it locally.
 from .._cext_loader import _cext
 
 __all__ = [
     "equally_sized_accessor",
     "extend_opview_class",
     "get_default_loc_context",
     "segmented_accessor",
 ]
 
 
 def extend_opview_class(ext_module):
   """Decorator to extend an OpView class from an extension module.
 
   Extension modules can expose various entry-points:
     Stand-alone class with the same name as a parent OpView class (i.e.
     "ReturnOp"). A name-based match is attempted first before falling back
     to a below mechanism.
 
     def select_opview_mixin(parent_opview_cls):
       If defined, allows an appropriate mixin class to be selected dynamically
       based on the parent OpView class. Should return NotImplemented if a
       decision is not made.
 
   Args:
     ext_module: A module from which to locate extensions. Can be None if not
       available.
 
   Returns:
     A decorator that takes an OpView subclass and further extends it as
     needed.
   """
 
   def class_decorator(parent_opview_cls: type):
     if ext_module is None:
       return parent_opview_cls
     mixin_cls = NotImplemented
     # First try to resolve by name.
     try:
       mixin_cls = getattr(ext_module, parent_opview_cls.__name__)
     except AttributeError:
       # Fall back to a select_opview_mixin hook.
       try:
         select_mixin = getattr(ext_module, "select_opview_mixin")
       except AttributeError:
         pass
       else:
         mixin_cls = select_mixin(parent_opview_cls)
 
     if mixin_cls is NotImplemented or mixin_cls is None:
       return parent_opview_cls
 
     # Have a mixin_cls. Create an appropriate subclass.
     try:
 
       class LocalOpView(mixin_cls, parent_opview_cls):
         pass
     except TypeError as e:
       raise TypeError(
           f"Could not mixin {mixin_cls} into {parent_opview_cls}") from e
     LocalOpView.__name__ = parent_opview_cls.__name__
     LocalOpView.__qualname__ = parent_opview_cls.__qualname__
     return LocalOpView
 
   return class_decorator
 
 
 def segmented_accessor(elements, raw_segments, idx):
   """
   Returns a slice of elements corresponding to the idx-th segment.
 
     elements: a sliceable container (operands or results).
     raw_segments: an mlir.ir.Attribute, of DenseIntElements subclass containing
         sizes of the segments.
     idx: index of the segment.
   """
   segments = _cext.ir.DenseIntElementsAttr(raw_segments)
   start = sum(segments[i] for i in range(idx))
   end = start + segments[idx]
   return elements[start:end]
 
 
 def equally_sized_accessor(elements, n_variadic, n_preceding_simple,
                            n_preceding_variadic):
   """
   Returns a starting position and a number of elements per variadic group
   assuming equally-sized groups and the given numbers of preceding groups.
 
     elements: a sequential container.
     n_variadic: the number of variadic groups in the container.
     n_preceding_simple: the number of non-variadic groups preceding the current
         group.
     n_preceding_variadic: the number of variadic groups preceding the current
         group.
   """
 
   total_variadic_length = len(elements) - n_variadic + 1
   # This should be enforced by the C++-side trait verifier.
   assert total_variadic_length % n_variadic == 0
 
   elements_per_group = total_variadic_length // n_variadic
   start = n_preceding_simple + n_preceding_variadic * elements_per_group
   return start, elements_per_group
 
 
 def get_default_loc_context(location=None):
   """
   Returns a context in which the defaulted location is created. If the location
   is None, takes the current location from the stack, raises ValueError if there
   is no location on the stack.
   """
   if location is None:
     # Location.current raises ValueError if there is no current location.
     return _cext.ir.Location.current.context
   return location.context

mlir/python/mlir/dialects/async_dialect/__init__.py

 #  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 .._async_dialect_ops_gen import *

mlir/python/mlir/dialects/async_dialect/passes/__init__.py

 #  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 ...._cext_loader import _load_extension
 _cextAsyncPasses = _load_extension("_mlirAsyncPasses")

mlir/python/mlir/dialects/builtin.py

 #  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 *

mlir/python/mlir/dialects/gpu/__init__.py

 #  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 .._gpu_ops_gen import *

mlir/python/mlir/dialects/gpu/passes/__init__.py

 #  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 ...._cext_loader import _load_extension
 _cextGPUPasses = _load_extension("_mlirGPUPasses")

mlir/python/mlir/dialects/linalg/__init__.py

 #  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
 
 # These are the backing OpView classes generated from the linalg tablegen
 # definitions following these steps:
 #   DSL -> YAML -> tblgen -> pytblgen -> build/.../_linalg_ops_gen.py.
 from .._linalg_ops_gen import *
 
 # These are the ground truth functions defined as:
 # ```
 #    @linalg_structured_op
 #    def matmul(A=TensorDef(T1, S.M, S.K),
 #               B=TensorDef(T2, S.K, S.N),
 #               C=TensorDef(U, S.M, S.N, output=True)):
 # ```
 # using the linalg-py eDSL.
 # The linalg-py eDSL builds a python representation (PyRepr) that is 
 # used in following ways:
 #  1. PyRepr -> YAML to generate the C++ and Python .td files. These
 #     then turn into the core C++ Op classes and Python OpView classes
 #     respectively (made available in _linalg_ops_gen). The generic OpView class 
 #     mechanism makes the C++ classes available to python through the CAPI.
 #     PyRepr -> YAML currently occurs before compiler compile time.
 #     The other steps in this category occur at compiler compile time.
 #  2. PyRepr -> linalg.core_named_ops calls: piggybacks on the 
 #     _linalg_ops_gen classes and the OpView mechanism to build IR at
 #     runtime in python:
 #       a. by default, the Named Op Form is emitted, e.g.:
 #          `linalg.matmul(lhs, rhs, outs=[out])` creates the following IR:
 #          ```
 #             %1 = linalg.matmul ins(%arg0, %arg1 : tensor<4x16xf32>, tensor<16x8xf32>) 
 #                               outs(%0 : tensor<4x8xf32>)
 #                  -> tensor<4x8xf32>   
 #          ```
 #       b. by setting emit_generic=True, the Generic Op Form is emitted, e.g.:
 #           `linalg.matmul(lhs, rhs, outs=[out], emit_generic=True)` creates the following IR:
 #          ```
 #             %1 = linalg.generic {indexing_maps = [...], iterator_types = [...]} 
 #               ins(%arg0, %arg1 : tensor<4x16xf32>, tensor<16x8xf32>) 
 #              outs(%0 : tensor<4x8xf32>) {
 #               ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):  
 #                  ...
 #                  linalg.yield %3 : f32
 #             } -> tensor<4x8xf32>  
 #          ```
 #  3. PyRepr -> Runtime Custom Op definitions: directly generates a
 #     linalg.generic form like in 2.b.
 #     !!!WARNING!!!: if one creates a runtime custom op with the same name 
 #     as an existing core named op, step 2. will likely take precedence.
 #     TODO: guard against surprises and fail create Runtime Custom Ops with 
 #     the same name as existing Core Named Ops.
 from .opdsl.ops.core_named_ops import *

mlir/python/mlir/dialects/linalg/opdsl/dump_oplib.py

 #!/usr/bin/which python
 # Command line tool to load an oplib module and dump all of the operations
 # it contains in some format.
 """Loads one or more modules containing op definitions and dumps them.
 
 The dump format can be:
 
 * `--dump_format=yaml` (default)
 * `--dump_format=repr`
 
 Positional arguments are interpreted as module names (optionally, relative to
 this module). Loose module files can be specified via `--file <filepath>`.
 
 Sample usage:
   # Dump the YAML op definitions for the core named ops (as in the dialect
   # source tree).
   python -m mlir.dialects.linalg.opdsl.dump_oplib .ops.core_named_ops
 
 Note: YAML output is emitted in "document list" format with each operation
 as its own "document". Practically, this means that each operation (or group
 of composite ops) is emitted with a "---" preceding it, which can be useful
 for testing.
 """
 
 import argparse
 import importlib
 
 from .lang import *
 from .lang.config import *
 from .lang.yaml_helper import *
 
 
 def create_arg_parser() -> argparse.ArgumentParser:
   p = argparse.ArgumentParser(description="Dump an oplib in various formats")
   p.add_argument("modules",
                  metavar="M",
                  type=str,
                  nargs="*",
                  help="Op module to dump")
   p.add_argument("--file",
                  metavar="F",
                  type=str,
                  nargs="*",
                  help="Python op file to dump")
   p.add_argument("--format",
                  type=str,
                  dest="format",
                  default="yaml",
                  choices=("yaml", "repr"),
                  help="Format in which to dump")
   return p
 
 
 def load_module_from_file(module_name, file_path):
   spec = importlib.util.spec_from_file_location(module_name, file_path)
   m = importlib.util.module_from_spec(spec)
   spec.loader.exec_module(m)
   return m
 
 
 def main(args):
   # Load all configs.
   configs = []
   modules = []
   for module_name in args.modules:
     modules.append(
         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:
     for attr_name, value in m.__dict__.items():
       # TODO: This class layering is awkward.
       if isinstance(value, DefinedOpCallable):
         try:
           linalg_config = LinalgOpConfig.from_linalg_op_def(value.model)
         except Exception as e:
           raise ValueError(
               f"Could not create LinalgOpConfig from {value.model}") from e
         configs.extend(linalg_config)
 
   # Print.
   if args.format == "yaml":
     print(yaml_dump_all(configs))
   elif args.format == "repr":
     for config in configs:
       print(repr(config))
 
 
 if __name__ == "__main__":
   main(create_arg_parser().parse_args())

mlir/python/mlir/dialects/linalg/opdsl/lang/__init__.py

 from .dsl import *

mlir/python/mlir/dialects/linalg/opdsl/lang/affine.py

 #  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
 """DSL for constructing affine expressions and maps.
 
 These python wrappers allow construction of affine expressions in a more
 pythonic fashion that is later instantiated as an IR AffineExpr. Separating the
 AST from construction of the map allows for manipulations of symbols and dims
 beyond the scope of one expression.
 
 Affine expression construction:
   >>> with _ir.Context():
   ...   s = AffineBuildState()
   ...   (S.K + S.M).build(s)
   ...   (S.K * S.M).build(s)
   ...   (S.K // S.M).build(s)
   ...   (S.K / S.M).build(s)
   ...   (S.K % 4).build(s)
   ...   (D.i + D.j * 4).build(s)
   ...   s
   AffineExpr(s0 + s1)
   AffineExpr(s0 * s1)
   AffineExpr(s0 floordiv s1)
   AffineExpr(s0 ceildiv s1)
   AffineExpr(s0 mod 4)
   AffineExpr(d0 + d1 * 4)
   AffineBuildState<
     symbols={'K': 0, 'M': 1}
     dims={'i': 0, 'j': 1}>
 
 In the DSL, dimensions and symbols are name-uniqued instances of DimDef and
 SymbolDef. There are shortcut "expando" instances that will create a
 corresponding DimDef/SymbolDef upon accessing an attribute:
 
 Referencing a named dimension:
 
   >>> D.i
   Dim(i)
   >>> D.a is D.b
   False
   >>> D.a is D.a
   True
 
 Referencing a named symbol:
 
   >>> S.foobar
   Symbol(foobar)
   >>> S.a is S.b
   False
   >>> S.a is S.a
   True
 """
 
 from typing import Callable, Dict, Optional, Tuple, Union
 
 from mlir import ir as _ir
 
 __all__ = [
     "AffineBuildState",
     "AffineExprDef",
     "D",
     "DimDef",
     "S",
     "SymbolDef",
 ]
 
 # Type aliases.
 SymbolPosMap = Dict[str, int]
 
 
 class AffineBuildState:
   """Internal state for the AffineExprDef._create impls.
 
   Note that a "local" AffineBuildState can be created relative to a "global"
   AffineBuildState. In that case, any affine expressions built will inherit
   symbol and dim bindings from the global state and will update both as new
   ones are discovered. This allows for building expressions across contexts
   which share a common symbol and dim space.
   """
 
   def __init__(self,
                *,
                global_state: "AffineBuildState" = None,
                allow_new_symbols: bool = True,
                allow_new_dims: bool = True):
     if not global_state:
       self.all_symbols = dict()  # type: Dict[str, int]
       self.all_dims = dict()  # type: Dict[str, int]
     else:
       # Alias the global dict.
       self.all_symbols = global_state.all_symbols
       self.all_dims = global_state.all_dims
 
     # Map of symbols and dims in the current build.
     self.local_symbols = dict()  # type: Dict[str, int]
     self.local_dims = dict()  # type: Dict[str, int]
     self.allow_new_symbols = allow_new_symbols
     self.allow_new_dims = allow_new_dims
 
   def get_dim(self, dimname: str) -> int:
     """Gets the dim position given a name."""
     pos = self.all_dims.get(dimname)
     if pos is None:
       if not self.allow_new_dims:
         raise ValueError(
             f"New dimensions not allowed in the current affine expression: "
             f"Requested '{dimname}', Availble: {self.all_dims}")
       pos = len(self.all_dims)
       self.all_dims[dimname] = pos
     self.local_dims[dimname] = pos
     return pos
 
   def get_symbol(self, symname: str) -> int:
     """Geta a symbol position given a name."""
     pos = self.all_symbols.get(symname)
     if pos is None:
       if not self.allow_new_symbols:
         raise ValueError(
             f"New symbols not allowed in the current affine expression: "
             f"Requested '{symname}', Availble: {self.all_symbols}")
       pos = len(self.all_symbols)
       self.all_symbols[symname] = pos
     self.local_symbols[symname] = pos
     return pos
 
   @property
   def local_dim_count(self) -> int:
     return len(self.local_dims)
 
   @property
   def local_symbol_count(self) -> int:
     return len(self.local_symbols)
 
   @property
   def dim_count(self) -> int:
     return len(self.all_dims)
 
   @property
   def symbol_count(self) -> int:
     return len(self.all_symbols)
 
   def __repr__(self):
     lines = [f"AffineBuildState<"]
     lines.append(f"  symbols={self.local_symbols}")
     lines.append(f"  dims={self.local_dims}>")
     return "\n".join(lines)
 
 
 class AffineExprDef:
   """Base class for an affine expression being defined."""
 
   def build(self, state: Optional[AffineBuildState] = None) -> _ir.AffineExpr:
     """Builds the corresponding _ir.AffineExpr from the definitions.
     """
     state = AffineBuildState() if state is None else state
     expr = self._create(state)
     return expr
 
   def _create(self, state: AffineBuildState) -> _ir.AffineExpr:
     raise NotImplementedError()
 
   @staticmethod
   def coerce_from(py_value):
     if isinstance(py_value, int):
       return AffineConstantExpr(py_value)
     assert isinstance(py_value, AffineExprDef)
     return py_value
 
   def visit_affine_exprs(self, callback):
     """Visits all AffineExprDefs including self."""
     callback(self)
 
   def __add__(lhs, rhs):
     rhs = AffineExprDef.coerce_from(rhs)
     return AffineBinaryExprDef(_ir.AffineAddExpr, lhs, rhs)
 
   def __mul__(lhs, rhs):
     rhs = AffineExprDef.coerce_from(rhs)
     return AffineBinaryExprDef(_ir.AffineMulExpr, lhs, rhs)
 
   def __mod__(lhs, rhs):
     rhs = AffineExprDef.coerce_from(rhs)
     return AffineBinaryExprDef(_ir.AffineModExpr, lhs, rhs)
 
   def __floordiv__(lhs, rhs):
     rhs = AffineExprDef.coerce_from(rhs)
     return AffineBinaryExprDef(_ir.AffineFloorDivExpr, lhs, rhs)
 
   def __truediv__(lhs, rhs):
     # TODO: Not really a ceil div - taking liberties for the DSL.
     rhs = AffineExprDef.coerce_from(rhs)
     return AffineBinaryExprDef(_ir.AffineCeilDivExpr, lhs, rhs)
 
 
 class AffineConstantExpr(AffineExprDef):
   """An affine constant being defined."""
 
   def __init__(self, value: int):
     assert isinstance(value, int)
     self.value = value
 
   def _create(self, state: AffineBuildState) -> _ir.AffineExpr:
     return _ir.AffineConstantExpr.get(self.value)
 
   def __repr__(self):
     return f"Const({self.value})"
 
 
 class AffineBinaryExprDef(AffineExprDef):
   """An affine binary expression being defined."""
 
   def __init__(self, ir_ctor, lhs: AffineExprDef, rhs: AffineExprDef):
     self.ir_ctor = ir_ctor
     self.lhs = lhs
     self.rhs = rhs
 
   def _create(self, state: AffineBuildState) -> _ir.AffineExpr:
     return self.ir_ctor.get(self.lhs._create(state), self.rhs._create(state))
 
   def visit_affine_exprs(self, callback):
     """Visits all AffineExprDefs including self."""
     super().visit_affine_exprs(callback)
     self.lhs.visit_affine_exprs(callback)
     self.rhs.visit_affine_exprs(callback)
 
   def __repr__(self):
     return f"{self.ir_ctor.__name__}({repr(self.lhs)}, {repr(self.rhs)})"
 
 
 class DimDef(AffineExprDef):
   """Represents a named dimension.
 
   """
   ALL_DIMS = dict()  # type: Dict[str, "DimDef"]
   dimname: str
 
   def __new__(cls, dimname: str):
     existing = cls.ALL_DIMS.get(dimname)
     if existing is not None:
       return existing
     new = super().__new__(cls)
     new.dimname = dimname
     cls.ALL_DIMS[dimname] = new
     return new
 
   def __repr__(self):
     return f"Dim({self.dimname})"
 
   def _create(self, state: AffineBuildState) -> _ir.AffineExpr:
     pos = state.get_dim(self.dimname)
     return _ir.AffineDimExpr.get(position=pos)
 
   @classmethod
   def create_expando(cls):
     """Create an expando class that creates unique symbols based on attr access.
     """
 
     class ExpandoDims:
 
       def __getattr__(self, n):
         return cls(n)
 
     return ExpandoDims()
 
 
 class SymbolDef(AffineExprDef):
   """Represents a named symbol.
 
     >>> s1 = SymbolDef("s1")
     >>> s1
     Symbol(s1)
     >>> s2 = SymbolDef("s2")
     >>> s1 is s2
     False
     >>> s1 is SymbolDef("s1")
     True
   """
   ALL_SYMBOLS = dict()  # type: Dict[str, "SymbolDef"]
   symname: str
 
   def __new__(cls, symname: str):
     existing = cls.ALL_SYMBOLS.get(symname)
     if existing is not None:
       return existing
     new = super().__new__(cls)
     new.symname = symname
     cls.ALL_SYMBOLS[symname] = new
     return new
 
   def __repr__(self):
     return f"Symbol({self.symname})"
 
   def _create(self, state: AffineBuildState) -> _ir.AffineExpr:
     pos = state.get_symbol(self.symname)
     return _ir.AffineSymbolExpr.get(position=pos)
 
   @classmethod
   def create_expando(cls):
     """Create an expando class that creates unique symbols based on attr access.
     """
 
     class ExpandoSymbols:
 
       def __getattr__(self, n):
         return cls(n)
 
     return ExpandoSymbols()
 
 
 # Global accessor for on-demand dims and symbols.
 D = DimDef.create_expando()
 S = SymbolDef.create_expando()

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py

 #  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
 """Model classes representing a tensor comprehension.
 
 These classes model the language more at an AST level as evaluated. Reasoning
 about it typically involves processing this form into config objects that
 represent actual op definitions (i.e. YAML).
 """
 
 from typing import Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
 
 from mlir import ir as _ir
 
 from .affine import *
 from .scalar_expr import *
 from .types import *
 from .yaml_helper import *
 
 # Type aliases.
 AffineDimList = Dict[str, _ir.AffineExpr]
 
 
 class TensorExpression:
   """An expression that can appear on the RHS of a comprehension."""
 
   def to_scalar_expression(self) -> ScalarExpression:
     raise NotImplementedError()
 
   def visit_affine_exprs(self, callback):
     """Visits all affine expressions reachable by the expression."""
     pass
 
   def _get_all_dim_defs(self) -> Set[DimDef]:
     """Recursively gets all DimDef affine expressions that are referenced."""
     results = set()
 
     def visitor(affine_expr):
       if isinstance(affine_expr, DimDef):
         results.add(affine_expr)
 
     self.visit_affine_exprs(visitor)
     return results
 
   def collect_uses(self, uses: Set["TensorUse"]):
     """Collects all TensorUses reachable through this expression."""
     pass
 
   def __add__(self, rhs: "TensorExpression") -> "TensorExpression":
     return PrimFn.add(self, rhs)
 
   def __mul__(self, rhs) -> "TensorExpression":
     return PrimFn.mul(self, rhs)
 
   def __sub__(self, rhs) -> "TensorExpression":
     return PrimFn.sub(self, rhs)
 
   def __hash__(self):
     return hash(id(self))
 
 
 class TensorUse(TensorExpression):
   """A used tensor represented by its (tensor_name, indices).
 
   Note that forming a comprehension via direct assignment is performed through
   __setitem__ on the TensorDef level. However, performing a reduction with
   compound ops (+=, *=, etc) is done by doing a:
     TensorDef.__getitem__
     TensorUse.__iadd__
     TensorDef.__setitem__
   """
 
   def __init__(self, tensor_def: "TensorDef", indices: Sequence[AffineExprDef]):
     self.tensor_def = tensor_def
     self.indices = tuple(indices)
 
   def to_scalar_expression(self) -> ScalarExpression:
     assert self.tensor_def.tensor_name is not None
     return ScalarArg(self.tensor_def.tensor_name).expr()
 
   @property
   def tensor_name(self) -> str:
     n = self.tensor_def.tensor_name
     assert n is not None, "TensorDef not attached"
     return n
 
   def visit_affine_exprs(self, callback):
     for ind in self.indices:
       ind.visit_affine_exprs(callback)
 
   def collect_uses(self, uses: Set["TensorUse"]):
     uses.add(self)
 
   def __iadd__(self, rhs: TensorExpression) -> TensorExpression:
     return ReduceFn.add(*self._compute_reduce_dims(rhs))(rhs)
 
   def _compute_reduce_dims(self, rhs: TensorExpression) -> Set[DimDef]:
     """For implicit reductions, computes default reduction dims.
 
     Assumes that the rhs is the expression being reduced and self is being
     reduced into. Any indices referenced on the rhs and not in self are
     considered reduction dims and will be ordered as encountered on the rhs.
     """
     rhs_dims = rhs._get_all_dim_defs()
     lhs_dims = self._get_all_dim_defs()
     return rhs_dims - lhs_dims
 
   def __repr__(self):
     return f"{self.tensor_name}[{', '.join([repr(i) for i in self.indices])}]"
 
 
 class TensorDef:
   """Bookkeeping of a single registered tensor, held in dict by name."""
 
   def __init__(self,
                type_var: TypeVar,
                *shape: AffineExprDef,
                indexing_map: Optional[_ir.AffineMap] = None,
                output: bool = False):
     if not isinstance(type_var, TypeVar):
       raise ValueError(f"TensorDef requires a TypeVar. Got: {repr(type_var)}")
     self.owner = None  # type: Optional["LinalgOpDef"]
     self.type_var = type_var
     self.shape = shape
     self.indexing_map = indexing_map
     self.output = output
     self.tensor_name = None  # type: Optional[str]
     self.registered_index = -1  # type: int
 
   @property
   def rank(self) -> int:
     """The rank of the tensor."""
     return len(self.shape)
 
   def attach(self, index: int, tensor_name: str, owner: "LinalgOpDef"):
     if self.owner:
       raise ValueError(f"TensorDef already registered with op: {self}")
     self.registered_index = index
     self.tensor_name = tensor_name
     self.owner = owner
 
   def __getitem__(self, dims) -> TensorUse:
     assert self.owner, "TensorDef is not attached to an op"
     state = AffineBuildState(global_state=self.owner._affine_state,
                              allow_new_symbols=False)
     if not isinstance(dims, tuple):
       dims = (dims,)  # Handle single subscript case.
     # Special case: (None) is a 0d-scalar use.
     if dims == (None,):
       dims = ()
 
     exprs = []
     for expr_def in dims:
       if not isinstance(expr_def, AffineExprDef):
         raise KeyError(
             "A TensorDef can only be subscripted by a tuple of affine dims")
       exprs.append(expr_def)
     return TensorUse(self, exprs)
 
   def __setitem__(self, dims, value):
     """Creates a new 1:1 comprehension by binding this tensor to an expression.
 
     Note that due to the way assignment works in Python, we have to capture
     direct assignment as a setitem on the TensorDef.
     """
     if not isinstance(value, TensorExpression):
       raise ValueError(f"Only TensorExpressions can be assigned to TensorDefs. "
                        f"Got: {repr(value)}")
     use = self[dims]
     comp = Comprehension((use, value))
     self.owner.comprehensions.append(comp)
 
   def __hash__(self):
     return hash(id(self))
 
   def __repr__(self):
     output = "OUTPUT " if self.output else ""
     return (f"{self.tensor_name}:TensorDef({output}{repr(self.type_var)}, "
             f"shape={self.shape})")
 
 
 class Comprehension:
   """Represents a single comprehension."""
 
   def __init__(self, *bindings: Tuple[TensorUse, TensorExpression]):
     self.definitions = list()  # List[TensorUse]
     self.values = list()  # List[TensorExpression]
 
     # Find the lhs to reduction rhs.
     for assign, value in bindings:
       if isinstance(value, ReduceApply):
         if value.lhs:
           raise ValueError(f"Reduction expression already assigns: {value}")
         value.lhs = assign
       self.definitions.append(assign)
       self.values.append(value)
 
   @property
   def all_reduction_dims(self) -> Set[Tuple[DimDef, ...]]:
     """Gets the reduction dims for the comprehension or None."""
     result = set()
     for use in self.values:
       if isinstance(use, ReduceApply):
         result.add(use.reduce.reduce_dims)
       else:
         result.add(tuple())
     return result
 
   def __repr__(self):
     if len(self.definitions) > 1:
       defs_repr = f"({', '.join(repr(d) for d in self.definitions)})"
       values_repr = f"({', '.join(repr(v) for v in self.values)})"
     else:
       defs_repr = f"{repr(self.definitions[0])}"
       values_repr = f"{repr(self.values[0])}"
 
     return f"{defs_repr} = {values_repr}"
 
 
 class PrimFnType:
   """Primitive operations."""
 
   def __init__(self, prim_name: str):
     self.prim_name = prim_name
 
   def __call__(self, *args):
     return PrimApply(self, args)
 
   def reduce(self, *reduce_dims: DimDef):
     """Shortcut to create a Reduce operation from this primitive."""
     return ReduceFnType(self, *reduce_dims)
 
   def __repr__(self):
     return f"{self.prim_name}"
 
 
 class PrimFn:
   add = PrimFnType("add")
   exp = PrimFnType("exp")
   log = PrimFnType("log")
   mul = PrimFnType("mul")
   max = PrimFnType("max")
   sub = PrimFnType("sub")
 
 
 class ReduceFnType:
   """A reduction operator that reduces into its LHS from its RHS."""
 
   def __init__(self, operator: PrimFnType, *reduce_dims: DimDef):
     """Initializes the ReduceFn with a primitive function and dims."""
     if not isinstance(operator, PrimFnType):
       raise ValueError(f"Reduce expected a Prim operator. Got: {operator}")
     self.operator = operator
     self.reduce_dims = tuple(reduce_dims)
 
   def __call__(self, *args: TensorExpression):
     return ReduceApply(self, args)
 
   def __repr__(self):
     return (f"reduce_{self.operator.prim_name}"
             f"({', '.join(repr(d) for d in self.reduce_dims)})")
 
 
 class ReduceFn:
   add = PrimFn.add.reduce
   mul = PrimFn.mul.reduce
   max = PrimFn.max.reduce
 
 
 class PrimApply(TensorExpression):
   """Application of a primitive."""
 
   def __init__(self, prim: PrimFnType, args: Sequence[TensorExpression]):
     self.prim = prim
     self.args = tuple(args)
 
   def to_scalar_expression(self) -> ScalarExpression:
     return ScalarApplyFn(self.prim.prim_name,
                          *[arg.to_scalar_expression() for arg in self.args
                           ]).expr()
 
   def visit_affine_exprs(self, callback):
     for arg in self.args:
       arg.visit_affine_exprs(callback)
 
   def collect_uses(self, uses: Set["TensorUse"]):
     for arg in self.args:
       arg.collect_uses(uses)
 
   def __repr__(self):
     return f"{repr(self.prim)}({', '.join(repr(a) for a in self.args)})"
 
 
 class cast(TensorExpression):
   """Casts the element type to a type (typically symbolic TypeVar)."""
 
   def __init__(self, to_type: TypeVar, operand: TensorExpression):
     self.to_type = to_type
     self.operand = operand
 
   def to_scalar_expression(self) -> ScalarExpression:
     return ScalarSymbolicCast(self.to_type,
                               self.operand.to_scalar_expression()).expr()
 
   def visit_affine_exprs(self, callback):
     self.operand.visit_affine_exprs(callback)
 
   def collect_uses(self, uses: Set["TensorUse"]):
     self.operand.collect_uses(uses)
 
   def __repr__(self):
     return f"cast({self.to_type}, {repr(self.operand)})"
 
 
 class ReduceApply(TensorExpression):
   """Application of a reduction.
 
   This captures the lhs separately (initial value) separately from the rhs.
   """
 
   def __init__(self, reduce: ReduceFnType, args: Sequence[TensorExpression]):
     self.reduce = reduce
     self.lhs = None  # type: Optional[TensorUse]
     self.args = tuple(args)
 
   def to_scalar_expression(self) -> ScalarExpression:
     if self.lhs is None:
       raise ValueError(f"Cannot scalarize a ReduceApply that has not been "
                        f"bound to its lhs: {self}")
     full_args = [self.lhs.to_scalar_expression()
                 ] + [arg.to_scalar_expression() for arg in self.args]
     return ScalarApplyFn(self.reduce.operator.prim_name, *full_args).expr()
 
   def visit_affine_exprs(self, callback):
     for ind in self.reduce.reduce_dims:
       ind.visit_affine_exprs(callback)
     for arg in self.args:
       arg.visit_affine_exprs(callback)
 
   def collect_uses(self, uses: Set["TensorUse"]):
     for arg in self.args:
       arg.collect_uses(uses)
 
   def __repr__(self):
     return f"{repr(self.reduce)}({', '.join(repr(a) for a in self.args)})"
 
 
 class OpInterfaceDef:
   """An interface that an op implements."""
 
   def __init__(self, cpp_name: str):
     self.cpp_name = cpp_name
 
 
 ContractionOpInterface = OpInterfaceDef("LinalgContractionOpInterface")
 
 
 class OpMetadataDef(YAMLObject):
   """Metadata about the op (generally not behavior impacting)."""
   yaml_tag = "!LinalgOpMetadata"
 
   def __init__(self, name: str, cpp_class_name: Optional[str], doc: Optional[str]):
     self.name = name
     self.cpp_class_name = cpp_class_name if cpp_class_name is not None else name
     self.doc = doc
     self.implements = []  # type: List[OpInterfaceDef]
 
   def to_yaml_custom_dict(self):
     d = dict(
         name=self.name,
         cpp_class_name=self.cpp_class_name,
         doc=self.doc,
     )
     if self.implements:
       d["implements"] = [intr.cpp_name for intr in self.implements]
     return d
 
 
 class LinalgOpDef:
   """Definition of a linalg op."""
 
   def __init__(self,
                name: str,
                cpp_class_name: Optional[str] = None,
                doc: Optional[str] = None):
     self.metadata = OpMetadataDef(name=name, cpp_class_name=cpp_class_name, doc=doc)
     self.registered_tensors = dict()  # type: Dict[str, TensorDef]
     self.comprehensions = list()  # type: List[Comprehension]
     self._affine_state = AffineBuildState()
 
   @property
   def inputs(self) -> Sequence[TensorDef]:
     return [t for t in self.registered_tensors.values() if not t.output]
 
   @property
   def outputs(self) -> Sequence[TensorDef]:
     return [t for t in self.registered_tensors.values() if t.output]
 
   def add_tensor(self, tensor_name: str, tensor: TensorDef):
     """Registers a tensor."""
     if tensor_name in self.registered_tensors:
       raise ValueError(f"Tensor {tensor_name} is already registered "
                        f"to {self.registered_tensors['tensor_name']}")
     tensor.attach(len(self.registered_tensors), tensor_name, self)
     self.registered_tensors[tensor_name] = tensor
 
   def tensor(self, name):
     """Gets a registered tensor by name."""
     try:
       return self.registered_tensors[name]
     except KeyError:
       raise KeyError(f"Tensor {name} is not registered")
 
   def __repr__(self):
     lines = [
         f"LinalgOpDef({self.metadata.name} -> {self.metadata.cpp_class_name},"
     ]
     for name, tensor in self.registered_tensors.items():
       lines.append(f"  {tensor}")
     if self.comprehensions:
       lines[-1] += " {"
       for comprehension in self.comprehensions:
         lines.append(f"    {comprehension}")
       lines.append("}")
     return "\n".join(lines)

mlir/python/mlir/dialects/linalg/opdsl/lang/config.py

 #  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
 """Represents configured ops as emitted for code generation.
 
 Classes in this module generally are directly serializable to YAML for use
 by the code generator.
 
 TODO: These should just be dumb containers or serialization code but they
 currently encode too many details of how the language is interpreted. Move this
 to helpers on the comprehension objects themselves.
 """
 
 from typing import Any, Dict, Optional
 
 from mlir import ir as _ir
 
 from .comprehension import *
 from .yaml_helper import *
 
 __all__ = [
     "LinalgStructuredOpConfig",
     "LinalgOpConfig",
     "TensorDefConfig",
 ]
 
 
 def _serialize_affine_map(affine_map: _ir.AffineMap) -> str:
   with affine_map.context:
     # Affine map printing/parsing is via an AffineMap attr.
     attr = _ir.AffineMapAttr.get(affine_map)
     return str(attr)
 
 
 class TensorUseConfig:
   """Wrapper around a TensorUse with additional context-bound state."""
 
   def __init__(self, tensor_use: TensorUse, indexing_map: _ir.AffineMap):
     self.tensor_use = tensor_use
     self.indexing_map = indexing_map
 
   def __repr__(self):
     return f"Use({self.tensor_use}, indexing_map={self.indexing_map})"
 
 
 class TensorDefConfig(YAMLObject):
   """Wrapper around a TensorDef with additional context-bound state."""
   yaml_tag = "LinalgTensorDef"
 
   def __init__(self, tensor_def: TensorDef, shape_map: _ir.AffineMap):
     self.tensor_def = tensor_def
     self.shape_map = shape_map
     self.indexing_map = None  # type: Optional[_ir.AffineMap]
 
   @property
   def usage(self) -> str:
     if self.tensor_def.output:
       return "output"
     else:
       return "input"
 
   def to_yaml_custom_dict(self):
     return dict(
         name=self.tensor_def.tensor_name,
         usage=self.usage,
         shape=_serialize_affine_map(self.shape_map),
         element_type_var=self.tensor_def.type_var.name,
     )
 
   def __repr__(self):
     return f"Def({self.tensor_def}, shape_map={self.shape_map}, indexing_map={self.indexing_map})"
 
 
 class LinalgIndexingMapsConfig(YAMLObject):
   """Abstracts the style of indexing maps that the op exports.
 
   Presently only static (tied to the op name) indexing maps are supported. In
   the future, it is expected that we will have additional variants:
     - Dynamic based on attributes
     - Dynamic based on operands
   Each is expected to require a different variant of specification.
   """
   yaml_tag = "!LinalgIndexingMapsConfig"
 
   def __init__(self,
                static_indexing_maps: Optional[Sequence[_ir.AffineMap]] = None):
     self.static_indexing_maps = static_indexing_maps
 
   def to_yaml_custom_dict(self):
     if self.static_indexing_maps is not None:
       return dict(static_indexing_maps=[
           _serialize_affine_map(m) for m in self.static_indexing_maps
       ])
     raise ValueError(
         f"LinalgIndexingMapsConfig must have one type of indexing map"
         f"(got none)")
 
 
 class LinalgStructuredOpConfig(YAMLObject):
   """Configuration for metadata sufficient to construct a linalg single
   contraction named op."""
 
   yaml_tag = "!LinalgStructuredOpConfig"
 
   def __init__(self,
                comprehension: Comprehension,
                context: Optional[_ir.Context] = None):
     self.context = context if context is not None else _ir.Context()
     self.affine_state = AffineBuildState()
     self.writes = list()  # type: List[Tuple[TensorUse, TensorExpression]]
     self.tensor_args = dict()  # type: Dict[TensorDef, TensorDefConfig]
     self.uses = dict()  # type: Dict[TensorUse, TensorUseConfig]
 
     # Compute the ordered set of writes.
     collected_uses = set()
     for write_use, read_use in zip(comprehension.definitions,
                                    comprehension.values):
       self.writes.append((write_use, read_use))
 
     for write_use, read_use in self.writes:
       collected_uses.add(write_use)
       read_use.collect_uses(collected_uses)
 
     # Need to add all definitions before uses, so process twice.
     for use in collected_uses:
       self.add_tensor_arg(use.tensor_def)
     for use in collected_uses:
       self.add_use(use)
 
     # Now normalize all defs and uses indexing maps now that full count of
     # dims and symbols are known.
     for cuse in self.uses.values():
       cuse.indexing_map = self._normalize_affine_map(cuse.indexing_map)
     for cdef in self.tensor_args.values():
       cdef.shape_map = self._normalize_affine_map(cdef.shape_map,
                                                   with_dims=False)
 
     # Now for each write use, propagate the indexing maps from the use to the
     # tensor, ensuring that there are not conflicts.
     for write_use, _ in self.writes:
       write_tensor_def = self.tensor_args[write_use.tensor_def]
       if write_tensor_def.indexing_map:
         raise ValueError(
             f"Unexpected multi-write to a single tensor: {write_tensor_def}")
       write_tensor_def.indexing_map = self.uses[write_use].indexing_map
 
     # For each read use, propagate the indexing maps from the use to the
     # tensor, ensuring that there are not conflicts.
     for _, read_expr in self.writes:
       read_uses = set()  # type: Set[TensorUse]
       read_expr.collect_uses(read_uses)
       for read_use in read_uses:
         read_tensor_def = self.tensor_args[read_use.tensor_def]
         if (read_tensor_def.indexing_map and
             read_tensor_def.indexing_map != self.uses[read_use].indexing_map):
           raise ValueError(
               f"Unexpected multi-read of a tensor with different accesses:"
               f"{read_tensor_def} vs {read_use}")
         read_tensor_def.indexing_map = self.uses[read_use].indexing_map
 
     # Sanity check that all defs have an indexing map.
     assert all(d.indexing_map for d in self.tensor_args.values()), (
         f"Missing indexing map on TensorDef: {self.tensor_args}")
 
     # Collect reduction dims and ensure all the same.
     all_reduction_dims = set(comprehension.all_reduction_dims)
     if len(all_reduction_dims) != 1:
       raise ValueError(
           f"All writes within a generic must have the same reduction "
           f"dims. Got: {all_reduction_dims}")
     self.reduction_dims = next(iter(all_reduction_dims))
 
     # Generate the scalar assignments (used to build a body).
     self.assignments = [
         ScalarAssign(write_use.tensor_name, read_expr.to_scalar_expression())
         for write_use, read_expr in self.writes
     ]
 
   @property
   def ordered_tensor_args(self) -> Sequence[TensorDefConfig]:
     return sorted(self.tensor_args.values(),
                   key=lambda tdc: tdc.tensor_def.registered_index)
 
   @property
   def ordered_tensor_uses(self) -> Sequence[TensorUseConfig]:
     return sorted(self.uses.values(),
                   key=lambda tuc: tuc.tensor_use.tensor_def.registered_index)
 
   @property
   def ordered_dims(self) -> Sequence[Tuple[str, int]]:
     """Gets the ordered list of dim bindings (symbolic name, position).
 
     TODO: The original parser relies on parse ordering to arrive at the
     iterator types, but that ordering is not defined on the Python side, so
     this may be ambiguous.
     """
     return list(self.affine_state.all_dims.items())
 
   @property
   def indexing_maps(self) -> Sequence[_ir.AffineMap]:
     return [use.indexing_map for use in self.ordered_tensor_uses]
 
   @property
   def iterator_types(self) -> Sequence[str]:
 
     def get_type(symbolic_name, position):
       for reduction_dim_expr in self.reduction_dims:
         if reduction_dim_expr.dimname == symbolic_name:
           return "reduction"
       return "parallel"
 
     return [get_type(*dim) for dim in self.ordered_dims]
 
   def add_tensor_arg(self, tensor_def: TensorDef):
     if tensor_def in self.tensor_args:
       return
     with self.context:
       local_state = AffineBuildState(global_state=self.affine_state,
                                      allow_new_dims=False)
       exprs = []
       for expr in tensor_def.shape:
         exprs.append(expr.build(state=local_state))
       assert local_state.local_dim_count == 0
       indexing_map = _ir.AffineMap.get(dim_count=0,
                                        symbol_count=local_state.symbol_count,
                                        exprs=exprs)
 
       def_config = TensorDefConfig(tensor_def, indexing_map)
       self.tensor_args[tensor_def] = def_config
 
   def add_use(self, tensor_use: TensorUse):
     if tensor_use in self.uses:
       return
     with self.context:
       local_state = AffineBuildState(global_state=self.affine_state,
                                      allow_new_symbols=False)
       exprs = []
       for expr in tensor_use.indices:
         exprs.append(expr.build(state=local_state))
       assert local_state.local_symbol_count == 0
       indexing_map = _ir.AffineMap.get(dim_count=local_state.dim_count,
                                        symbol_count=local_state.symbol_count,
                                        exprs=exprs)
 
       use_config = TensorUseConfig(tensor_use, indexing_map)
       self.uses[tensor_use] = use_config
 
   def _normalize_affine_map(self,
                             affine_map: _ir.AffineMap,
                             with_dims: bool = True) -> _ir.AffineMap:
     """Normalizes an indexing map to have the max known symbols and dims."""
     with self.context:
       return _ir.AffineMap.get(
           dim_count=self.affine_state.dim_count if with_dims else 0,
           symbol_count=self.affine_state.symbol_count,
           exprs=list(affine_map.results))
 
   def to_yaml_custom_dict(self):
     self_dict = dict(
         args=self.ordered_tensor_args,
         # TODO: Refactor the hierarchy internally when supporting more
         # than static (preserving this serialized form).
         indexing_maps=LinalgIndexingMapsConfig(
             static_indexing_maps=self.indexing_maps),
         iterator_types=self.iterator_types,
         assignments=self.assignments,
     )
     return self_dict
 
   def __repr__(self):
     lines = [f"LinalgGenericOpConfig(reduction_dims={self.reduction_dims},"]
     lines.append("tensor_args=[")
     for def_config in self.ordered_tensor_args:
       lines.append(f"  {repr(def_config)}")
     lines.append("], indexing_maps=[")
     for m in self.indexing_maps:
       lines.append(f"  {repr(m)}")
     lines.append(f"], iterator_types=[")
     for t in self.iterator_types:
       lines.append(f"  {t}")
     lines.append("])")
     return "\n".join(lines)
 
 
 class LinalgOpConfig(YAMLObject):
   """Container for any supported linalg op type.
 
   This includes the concrete type by name for ease of parsing by systems
   that ignore tags.
   """
   yaml_tag = "!LinalgOpConfig"
 
   def __init__(self,
                metadata: OpMetadataDef,
                *,
                structured_op: Optional[LinalgStructuredOpConfig] = None):
     self.metadata = metadata
     self.structured_op = structured_op
 
   def to_yaml_custom_dict(self):
     self_dict = dict(metadata=self.metadata,)
     if self.structured_op:
       self_dict["structured_op"] = self.structured_op
     return self_dict
 
   @staticmethod
   def from_linalg_op_def(
       tc_op_def: LinalgOpDef,
       context: Optional[_ir.Context] = None) -> Sequence["LinalgOpConfig"]:
     """Expands a LinalgOpDef into corresponding Linalg configured ops."""
     # TODO: Many LinalgOpDef patterns need to expand to multiple generics.
     assert len(
         tc_op_def.comprehensions) == 1, "Only one comprehension supported"
     return [
         LinalgOpConfig(tc_op_def.metadata,
                        structured_op=LinalgStructuredOpConfig(
                            tc_op_def.comprehensions[0], context)),
     ]
 
   def __repr__(self):
     return (f"LinalgOpConfig(metadata={self.metadata},\n"
             f"structured_op={self.structured_op})")

mlir/python/mlir/dialects/linalg/opdsl/lang/dsl.py

 #  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 typing import Dict, List
 
 from contextlib import contextmanager
 import functools
 import inspect
 import threading
 
 from mlir import ir
 from .comprehension import *
 from .config import *
 from .emitter import *
 
 _CONTEXT = threading.local()
 
 
 @contextmanager
 def bind_op_def(model: LinalgOpDef):
   if hasattr(_CONTEXT, "current_op_def"):
     raise ValueError("Cannot recursively define an operation")
   _CONTEXT.current_op_def = model
   try:
     yield model
   finally:
     del _CONTEXT.current_op_def
 
 
 def current_op_def() -> LinalgOpDef:
   try:
     return _CONTEXT.current_op_def
   except AttributeError:
     raise ValueError(
         "Attempt to access the current op definition being defined "
         "but none is set. Did you mean to call this in an op definition?")
 
 
 class DefinedOpCallable:
   """Callable that wraps any defined op function."""
 
   def __init__(self, op_name: str, model: LinalgOpDef):
     self.op_name = op_name
     self.model = model
 
   def __call__(self, *args, emit_generic: bool = False, **kwargs):
     """Emits the corresponding op definition as IR.
 
     Most arguments are passed through to the underlying emitter. The following
     are interpreted here:
       emit_generic: Emits a generic form as appropriate (default True). If
         False, a named form is emitted (which must have been built in to the
         compiler).
     """
     op_configs = LinalgOpConfig.from_linalg_op_def(self.model,
                                                    context=ir.Context.current)
 
     if len(op_configs) != 1:
       # TODO: Support composite ops.
       raise NotImplementedError(
           f"Emission of composite linalg ops not supported: {op_configs}")
 
     ctx = ir.Context.current
     linalgDialect = ctx.get_dialect_descriptor("linalg")
     fully_qualified_name = 'linalg.' + self.op_name
     emit_generic = (emit_generic or not ctx.is_registered_operation(fully_qualified_name))
 
     op_config = op_configs[0]
     if op_config.structured_op:
       if emit_generic:
         return emit_generic_structured_op(op_config.structured_op, *args,
                                           **kwargs)
       else:
         return emit_named_structured_op(
           op_config.structured_op, self.op_name,
           self.model.metadata.cpp_class_name, *args, **kwargs)
 
     raise NotImplementedError(
         f"Emission of linalg op type not supported: {op_config}")
 
 
 def linalg_structured_op(dsl_func=None,
                          *,
                          op_name=None,
                          op_class_name=None) -> DefinedOpCallable:
   if dsl_func is None:
     # Curry the keyword args in for delayed application.
     return functools.partial(tc_def_op,
                              op_name=op_name,
                              op_class_name=op_class_name)
   # Determine default names by introspecting the function.
   if op_name is None:
     op_name = dsl_func.__name__
   if op_class_name is None:
     # Camel case it.
     op_class_name = f"{''.join(x.title() for x in op_name.split('_'))}Op"
 
   tc_model = LinalgOpDef(name=op_name,
                          cpp_class_name=op_class_name,
                          doc=inspect.getdoc(dsl_func))
 
   # Extract arguments and TensorDefs from the signature.
   dsl_func_args = list()
   sig = inspect.signature(dsl_func)
   for param_name, param in sig.parameters.items():
     param_default = param.default
     if not isinstance(param_default, TensorDef):
       raise ValueError(f"@tc_def_op function parameters must be defaulted as "
                        f"TensorDef(...): Found {param_name}: {param_default}")
     dsl_func_args.append(param_default)
     tc_model.add_tensor(param_name, param_default)
 
   # Invoke the DSL func to finish populating the model.
   with bind_op_def(tc_model):
     dsl_func(*dsl_func_args)
 
   # TODO: The returned callable should be an IR emitter but that is not
   # upstreamed yet.
   return DefinedOpCallable(op_name, tc_model)
 
 
 def implements(*interfaces: OpInterfaceDef):
   current_op_def().metadata.implements.extend(interfaces)

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

 #  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 typing import Dict, Sequence
 
 from mlir.ir import *
 from mlir.dialects import linalg
 from mlir.dialects import std
 # TODO: resolve name collision for Linalg functionality that is injected inside
 # the _mlir.dialects.linalg directly via pybind.
 from _mlir.dialects.linalg import fill_builtin_region
 
 from .scalar_expr import *
 from .config import *
 
 __all__ = [
     "emit_generic_structured_op",
     "emit_named_structured_op",
 ]
 
 def isa(cls : Type, ty : Type):
   try:
     cls(ty)
     return True
   except ValueError:
     return False
 
 def prepare_common_structured_op(op_config: LinalgStructuredOpConfig,
                                  *ins: Value,
                                  outs: Value):
   all_arg_defs = op_config.ordered_tensor_args
   in_arg_defs = [arg for arg in all_arg_defs if arg.usage == "input"]
   out_arg_defs = [arg for arg in all_arg_defs if arg.usage == "output"]
 
   # Arity validation.
   if len(ins) != len(in_arg_defs):
     raise ValueError(f"Expected {len(in_arg_defs)} inputs but got "
                      f"{len(ins)} for {op_config}")
   if outs and len(outs) != len(out_arg_defs):
     raise ValueError(f"Expected {len(out_arg_defs)} outputs but got "
                      f"{len(outs)} for {op_config}")
 
   outs, out_types = _infer_structured_outs(op_config, in_arg_defs, ins,
                                            out_arg_defs, outs)
 
   result_types = [t for t in out_types if isa(RankedTensorType, t)]
 
   # Extract type vars for input/output based types.
   type_mapping = dict()  # type: Dict[str, Type]
   for arg_def, arg_element_type in zip(
       in_arg_defs + out_arg_defs,
       _get_shaped_element_types_from_values(*ins, *outs)):
     tv_name = arg_def.tensor_def.type_var.name
     type_mapping[tv_name] = arg_element_type
 
   # Emit the generic op.
   # TODO: Support emission of pure memref form.
   indexing_maps_attr = ArrayAttr.get(
       [AffineMapAttr.get(am)
        # TODO: linalg verification does not currently allow symbols.
        # Compress them for now.
        for am in AffineMap.compress_unused_symbols(op_config.indexing_maps, Context.current)])
   iterator_types_attr = ArrayAttr.get(
       [StringAttr.get(s) for s in op_config.iterator_types])
   sparse_attr = ArrayAttr.get(
       [BoolAttr.get(False) for s in list(ins) + list(outs) if isa(RankedTensorType, s.type)])
   if len(sparse_attr) == 0:
     sparse_attr = None
 
   return (all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types,
           type_mapping, indexing_maps_attr, iterator_types_attr, sparse_attr)
 
 
 def emit_generic_structured_op(op_config: LinalgStructuredOpConfig,
                                *ins: Value,
                                outs: Value = ()):
   all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, \
   type_mapping, indexing_maps_attr, iterator_types_attr, sparse_attr =   \
      prepare_common_structured_op(op_config, *ins, outs = outs)
 
   generic_op = linalg.GenericOp(
       result_tensors=result_types,
       inputs=ins,
       outputs=outs,
       indexing_maps=indexing_maps_attr,
       iterator_types=iterator_types_attr,
       doc=None,  # TODO: Make optional.
       library_call=None,  # TODO: Make optional.
       sparse=sparse_attr)  # TODO: Make optional.
 
   # Construct the body.
   block_arg_names = _get_tensor_def_names(*in_arg_defs, *out_arg_defs)
   block_arg_types = _get_shaped_element_types_from_values(*ins, *outs)
   block = generic_op.regions[0].blocks.append(*block_arg_types)
   block_arg_mapping = dict(zip(block_arg_names, block.arguments))
   with InsertionPoint(block):
     body_builder = _BodyBuilder(type_mapping, block_arg_mapping)
     for assignment in op_config.assignments:
       body_builder.assign(assignment)
     body_builder.yield_outputs(*_get_tensor_def_names(*out_arg_defs))
 
   if len(result_types) == 1:
     return generic_op.result
   else:
     return generic_op.results
 
 
 def emit_named_structured_op(op_config: LinalgStructuredOpConfig,
                              op_name: str,
                              op_class_name: str,
                              *ins: Value,
                              outs: Value = ()):
   all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, \
   type_mapping, indexing_maps_attr, iterator_types_attr, sparse_attr =   \
      prepare_common_structured_op(op_config, *ins, outs = outs)
 
   # If we get here, there must exist a builtin class `op_class_name`.
   ctx = Context.current
   fully_qualified_name = 'linalg.' + op_name
   if (not ctx.is_registered_operation(fully_qualified_name) or
       not op_class_name in linalg.__dict__.keys()):
     raise NotImplementedError(
         f"Unknown named op_name / op_class_name: {op_name} / {op_class_name}")
 
   named_op = getattr(linalg, op_class_name)(ins, outs, result_types)
   linalgDialect = ctx.get_dialect_descriptor("linalg")
   fill_builtin_region(linalgDialect, named_op.operation)
   # Note: mlir-linalg-ods-yaml-gen.cpp uses a special linalg.memoized_indexing_maps
   # attribute that the non-yaml path does not. The non-yaml path hardcodes the 
   # indexing_maps in C++ directly.
   named_op.operation.attributes["linalg.memoized_indexing_maps"] = indexing_maps_attr
   # iterator_types are hardcoded in C++ both in the yaml and non-yaml path.
 
   if len(result_types) == 1:
     return named_op.result
   else:
     return named_op.results
 
 
 class _BodyBuilder:
   """Constructs a structured op body by evaluating assignments."""
 
   def __init__(self, type_mapping: Dict[str, Type],
                block_arg_mapping: Dict[str, Value]):
     self.type_mapping = type_mapping
     self.block_arg_mapping = block_arg_mapping
     self.yield_mapping = dict()  # type: Dict[str, Value]
 
   def assign(self, assignment: ScalarAssign):
     if assignment.arg in self.yield_mapping:
       raise ValueError(
           f"Multiple assignments to the same argument are forbidden: "
           f"{assignment}")
     self.yield_mapping[assignment.arg] = self.expression(assignment.value)
 
   def expression(self, expr: ScalarExpression) -> Value:
     if expr.scalar_arg:
       try:
         return self.block_arg_mapping[expr.scalar_arg.arg]
       except KeyError:
         raise ValueError(f"Argument {expr.scalar_arg.arg} is not bound for "
                          f"this structured op.")
     elif expr.scalar_apply:
       try:
         fn = getattr(self, f"_eval_{expr.scalar_apply.fn_name}")
       except AttributeError:
         raise ValueError(
             f"Function '{expr.scalar_apply.fn_name}' is not a known "
             "scalar body function")
       operand_values = [
           self.expression(operand) for operand in expr.scalar_apply.operands
       ]
       return fn(*operand_values)
     elif expr.symbolic_cast:
       operand_value = self.expression(expr.symbolic_cast.operand)
       return self.cast(expr.symbolic_cast.to_type.name, operand_value)
     raise NotImplementedError(f"Unimplemented scalar body expression: {expr}")
 
   def cast(self, type_var_name: str, operand: Value) -> Value:
     try:
       to_type = self.type_mapping[type_var_name]
     except KeyError:
       raise ValueError(f"Unbound type variable '{type_var_name}' ("
                        f"expected one of {self.type_mappings.keys()}")
     if operand.type == to_type:
       return operand
     if _is_integer_type(to_type):
       return self._cast_to_integer(to_type, operand)
     elif _is_floating_point_type(to_type):
       return self._cast_to_floating_point(to_type, operand)
 
     raise ValueError(f"Unable to cast body expression from {operand.type} to "
                      f"{to_type}")
 
   def _cast_to_integer(self, to_type: Type, operand: Value) -> Value:
     to_width = IntegerType(to_type).width
     operand_type = operand.type
     if _is_floating_point_type(operand_type):
       return std.FPToSIOp(to_type, operand).result
     # Assume integer.
     from_width = IntegerType(operand_type).width
     if to_width > from_width:
       return std.SignExtendIOp(to_type, operand).result
     elif to_width < from_width:
       return std.TruncateIOp(to_type, operand).result
     raise ValueError(f"Unable to cast body expression from {operand_type} to "
                      f"{to_type}")
 
   def _cast_to_floating_point(self, to_type: Type, operand: Value) -> Value:
     operand_type = operand.type
     if _is_integer_type(operand_type):
       return std.SIToFPOp(to_type, operand).result
     # Assume FloatType.
     to_width = _get_floating_point_width(to_type)
     from_width = _get_floating_point_width(operand_type)
     if to_width > from_width:
       return std.FPExtOp(to_type, operand).result
     elif to_width < from_width:
       return std.FPTruncOp(to_type, operand).result
     raise ValueError(f"Unable to cast body expression from {operand_type} to "
                      f"{to_type}")
 
   def yield_outputs(self, *output_names: str):
     output_values = []
     for n in output_names:
       try:
         output_values.append(self.yield_mapping[n])
       except KeyError:
         raise ValueError(f"Body assignments do not assign all outputs: "
                          f"missing '{n}'")
     linalg.YieldOp(output_values)
 
   def _eval_add(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return std.AddFOp(lhs.type, lhs, rhs).result
     if _is_integer_type(lhs.type):
       return std.AddIOp(lhs.type, lhs, rhs).result
     raise NotImplementedError("Unsupported 'add' operand: {lhs}")
 
   def _eval_mul(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
       return std.MulFOp(lhs.type, lhs, rhs).result
     if _is_integer_type(lhs.type):
       return std.MulIOp(lhs.type, lhs, rhs).result
     raise NotImplementedError("Unsupported 'mul' operand: {lhs}")
 
 
 def _infer_structured_outs(op_config: LinalgStructuredOpConfig,
                            in_arg_defs: Sequence[TensorDefConfig],
                            ins: Sequence[Value],
                            out_arg_defs: Sequence[TensorDefConfig],
                            outs: Sequence[Value]):
   """Infers implicit outs and output types.
 
   Respects existing contents of outs if not empty.
 
   Returns:
     normalized outs, output types
   """
   # If outs were explicitly provided, we accept them verbatim.
   if outs:
     return outs, [out.type for out in outs]
 
   raise NotImplementedError(f"Output tensor inference not yet supported for "
                             "structured ops")
 
 
 def _get_shaped_element_types_from_values(*values: Value) -> Sequence[Type]:
   types = []
   for v in values:
     try:
       t = ShapedType(v.type)
     except Exception as e:
       raise ValueError(f"Expected ShapedType but got {v}") from e
     types.append(t.element_type)
   return types
 
 
 def _get_tensor_def_names(
     *tensor_def_configs: TensorDefConfig) -> Sequence[str]:
   return [tdc.tensor_def.tensor_name for tdc in tensor_def_configs]
 
 
 def _is_floating_point_type(t: Type) -> bool:
   # TODO: Create a FloatType in the Python API and implement the switch
   # there.
   return (F64Type.isinstance(t) or F32Type.isinstance(t) or
           F16Type.isinstance(t) or BF16Type.isinstance(t))
 
 
 def _is_integer_type(t: Type) -> bool:
   return IntegerType.isinstance(t)
 
 
 def _get_floating_point_width(t: Type) -> int:
   # TODO: Create a FloatType in the Python API and implement the switch
   # there.
   if F64Type.isinstance(t):
     return 64
   if F32Type.isinstance(t):
     return 32
   if F16Type.isinstance(t):
     return 16
   if BF16Type.isinstance(t):
     return 16
   raise NotImplementedError(f"Unhandled floating point type switch {t}")

mlir/python/mlir/dialects/linalg/opdsl/lang/scalar_expr.py

 #  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
 """Models DAGs of scalar math expressions.
 
 Used for generating region bodies at the "math" level where they are still type
 polymorphic. This is modeled to be polymorphic by attribute name for interop
 with serialization schemes that are just plain-old-dicts.
 
 These classes are typically not user accessed and are created as a by-product
 of interpreting a comprehension DSL and model the operations to perform in the
 op body. The class hierarchy is laid out to map well to a form of YAML that
 can be easily consumed from the C++ side, not necessarily for ergonomics.
 """
 
 from typing import Optional, Sequence
 
 from .yaml_helper import *
 from .types import *
 
 __all__ = [
     "ScalarAssign",
     "ScalarApplyFn",
     "ScalarArg",
     "ScalarExpression",
     "ScalarSymbolicCast",
 ]
 
 
 class ScalarApplyFn:
   """A type of ScalarExpression that applies a named function to operands."""
 
   def __init__(self, fn_name: str, *operands: "ScalarExpression"):
     self.fn_name = fn_name
     self.operands = operands
 
   def expr(self) -> "ScalarExpression":
     return ScalarExpression(scalar_apply=self)
 
   def __repr__(self):
     return f"ScalarApplyFn<{self.fn_name}>({', '.join(self.operands)})"
 
 
 class ScalarArg:
   """A type of ScalarExpression that references a named argument."""
 
   def __init__(self, arg: str):
     self.arg = arg
 
   def expr(self) -> "ScalarExpression":
     return ScalarExpression(scalar_arg=self)
 
   def __repr__(self):
     return f"(ScalarArg({self.arg})"
 
 
 class ScalarSymbolicCast:
   """A type of ScalarExpression that symbolically casts an operand to a TypeVar.
   """
 
   def __init__(self, to_type: TypeVar, operand: "ScalarExpression"):
     self.to_type = to_type
     self.operand = operand
 
   def expr(self) -> "ScalarExpression":
     return ScalarExpression(symbolic_cast=self)
 
   def __repr__(self):
     return f"ScalarSymbolicCast({self.to_type}, {self.operand})"
 
 
 class ScalarExpression(YAMLObject):
   """An expression on scalar values.
 
   Can be one of:
     - ScalarApplyFn
     - ScalarArg
     - ScalarSymbolicCast
   """
   yaml_tag = "!ScalarExpression"
 
   def __init__(self,
                scalar_apply: Optional[ScalarApplyFn] = None,
                scalar_arg: Optional[ScalarArg] = None,
                symbolic_cast: Optional[ScalarSymbolicCast] = None):
     if (bool(scalar_apply) + bool(scalar_arg) + bool(symbolic_cast)) != 1:
       raise ValueError(
           "One of 'scalar_apply', 'scalar_block_arg', 'symbolic_cast' must be "
           "specified")
     self.scalar_apply = scalar_apply
     self.scalar_arg = scalar_arg
     self.symbolic_cast = symbolic_cast
 
   def to_yaml_custom_dict(self):
     if self.scalar_apply:
       return dict(scalar_apply=dict(
           fn_name=self.scalar_apply.fn_name,
           operands=list(self.scalar_apply.operands),
       ))
     elif self.scalar_arg:
       return dict(scalar_arg=self.scalar_arg.arg)
     elif self.symbolic_cast:
       # Note that even though operands must be arity 1, we write it the
       # same way as for apply because it allows handling code to be more
       # generic vs having a special form.
       return dict(symbolic_cast=dict(type_var=self.symbolic_cast.to_type.name,
                                      operands=[self.symbolic_cast.operand]))
     else:
       raise ValueError(f"Unexpected ScalarExpression type: {self}")
 
 
 class ScalarAssign(YAMLObject):
   """An assignment to a named argument (LHS of a comprehension)."""
   yaml_tag = "!ScalarAssign"
 
   def __init__(self, arg: str, value: ScalarExpression):
     self.arg = arg
     self.value = value
 
   def to_yaml_custom_dict(self):
     return dict(arg=self.arg, value=self.value)
 
   def __repr__(self):
     return f"ScalarAssign({self.arg}, {self.value})"

mlir/python/mlir/dialects/linalg/opdsl/lang/types.py

 #  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
 """Facility for symbolically referencing type variables.
 
 Type variables are instances of the TypeVar class, which is uniqued by name.
 An "expando" accessor `TV` is provided that generates a named TypeVar for
 any attribute access:
 
   >>> TV.T
   TypeVar(T)
   >>> TV.T is TV.U
   False
   >>> TV.T is TV.T
   True
 """
 
 from enum import Enum
 from typing import Dict
 
 __all__ = [
     "TypeVar",
     "TV",
 
     # TypeVar aliases.
     "T",
     "U",
     "V",
 ]
 
 
 class TypeVar:
   """A replaceable type variable.
 
   Type variables are uniqued by name.
   """
   ALL_TYPEVARS = dict()  # type: Dict[str, "TypeVar"]
 
   def __new__(cls, name: str):
     existing = cls.ALL_TYPEVARS.get(name)
     if existing is not None:
       return existing
     new = super().__new__(cls)
     new.name = name
     cls.ALL_TYPEVARS[name] = new
     return new
 
   def __repr__(self):
     return f"TypeVar({self.name})"
 
   @classmethod
   def create_expando(cls):
     """Create an expando class that creates unique type vars on attr access."""
 
     class ExpandoTypeVars:
 
       def __getattr__(self, n):
         return cls(n)
 
     return ExpandoTypeVars()
 
 
 # Expando access via TV.foo
 TV = TypeVar.create_expando()
 
 # Some common type name aliases.
 T = TV.T
 U = TV.U
 V = TV.V

mlir/python/mlir/dialects/linalg/opdsl/lang/yaml_helper.py

 #  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
 """YAML serialization is routed through here to centralize common logic."""
 
 import sys
 
 try:
   import yaml
 except ModuleNotFoundError as e:
   raise ModuleNotFoundError(
       f"This tool requires PyYAML but it was not installed. "
       f"Recommend: {sys.executable} -m pip install PyYAML") from e
 
 __all__ = [
     "yaml_dump",
     "yaml_dump_all",
     "YAMLObject",
 ]
 
 
 class YAMLObject(yaml.YAMLObject):
 
   @classmethod
   def to_yaml(cls, dumper, self):
     """Default to a custom dictionary mapping."""
     return dumper.represent_mapping(cls.yaml_tag, self.to_yaml_custom_dict())
 
   def to_yaml_custom_dict(self):
     raise NotImplementedError()
 
   def as_linalg_yaml(self):
     return yaml_dump(self)
 
 
 def multiline_str_representer(dumper, data):
   if len(data.splitlines()) > 1:
     return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')
   else:
     return dumper.represent_scalar('tag:yaml.org,2002:str', data)
 
 
 yaml.add_representer(str, multiline_str_representer)
 
 
 def yaml_dump(data, sort_keys=False, **kwargs):
   return yaml.dump(data, sort_keys=sort_keys, **kwargs)
 
 
 def yaml_dump_all(data, sort_keys=False, explicit_start=True, **kwargs):
   return yaml.dump_all(data,
                        sort_keys=sort_keys,
                        explicit_start=explicit_start,
                        **kwargs)

mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py

 from ..lang import *
 
 T1 = TV.T1
 T2 = TV.T2
 
 Batch = S.Batch
 
 
 @linalg_structured_op
 def matmul(A=TensorDef(T1, S.M, S.K),
            B=TensorDef(T2, S.K, S.N),
            C=TensorDef(U, S.M, S.N, output=True)):
   """Performs a matrix multiplication of two 2D inputs.
 
   Numeric casting is performed on the operands to the inner multiply, promoting
   them to the same data type as the accumulator/output.
   """
   implements(ContractionOpInterface)
   C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
 
 
 @linalg_structured_op
 def batch_matmul(A=TensorDef(T1, Batch, S.M, S.K),
                  B=TensorDef(T2, Batch, S.K, S.N),
                  C=TensorDef(U, Batch, S.M, S.N, output=True)):
   """Performs a batched matrix multiplication of two 3D inputs.
 
   Numeric casting is performed on the operands to the inner multiply, promoting
   them to the same data type as the accumulator/output.
   """
   implements(ContractionOpInterface)
   C[D.b, D.m, D.n] += cast(U, A[D.b, D.m, D.k]) * cast(U, B[D.b, D.k, D.n])
 
 
 @linalg_structured_op
 def matvec(A=TensorDef(T1, S.M, S.N),
            y=TensorDef(T2, S.N),
            x=TensorDef(U, S.M, output=True)):
   """Performs a matrix-vector multiplication.
 
   Numeric casting is performed on the operands to the inner multiply, promoting
   them to the same data type as the accumulator/output.
   """
   implements(ContractionOpInterface)
   x[D.m] += cast(U, A[D.m, D.n]) * cast(U, y[D.n])
 
 
 @linalg_structured_op
 def vecmat(y=TensorDef(T1, S.M),
            A=TensorDef(T2, S.M, S.N),
            x=TensorDef(U, S.N, output=True)):
   """Performs a vector-matrix multiplication.
 
   Numeric casting is performed on the operands to the inner multiply, promoting
   them to the same data type as the accumulator/output.
   """
   implements(ContractionOpInterface)
   x[D.n] += cast(U, y[D.m]) * cast(U, A[D.m, D.n])
 
 
 @linalg_structured_op
 def dot(A=TensorDef(T1, S.M), B=TensorDef(T2, S.M), C=TensorDef(U,
                                                                 output=True)):
   """Performs a dot product of two vectors to a scalar result.
 
   Numeric casting is performed on the operands to the inner multiply, promoting
   them to the same data type as the accumulator/output.
   """
   implements(ContractionOpInterface)
   C[None] += cast(U, A[D.m]) * cast(U, B[D.m])

mlir/python/mlir/dialects/linalg/passes/__init__.py

 #  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 ...._cext_loader import _load_extension
 _cextLinalgPasses = _load_extension("_mlirLinalgPasses")

mlir/python/mlir/dialects/memref.py

 #  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 ._memref_ops_gen import *

mlir/python/mlir/dialects/python_test.py

 #  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 *

mlir/python/mlir/dialects/shape.py

 #  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 *

mlir/python/mlir/dialects/std.py

 #  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 *

mlir/python/mlir/dialects/tensor.py

 #  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 *

mlir/python/mlir/execution_engine.py

 #  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
 
 # Simply a wrapper around the extension module of the same name.
 from ._cext_loader import _cext
 import ctypes
 
 class ExecutionEngine(_cext.execution_engine.ExecutionEngine):
 
   def lookup(self, name):
     """Lookup a function emitted with the `llvm.emit_c_interface`
     attribute and returns a ctype callable.
     Raise a RuntimeError if the function isn't found.
     """
     func = self.raw_lookup("_mlir_ciface_" + name)
     if not func:
       raise RuntimeError("Unknown function " + name)
     prototype = ctypes.CFUNCTYPE(None, ctypes.c_void_p)
     return prototype(func)
 
   def invoke(self, name, *ctypes_args):
     """Invoke a function with the list of ctypes arguments.
     All arguments must be pointers.
     Raise a RuntimeError if the function isn't found.
     """
     func = self.lookup(name)
     packed_args = (ctypes.c_void_p * len(ctypes_args))()
     for argNum in range(len(ctypes_args)):
       packed_args[argNum] = ctypes.cast(ctypes_args[argNum], ctypes.c_void_p)
     func(packed_args)
 
   def register_runtime(self, name, ctypes_callback):
     """Register a runtime function available to the jitted code
     under the provided `name`. The `ctypes_callback` must be a
     `CFuncType` that outlives the execution engine.
     """
     callback = ctypes.cast(ctypes_callback, ctypes.c_void_p).value
     self.raw_register_runtime("_mlir_ciface_" + name, callback)

mlir/python/mlir/ir.py

 #  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
 
 # Simply a wrapper around the extension module of the same name.
 from ._cext_loader import _reexport_cext
 _reexport_cext("ir", __name__)
 del _reexport_cext
 

mlir/python/mlir/passmanager.py

 #  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
 
 # Simply a wrapper around the extension module of the same name.
 from ._cext_loader import _reexport_cext
 _reexport_cext("passmanager", __name__)
 del _reexport_cext

mlir/python/mlir/runtime/__init__.py

 from .np_to_memref import *

mlir/python/mlir/runtime/np_to_memref.py

 #  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 file contains functions to convert between Memrefs and NumPy arrays and vice-versa.
 
 import numpy as np
 import ctypes
 
 
 def make_nd_memref_descriptor(rank, dtype):
     class MemRefDescriptor(ctypes.Structure):
         """
         Build an empty descriptor for the given rank/dtype, where rank>0.
         """
 
         _fields_ = [
             ("allocated", ctypes.c_longlong),
             ("aligned", ctypes.POINTER(dtype)),
             ("offset", ctypes.c_longlong),
             ("shape", ctypes.c_longlong * rank),
             ("strides", ctypes.c_longlong * rank),
         ]
 
     return MemRefDescriptor
 
 
 def make_zero_d_memref_descriptor(dtype):
     class MemRefDescriptor(ctypes.Structure):
         """
         Build an empty descriptor for the given dtype, where rank=0.
         """
 
         _fields_ = [
             ("allocated", ctypes.c_longlong),
             ("aligned", ctypes.POINTER(dtype)),
             ("offset", ctypes.c_longlong),
         ]
 
     return MemRefDescriptor
 
 
 class UnrankedMemRefDescriptor(ctypes.Structure):
     """ Creates a ctype struct for memref descriptor"""
 
     _fields_ = [("rank", ctypes.c_longlong), ("descriptor", ctypes.c_void_p)]
 
 
 def get_ranked_memref_descriptor(nparray):
     """
     Return a ranked memref descriptor for the given numpy array.
     """
     if nparray.ndim == 0:
         x = make_zero_d_memref_descriptor(np.ctypeslib.as_ctypes_type(nparray.dtype))()
         x.allocated = nparray.ctypes.data
         x.aligned = nparray.ctypes.data_as(
             ctypes.POINTER(np.ctypeslib.as_ctypes_type(nparray.dtype))
         )
         x.offset = ctypes.c_longlong(0)
         return x
 
     x = make_nd_memref_descriptor(
         nparray.ndim, np.ctypeslib.as_ctypes_type(nparray.dtype)
     )()
     x.allocated = nparray.ctypes.data
     x.aligned = nparray.ctypes.data_as(
         ctypes.POINTER(np.ctypeslib.as_ctypes_type(nparray.dtype))
     )
     x.offset = ctypes.c_longlong(0)
     x.shape = nparray.ctypes.shape
 
     # Numpy uses byte quantities to express strides, MLIR OTOH uses the
     # torch abstraction which specifies strides in terms of elements.
     strides_ctype_t = ctypes.c_longlong * nparray.ndim
     x.strides = strides_ctype_t(*[x // nparray.itemsize for x in nparray.strides])
     return x
 
 
 def get_unranked_memref_descriptor(nparray):
     """
     Return a generic/unranked memref descriptor for the given numpy array.
     """
     d = UnrankedMemRefDescriptor()
     d.rank = nparray.ndim
     x = get_ranked_memref_descriptor(nparray)
     d.descriptor = ctypes.cast(ctypes.pointer(x), ctypes.c_void_p)
     return d
 
 
 def unranked_memref_to_numpy(unranked_memref, np_dtype):
     """
     Converts unranked memrefs to numpy arrays.
     """
     descriptor = make_nd_memref_descriptor(
         unranked_memref[0].rank, np.ctypeslib.as_ctypes_type(np_dtype)
     )
     val = ctypes.cast(unranked_memref[0].descriptor, ctypes.POINTER(descriptor))
     np_arr = np.ctypeslib.as_array(val[0].aligned, shape=val[0].shape)
     strided_arr = np.lib.stride_tricks.as_strided(
         np_arr,
         np.ctypeslib.as_array(val[0].shape),
         np.ctypeslib.as_array(val[0].strides) * np_arr.itemsize,
     )
     return strided_arr
 
 
 def ranked_memref_to_numpy(ranked_memref):
     """
     Converts ranked memrefs to numpy arrays.
     """
     np_arr = np.ctypeslib.as_array(
         ranked_memref[0].aligned, shape=ranked_memref[0].shape
     )
     strided_arr = np.lib.stride_tricks.as_strided(
         np_arr,
         np.ctypeslib.as_array(ranked_memref[0].shape),
         np.ctypeslib.as_array(ranked_memref[0].strides) * np_arr.itemsize,
     )
     return strided_arr

mlir/python/mlir/transforms/__init__.py

 #  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
 
 # Expose the corresponding C-Extension module with a well-known name at this
 # level.
 from .._cext_loader import _load_extension
 _cextTransforms = _load_extension("_mlirTransforms")

mlir/python/requirements.txt

 numpy
 pybind11>=2.6.0
 PyYAML

mlir/test/Bindings/CMakeLists.txt

-if(MLIR_BINDINGS_PYTHON_ENABLED)
-  add_subdirectory(Python)
-endif()

mlir/test/Bindings/Python/.style.yapf

-[style]
-  based_on_style = google
-  column_limit = 80
-  indent_width = 2

mlir/test/CMakeLists.txt

-add_subdirectory(Bindings)
 add_subdirectory(CAPI)
 add_subdirectory(EDSC)
 add_subdirectory(SDBM)
 add_subdirectory(lib)
 
+if(MLIR_BINDINGS_PYTHON_ENABLED)
+  add_subdirectory(python)
+endif()
+
 llvm_canonicalize_cmake_booleans(
   MLIR_BINDINGS_PYTHON_ENABLED
   LLVM_BUILD_EXAMPLES
   MLIR_CUDA_CONVERSIONS_ENABLED
   MLIR_CUDA_RUNNER_ENABLED
   MLIR_ROCM_CONVERSIONS_ENABLED
   MLIR_ROCM_RUNNER_ENABLED
   MLIR_SPIRV_CPU_RUNNER_ENABLED
if(MLIR_VULKAN_RUNNER_ENABLED)
   list(APPEND MLIR_TEST_DEPENDS
     mlir-vulkan-runner
   )
 endif()
 
 if(MLIR_BINDINGS_PYTHON_ENABLED)
   list(APPEND MLIR_TEST_DEPENDS
     MLIRBindingsPythonExtension
+    MLIRBindingsPythonSources
     MLIRBindingsPythonTestOps
     MLIRTransformsBindingsPythonExtension
     MLIRConversionsBindingsPythonExtension
   )
 endif()
 
 add_lit_testsuite(check-mlir "Running the MLIR regression tests"
   ${CMAKE_CURRENT_BINARY_DIR}
   DEPENDS ${MLIR_TEST_DEPENDS}
   )
 set_target_properties(check-mlir PROPERTIES FOLDER "Tests")
 
 add_lit_testsuites(MLIR ${CMAKE_CURRENT_SOURCE_DIR}
   DEPENDS ${MLIR_TEST_DEPENDS}
 )

mlir/test/python/CMakeLists.txt

-include(AddMLIRPythonExtension)
+include(AddMLIRPython)
+
 add_mlir_dialect_python_bindings(MLIRBindingsPythonTestOps
   TD_FILE python_test_ops.td
   DIALECT_NAME python_test)

mlir/test/python/dialects/async_dialect.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 import mlir.dialects.async_dialect
 import mlir.dialects.async_dialect.passes
 from mlir.passmanager import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
 
 def testAsyncPass():
   with Context() as context:
     PassManager.parse('async-to-async-runtime')
   print('SUCCESS')
 
 # CHECK-LABEL: testAsyncPass
 #       CHECK: SUCCESS
 run(testAsyncPass)

mlir/test/python/dialects/builtin.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 import mlir.dialects.builtin as builtin
 import mlir.dialects.std as std
 
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   return f
 
 
 # CHECK-LABEL: TEST: testFromPyFunc
 @run
 def testFromPyFunc():
   with Context() as ctx, Location.unknown() as loc:
     m = builtin.ModuleOp()
     f32 = F32Type.get()
     f64 = F64Type.get()
     with InsertionPoint(m.body):
       # CHECK-LABEL: func @unary_return(%arg0: f64) -> f64
       # CHECK: return %arg0 : f64
       @builtin.FuncOp.from_py_func(f64)
       def unary_return(a):
         return a
 
       # CHECK-LABEL: func @binary_return(%arg0: f32, %arg1: f64) -> (f32, f64)
       # CHECK: return %arg0, %arg1 : f32, f64
       @builtin.FuncOp.from_py_func(f32, f64)
       def binary_return(a, b):
         return a, b
 
       # CHECK-LABEL: func @none_return(%arg0: f32, %arg1: f64)
       # CHECK: return
       @builtin.FuncOp.from_py_func(f32, f64)
       def none_return(a, b):
         pass
 
       # CHECK-LABEL: func @call_unary
       # CHECK: %0 = call @unary_return(%arg0) : (f64) -> f64
       # CHECK: return %0 : f64
       @builtin.FuncOp.from_py_func(f64)
       def call_unary(a):
         return unary_return(a)
 
       # CHECK-LABEL: func @call_binary
       # CHECK: %0:2 = call @binary_return(%arg0, %arg1) : (f32, f64) -> (f32, f64)
       # CHECK: return %0#0, %0#1 : f32, f64
       @builtin.FuncOp.from_py_func(f32, f64)
       def call_binary(a, b):
         return binary_return(a, b)
 
       # CHECK-LABEL: func @call_none
       # CHECK: call @none_return(%arg0, %arg1) : (f32, f64) -> ()
       # CHECK: return
       @builtin.FuncOp.from_py_func(f32, f64)
       def call_none(a, b):
         return none_return(a, b)
 
       ## Variants and optional feature tests.
       # CHECK-LABEL: func @from_name_arg
       @builtin.FuncOp.from_py_func(f32, f64, name="from_name_arg")
       def explicit_name(a, b):
         return b
 
       @builtin.FuncOp.from_py_func(f32, f64)
       def positional_func_op(a, b, func_op):
         assert isinstance(func_op, builtin.FuncOp)
         return b
 
       @builtin.FuncOp.from_py_func(f32, f64)
       def kw_func_op(a, b=None, func_op=None):
         assert isinstance(func_op, builtin.FuncOp)
         return b
 
       @builtin.FuncOp.from_py_func(f32, f64)
       def kwargs_func_op(a, b=None, **kwargs):
         assert isinstance(kwargs["func_op"], builtin.FuncOp)
         return b
 
       # CHECK-LABEL: func @explicit_results(%arg0: f32, %arg1: f64) -> f64
       # CHECK: return %arg1 : f64
       @builtin.FuncOp.from_py_func(f32, f64, results=[f64])
       def explicit_results(a, b):
         std.ReturnOp([b])
 
   print(m)
 
 
 # CHECK-LABEL: TEST: testFromPyFuncErrors
 @run
 def testFromPyFuncErrors():
   with Context() as ctx, Location.unknown() as loc:
     m = builtin.ModuleOp()
     f32 = F32Type.get()
     f64 = F64Type.get()
     with InsertionPoint(m.body):
       try:
 
         @builtin.FuncOp.from_py_func(f64, results=[f64])
         def unary_return(a):
           return a
       except AssertionError as e:
         # CHECK: Capturing a python function with explicit `results=` requires that the wrapped function returns None.
         print(e)
 
 
 # CHECK-LABEL: TEST: testBuildFuncOp
 @run
 def testBuildFuncOp():
   ctx = Context()
   with Location.unknown(ctx) as loc:
     m = builtin.ModuleOp()
 
     f32 = F32Type.get()
     tensor_type = RankedTensorType.get((2, 3, 4), f32)
     with InsertionPoint.at_block_begin(m.body):
       func = builtin.FuncOp(name="some_func",
                             type=FunctionType.get(
                                 inputs=[tensor_type, tensor_type],
                                 results=[tensor_type]),
                             visibility="nested")
       # CHECK: Name is: "some_func"
       print("Name is: ", func.name)
 
       # CHECK: Type is: (tensor<2x3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32>
       print("Type is: ", func.type)
 
       # CHECK: Visibility is: "nested"
       print("Visibility is: ", func.visibility)
 
       try:
         entry_block = func.entry_block
       except IndexError as e:
         # CHECK: External function does not have a body
         print(e)
 
       with InsertionPoint(func.add_entry_block()):
         std.ReturnOp([func.entry_block.arguments[0]])
         pass
 
       try:
         func.add_entry_block()
       except IndexError as e:
         # CHECK: The function already has an entry block!
         print(e)
 
       # Try the callback builder and passing type as tuple.
       func = builtin.FuncOp(name="some_other_func",
                             type=([tensor_type, tensor_type], [tensor_type]),
                             visibility="nested",
                             body_builder=lambda func: std.ReturnOp(
                                 [func.entry_block.arguments[0]]))
 
   # CHECK: module  {
   # CHECK:  func nested @some_func(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
   # CHECK:   return %arg0 : tensor<2x3x4xf32>
   # CHECK:  }
   # CHECK:  func nested @some_other_func(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
   # CHECK:   return %arg0 : tensor<2x3x4xf32>
   # CHECK:  }
   print(m)

mlir/test/python/dialects/gpu.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 import mlir.dialects.gpu
 import mlir.dialects.gpu.passes
 from mlir.passmanager import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
 
 def testGPUPass():
   with Context() as context:
     PassManager.parse('gpu-kernel-outlining')
   print('SUCCESS')
 
 # CHECK-LABEL: testGPUPass
 #       CHECK: SUCCESS
 run(testGPUPass)

mlir/test/python/dialects/linalg/opdsl/assignments.py

 # RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s | FileCheck %s
 
 from mlir.dialects.linalg.opdsl.lang import *
 
 # CHECK: ---
 # CHECK-LABEL: matmul
 # CHECK: assignments:
 # CHECK:  -
 # CHECK:    arg: C
 # CHECK:    value:
 # CHECK:      scalar_apply:
 # CHECK:        fn_name: add
 # CHECK:        operands:
 # CHECK:          scalar_apply:
 # CHECK:            fn_name: mul
 # CHECK:            operands:
 # CHECK:              symbolic_cast:
 # CHECK:                type_var: U
 # CHECK:                operands:
 # CHECK:                  scalar_arg: A
 # CHECK:              symbolic_cast:
 # CHECK:                type_var: U
 # CHECK:                operands:
 # CHECK:                  scalar_arg: B
 @linalg_structured_op
 def matmul(A=TensorDef(T, S.M, S.K),
            B=TensorDef(T, S.K, S.N),
            C=TensorDef(U, S.M, S.N, output=True)):
   C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])

mlir/test/python/dialects/linalg/opdsl/doctests.py

 # RUN: %PYTHON %s
 
 import doctest
 import importlib
 
 def test_module(module_name):
   print(f"--- Testing module: {module_name}")
   m = importlib.import_module(module_name)
   doctest.testmod(m, verbose=True, raise_on_error=True, report=True)
 
 
 test_module("mlir.dialects.linalg.opdsl.lang.affine")
 test_module("mlir.dialects.linalg.opdsl.lang.types")

mlir/test/python/dialects/linalg/opdsl/emit_structured_generic.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from typing import Optional, Sequence
 
 from mlir.ir import *
 from mlir.dialects import builtin
 from mlir.dialects import linalg
 from mlir.dialects import std
 
 from mlir.dialects.linalg.opdsl.lang import *
 
 
 @linalg_structured_op
 def matmul_mono(A=TensorDef(T, S.M, S.K),
                 B=TensorDef(T, S.K, S.N),
                 C=TensorDef(T, S.M, S.N, output=True)):
   C[D.m, D.n] += A[D.m, D.k] * B[D.k, D.n]
 
 
 @linalg_structured_op
 def matmul_poly(A=TensorDef(TV.T1, S.M, S.K),
                 B=TensorDef(TV.T2, S.K, S.N),
                 C=TensorDef(U, S.M, S.N, output=True)):
   C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
 
 
 with Context() as ctx, Location.unknown():
   module = Module.create()
   f16 = F16Type.get()
   f32 = F32Type.get()
   f64 = F64Type.get()
   i8 = IntegerType.get_signless(8)
   i16 = IntegerType.get_signless(16)
   i32 = IntegerType.get_signless(32)
   with InsertionPoint(module.body):
 
     # Note that these all have the same indexing maps. We verify the first and
     # then do more permutation tests on casting and body generation
     # behavior.
     # CHECK: #[[$MAPA:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
     # CHECK: #[[$MAPB:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
     # CHECK: #[[$MAPC:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
 
     # CHECK-LABEL: func @test_matmul_mono
     # CHECK-SAME:  %[[A:.+]]: tensor<4x16xf32>
     # CHECK-SAME: %[[B:.+]]: tensor<16x8xf32>
 
     # CHECK: %[[INITC:.+]] = linalg.init_tensor [4, 8] : tensor<4x8xf32>
     # CHECK: linalg.generic
     # CHECK-SAME: indexing_maps = [#[[$MAPA]], #[[$MAPB]], #[[$MAPC]]]
     # CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
     # CHECK-SAME: ins(%[[A]], %[[B]]
     # CHECK-SAME: outs(%[[INITC]]
 
     @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32),
                                  RankedTensorType.get((16, 8), f32))
     def test_matmul_mono(lhs, rhs):
       init_result = linalg.InitTensorOp([4, 8], f32)
       return matmul_mono(lhs, rhs, outs=[init_result.result])
 
     # CHECK-LABEL: @test_i8i8i32_matmul
     # CHECK:      ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i8, %[[C_ARG:.+]]: i32)
     # CHECK-NEXT:   %[[A_CAST:.+]] = sexti %[[A_ARG]] : i8 to i32
     # CHECK-NEXT:   %[[B_CAST:.+]] = sexti %[[B_ARG]] : i8 to i32
     # CHECK-NEXT:   %[[MUL:.+]] = muli %[[A_CAST]], %[[B_CAST]] : i32
     # CHECK-NEXT:   %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i32
     # CHECK-NEXT:   linalg.yield %[[ADD]] : i32
     # CHECK-NEXT: -> tensor<4x8xi32>
     @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), i8),
                                  RankedTensorType.get((16, 8), i8),
                                  RankedTensorType.get((4, 8), i32))
     def test_i8i8i32_matmul(lhs, rhs, init_result):
       return matmul_poly(lhs, rhs, outs=[init_result])
 
     # CHECK-LABEL: @test_i8i16i32_matmul
     # CHECK:      ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i16, %[[C_ARG:.+]]: i32)
     # CHECK-NEXT:   %[[A_CAST:.+]] = sexti %[[A_ARG]] : i8 to i32
     # CHECK-NEXT:   %[[B_CAST:.+]] = sexti %[[B_ARG]] : i16 to i32
     # CHECK-NEXT:   %[[MUL:.+]] = muli %[[A_CAST]], %[[B_CAST]] : i32
     # CHECK-NEXT:   %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i32
     # CHECK-NEXT:   linalg.yield %[[ADD]] : i32
     # CHECK-NEXT: -> tensor<4x8xi32>
     @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), i8),
                                  RankedTensorType.get((16, 8), i16),
                                  RankedTensorType.get((4, 8), i32))
     def test_i8i16i32_matmul(lhs, rhs, init_result):
       return matmul_poly(lhs, rhs, outs=[init_result])
 
     # CHECK-LABEL: @test_i32i32i16_matmul
     # CHECK:      ^{{.*}}(%[[A_ARG:.+]]: i32, %[[B_ARG:.+]]: i32, %[[C_ARG:.+]]: i16)
     # CHECK-NEXT:   %[[A_CAST:.+]] = trunci %[[A_ARG]] : i32 to i16
     # CHECK-NEXT:   %[[B_CAST:.+]] = trunci %[[B_ARG]] : i32 to i16
     # CHECK-NEXT:   %[[MUL:.+]] = muli %[[A_CAST]], %[[B_CAST]] : i16
     # CHECK-NEXT:   %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i16
     # CHECK-NEXT:   linalg.yield %[[ADD]] : i16
     # CHECK-NEXT: -> tensor<4x8xi16>
     @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), i32),
                                  RankedTensorType.get((16, 8), i32),
                                  RankedTensorType.get((4, 8), i16))
     def test_i32i32i16_matmul(lhs, rhs, init_result):
       return matmul_poly(lhs, rhs, outs=[init_result])
 
     # CHECK-LABEL: @test_i8i8f32_matmul
     # CHECK:      ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i8, %[[C_ARG:.+]]: f32)
     # CHECK-NEXT:   %[[A_CAST:.+]] = sitofp %[[A_ARG]] : i8 to f32
     # CHECK-NEXT:   %[[B_CAST:.+]] = sitofp %[[B_ARG]] : i8 to f32
     # CHECK-NEXT:   %[[MUL:.+]] = mulf %[[A_CAST]], %[[B_CAST]] : f32
     # CHECK-NEXT:   %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32
     # CHECK-NEXT:   linalg.yield %[[ADD]] : f32
     # CHECK-NEXT: -> tensor<4x8xf32>
     @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), i8),
                                  RankedTensorType.get((16, 8), i8),
                                  RankedTensorType.get((4, 8), f32))
     def test_i8i8f32_matmul(lhs, rhs, init_result):
       return matmul_poly(lhs, rhs, outs=[init_result])
 
     # CHECK-LABEL: @test_f16f16f32_matmul
     # CHECK:      ^{{.*}}(%[[A_ARG:.+]]: f16, %[[B_ARG:.+]]: f16, %[[C_ARG:.+]]: f32)
     # CHECK-NEXT:   %[[A_CAST:.+]] = fpext %[[A_ARG]] : f16 to f32
     # CHECK-NEXT:   %[[B_CAST:.+]] = fpext %[[B_ARG]] : f16 to f32
     # CHECK-NEXT:   %[[MUL:.+]] = mulf %[[A_CAST]], %[[B_CAST]] : f32
     # CHECK-NEXT:   %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32
     # CHECK-NEXT:   linalg.yield %[[ADD]] : f32
     # CHECK-NEXT: -> tensor<4x8xf32>
     @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f16),
                                  RankedTensorType.get((16, 8), f16),
                                  RankedTensorType.get((4, 8), f32))
     def test_f16f16f32_matmul(lhs, rhs, init_result):
       return matmul_poly(lhs, rhs, outs=[init_result])
 
     # CHECK-LABEL: @test_f64f64f32_matmul
     # CHECK:      ^{{.*}}(%[[A_ARG:.+]]: f64, %[[B_ARG:.+]]: f64, %[[C_ARG:.+]]: f32)
     # CHECK-NEXT:   %[[A_CAST:.+]] = fptrunc %[[A_ARG]] : f64 to f32
     # CHECK-NEXT:   %[[B_CAST:.+]] = fptrunc %[[B_ARG]] : f64 to f32
     # CHECK-NEXT:   %[[MUL:.+]] = mulf %[[A_CAST]], %[[B_CAST]] : f32
     # CHECK-NEXT:   %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32
     # CHECK-NEXT:   linalg.yield %[[ADD]] : f32
     # CHECK-NEXT: -> tensor<4x8xf32>
     @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f64),
                                  RankedTensorType.get((16, 8), f64),
                                  RankedTensorType.get((4, 8), f32))
     def test_f64f64f32_matmul(lhs, rhs, init_result):
       return matmul_poly(lhs, rhs, outs=[init_result])
 
 
 print(module)

mlir/test/python/dialects/linalg/opdsl/interfaces.py

 # RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s | FileCheck %s
 
 from mlir.dialects.linalg.opdsl.lang import *
 
 # CHECK: ---
 # CHECK-LABEL: matmul
 # CHECK:      implements:
 # CHECK-NEXT: - LinalgContractionOpInterface
 @linalg_structured_op
 def matmul(A=TensorDef(T, S.M, S.K),
            B=TensorDef(T, S.K, S.N),
            C=TensorDef(U, S.M, S.N, output=True)):
   implements(ContractionOpInterface)
   C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])

mlir/test/python/dialects/linalg/opdsl/lit.local.cfg

 # TODO: This tool requires PyYAML, which is not yet a required build/test
 # dependency. Remove this exclusion once it is a required dep.
 
 # Since both lit and the python bindings use the same python interpreter,
 # we can just check whether yaml can be imported here and exclude if not.
 try:
   import yaml
 except ModuleNotFoundError:
   config.unsupported = True

mlir/test/python/dialects/linalg/opdsl/shape_maps_iteration.py

 # RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s | FileCheck %s
 
 from mlir.dialects.linalg.opdsl.lang import *
 
 
 # Verify that simple case with iteration order defined lexically and reduction
 # dims auto discovered emits the right shape, indexing maps and iterator types.
 # CHECK: ---
 # CHECK-LABEL: matmul
 # CHECK: shape: affine_map<()[s0, s1, s2] -> (s0, s2)>
 # CHECK: shape: affine_map<()[s0, s1, s2] -> (s2, s1)>
 # CHECK: shape: affine_map<()[s0, s1, s2] -> (s0, s1)>
 # CHECK: static_indexing_maps:
 # CHECK-NEXT: - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d2)>
 # CHECK-NEXT: - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d2, d1)>
 # CHECK-NEXT: - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d1)>
 # CHECK: iterator_types:
 # CHECK-NEXT: - parallel
 # CHECK-NEXT: - parallel
 # CHECK-NEXT: - reduction
 @linalg_structured_op
 def matmul(A=TensorDef(T, S.M, S.K),
            B=TensorDef(T, S.K, S.N),
            C=TensorDef(U, S.M, S.N, output=True)):
   C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
 
 
 # Verifies that assignment to a scalar (represented as [None]) is represented
 # correctly.
 # CHECK: ---
 # CHECK-LABEL: dot
 # CHECK: shape: affine_map<()[s0] -> (s0)>
 # CHECK: shape: affine_map<()[s0] -> (s0)>
 # CHECK: shape: affine_map<()[s0] -> ()>
 # CHECK: static_indexing_maps:
 # CHECK-NEXT: - affine_map<(d0)[s0] -> (d0)>
 # CHECK-NEXT: - affine_map<(d0)[s0] -> (d0)>
 # CHECK-NEXT: - affine_map<(d0)[s0] -> ()>
 # CHECK: iterator_types:
 # CHECK-NEXT: - reduction
 @linalg_structured_op
 def dot(A=TensorDef(T, S.M), B=TensorDef(T, S.M), C=TensorDef(U, output=True)):
   C[None] += cast(U, A[D.m]) * cast(U, B[D.m])

mlir/test/python/dialects/linalg/opdsl/test_core_named_ops.py

 # 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

mlir/test/python/dialects/linalg/ops.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 from mlir.dialects import builtin
 from mlir.dialects import linalg
 from mlir.dialects import std
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   return f
 
 
 # CHECK-LABEL: TEST: testInitTensor
 @run
 def testInitTensor():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
     with InsertionPoint(module.body):
       # CHECK-LABEL: func @static_sizes
       # CHECK: %0 = linalg.init_tensor [3, 4] : tensor<3x4xf32>
       @builtin.FuncOp.from_py_func()
       def static_sizes():
         return linalg.InitTensorOp([3, 4], f32)
 
       # CHECK-LABEL: func @dynamic_sizes
       # CHECK: %0 = linalg.init_tensor [%arg0, %arg1] : tensor<?x?xf32>
       @builtin.FuncOp.from_py_func(IndexType.get(), IndexType.get())
       def dynamic_sizes(d0, d1):
         return linalg.InitTensorOp([d0, d1], f32)
 
       # CHECK-LABEL: func @zero_d
       # CHECK: %0 = linalg.init_tensor [] : tensor<f32>
       @builtin.FuncOp.from_py_func()
       def zero_d():
         return linalg.InitTensorOp([], f32)
 
   print(module)
 
 # CHECK-LABEL: TEST: testFill
 @run
 def testFill():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
     with InsertionPoint(module.body):
       # CHECK-LABEL: func @fill_tensor
       #  CHECK-SAME:   %[[OUT:[0-9a-z]+]]: tensor<12x?xf32>
       #  CHECK-NEXT: %[[CST:.*]] = constant 0.0{{.*}} : f32
       #  CHECK-NEXT: %[[RES:.*]] = linalg.fill(%[[OUT]], %[[CST]]) : tensor<12x?xf32>, f32 -> tensor<12x?xf32>
       #  CHECK-NEXT: return %[[RES]] : tensor<12x?xf32>
       @builtin.FuncOp.from_py_func(
           RankedTensorType.get((12, -1), f32))
       def fill_tensor(out):
         zero = std.ConstantOp(value=FloatAttr.get(f32, 0.), result=f32).result
         # TODO: FillOp.result is None. When len(results) == 1 we expect it to
         # be results[0] as per _linalg_ops_gen.py. This seems like an
         # orthogonal bug in the generator of _linalg_ops_gen.py.
         return linalg.FillOp(output=out, value=zero).results[0]
 
       # CHECK-LABEL: func @fill_buffer
       #  CHECK-SAME:   %[[OUT:[0-9a-z]+]]: memref<12x?xf32>
       #  CHECK-NEXT: %[[CST:.*]] = constant 0.0{{.*}} : f32
       #  CHECK-NEXT: linalg.fill(%[[OUT]], %[[CST]]) : memref<12x?xf32>, f32
       #  CHECK-NEXT: return
       @builtin.FuncOp.from_py_func(
           MemRefType.get((12, -1), f32))
       def fill_buffer(out):
         zero = std.ConstantOp(value=FloatAttr.get(f32, 0.), result=f32).result
         linalg.FillOp(output=out, value=zero)
 
   print(module)
 
 
 # CHECK-LABEL: TEST: testStructuredOpOnTensors
 @run
 def testStructuredOpOnTensors():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
     tensor_type = RankedTensorType.get((2, 3, 4), f32)
     with InsertionPoint(module.body):
       func = builtin.FuncOp(name="matmul_test",
                             type=FunctionType.get(
                                 inputs=[tensor_type, tensor_type],
                                 results=[tensor_type]))
       with InsertionPoint(func.add_entry_block()):
         lhs, rhs = func.entry_block.arguments
         result = linalg.MatmulOp([lhs, rhs], results=[tensor_type]).result
         std.ReturnOp([result])
 
   # CHECK: %[[R:.*]] = linalg.matmul ins(%arg0, %arg1 : tensor<2x3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32>
   print(module)
 
 
 # CHECK-LABEL: TEST: testStructuredOpOnBuffers
 @run
 def testStructuredOpOnBuffers():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
     memref_type = MemRefType.get((2, 3, 4), f32)
     with InsertionPoint(module.body):
       func = builtin.FuncOp(name="matmul_test",
                             type=FunctionType.get(
                                 inputs=[memref_type, memref_type, memref_type],
                                 results=[]))
       with InsertionPoint(func.add_entry_block()):
         lhs, rhs, result = func.entry_block.arguments
         # TODO: prperly hook up the region.
         linalg.MatmulOp([lhs, rhs], outputs=[result])
         std.ReturnOp([])
 
   # CHECK: linalg.matmul ins(%arg0, %arg1 : memref<2x3x4xf32>, memref<2x3x4xf32>) outs(%arg2 : memref<2x3x4xf32>)
   print(module)
 
 # CHECK-LABEL: TEST: testNamedStructuredOpCustomForm
 @run
 def testNamedStructuredOpCustomForm():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
     with InsertionPoint(module.body):
       @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32),
                                    RankedTensorType.get((16, 8), f32))
       def named_form(lhs, rhs):
         init_result = linalg.InitTensorOp([4, 8], f32)
         # First check the named form with custom format
         #      CHECK: linalg.matmul
         #  CHECK-NOT: linalg.memoized_indexing_maps
         # CHECK-SAME:    ins(%{{.*}} : tensor<4x16xf32>, tensor<16x8xf32>)
         # CHECK-SAME:   outs(%{{.*}} : tensor<4x8xf32>)
         # CHECK-SAME:   -> tensor<4x8xf32>
         # CHECK-NEXT: return
         return linalg.matmul(lhs, rhs, outs=[init_result.result])
 
   print(module)
 
 # CHECK-LABEL: TEST: testNamedStructuredOpGenericForm
 @run
 def testNamedStructuredOpGenericForm():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
     with InsertionPoint(module.body):
       @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32),
                                    RankedTensorType.get((16, 8), f32))
       def named_form(lhs, rhs):
         init_result = linalg.InitTensorOp([4, 8], f32)
         #      CHECK: "linalg.matmul"(%{{.*}})
         # CHECK-NEXT:  ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32):
         # CHECK-NEXT:    std.mulf{{.*}} (f32, f32) -> f32
         # CHECK-NEXT:    std.addf{{.*}} (f32, f32) -> f32
         # CHECK-NEXT:    linalg.yield{{.*}} (f32) -> ()
         # CHECK-NEXT:    {linalg.memoized_indexing_maps{{.*}}operand_segment_sizes = dense<[2, 1]> : vector<2xi32>} : 
         # CHECK-SAME: (tensor<4x16xf32>, tensor<16x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32>
         return linalg.matmul(lhs, rhs, outs=[init_result.result])
 
   module.operation.print(print_generic_op_form=True)
 
 # CHECK-LABEL: TEST: testNamedStructuredAsGenericOp
 @run
 def testNamedStructuredAsGenericOp():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
     with InsertionPoint(module.body):
       @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32),
                                    RankedTensorType.get((16, 8), f32))
       def generic_form(lhs, rhs):
         init_result = linalg.InitTensorOp([4, 8], f32)
         # CHECK: linalg.generic
         return linalg.matmul(lhs, rhs, outs=[init_result.result], emit_generic=True)
 
   print(module)

mlir/test/python/dialects/linalg/opsrun.py

 # RUN: %PYTHON %s 2>&1 | FileCheck %s
 
 import sys
 from mlir.ir import *
 from mlir.dialects import builtin
 from mlir.dialects import linalg
 from mlir.dialects import std
 from mlir.passmanager import *
 from mlir.execution_engine import *
 
 # Log everything to stderr and flush so that we have a unified stream to match
 # errors/info emitted by MLIR to stderr.
 def log(*args):
   print(*args, file=sys.stderr)
   sys.stderr.flush()
 
 boilerplate = """
 func @main() -> f32 attributes {llvm.emit_c_interface} {
   %v0 = constant 0.0 : f32
   %v1 = constant 1.0 : f32
   %v2 = constant 2.0 : f32
 
   %A = memref.alloc() : memref<4x16xf32>
   %B = memref.alloc() : memref<16x8xf32>
   %C = memref.alloc() : memref<4x8xf32>
   linalg.fill(%A, %v1) : memref<4x16xf32>, f32
   linalg.fill(%B, %v2) : memref<16x8xf32>, f32
   linalg.fill(%C, %v0) : memref<4x8xf32>, f32
 
   call @matmul_on_buffers(%A, %B, %C) : 
     (memref<4x16xf32>, memref<16x8xf32>, memref<4x8xf32>) -> ()
 
   %c0 = constant 0 : index
   %0 = memref.load %C[%c0, %c0] : memref<4x8xf32>
 
   // TODO: FFI-based solution to allow testing and printing with python code.
   return %0 : f32
 }
 """
 
 def transform(module):
   import mlir.conversions
   import mlir.dialects.linalg.passes
   import mlir.transforms
 
   # TODO: Allow cloning functions from one module to another.
   # Atm we have to resort to string concatenation.
   mod = Module.parse(
     str(module.operation.regions[0].blocks[0].operations[0].operation) +
     boilerplate)
   pm = PassManager.parse("func(convert-linalg-to-loops, convert-scf-to-std)," + 
                          "convert-vector-to-llvm," + 
                          "convert-std-to-llvm")
   pm.run(mod)
   return mod
 
 def test_builtin():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
     with InsertionPoint(module.body):
       @builtin.FuncOp.from_py_func(MemRefType.get((4, 16), f32),
                                    MemRefType.get((16, 8), f32),
                                    MemRefType.get((4, 8), f32))
       def matmul_on_buffers(lhs, rhs, out):
         linalg.matmul(lhs, rhs, outs=[out])
     
     execution_engine = ExecutionEngine(transform(module))
 
     # TODO: FFI-based solution to allow testing and printing with python code.
     # Prepare arguments: one result f32.
     # Arguments must be passed as pointers.
     c_float_p = ctypes.c_float * 1
     res = c_float_p(-1.)
     execution_engine.invoke("main", res)
 
     log('RESULT: ', res[0])
     # CHECK: RESULT: 32.0
 
 test_builtin()
 
 def test_generic():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
     with InsertionPoint(module.body):
       @builtin.FuncOp.from_py_func(MemRefType.get((4, 16), f32),
                                    MemRefType.get((16, 8), f32),
                                    MemRefType.get((4, 8), f32))
       def matmul_on_buffers(lhs, rhs, out):
         linalg.matmul(lhs, rhs, outs=[out], emit_generic=True)
     
     execution_engine = ExecutionEngine(transform(module))
 
     # TODO: FFI-based solution to allow testing and printing with python code.
     # Prepare arguments: one result f32.
     # Arguments must be passed as pointers.
     c_float_p = ctypes.c_float * 1
     res = c_float_p(-1.)
     execution_engine.invoke("main", res)
 
     log('RESULT: ', res[0])
     # CHECK: RESULT: 32.0
 
 test_generic()

mlir/test/python/dialects/memref.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 import mlir.dialects.std as std
 import mlir.dialects.memref as memref
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
 
 # CHECK-LABEL: TEST: testSubViewAccessors
 def testSubViewAccessors():
   ctx = Context()
   module = Module.parse(r"""
     func @f1(%arg0: memref<?x?xf32>) {
       %0 = constant 0 : index
       %1 = constant 1 : index
       %2 = constant 2 : index
       %3 = constant 3 : index
       %4 = constant 4 : index
       %5 = constant 5 : index
       memref.subview %arg0[%0, %1][%2, %3][%4, %5] : memref<?x?xf32> to memref<?x?xf32, offset: ?, strides: [?, ?]>
       return
     }
   """, ctx)
   func_body = module.body.operations[0].regions[0].blocks[0]
   subview = func_body.operations[6]
 
   assert subview.source == subview.operands[0]
   assert len(subview.offsets) == 2
   assert len(subview.sizes) == 2
   assert len(subview.strides) == 2
   assert subview.result == subview.results[0]
 
   # CHECK: SubViewOp
   print(type(subview).__name__)
 
   # CHECK: constant 0
   print(subview.offsets[0])
   # CHECK: constant 1
   print(subview.offsets[1])
   # CHECK: constant 2
   print(subview.sizes[0])
   # CHECK: constant 3
   print(subview.sizes[1])
   # CHECK: constant 4
   print(subview.strides[0])
   # CHECK: constant 5
   print(subview.strides[1])
 
 
 run(testSubViewAccessors)
 

mlir/test/python/dialects/ods_helpers.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
 
 
 def add_dummy_value():
   return Operation.create(
       "custom.value",
       results=[IntegerType.get_signless(32)]).result
 
 
 def testOdsBuildDefaultImplicitRegions():
 
   class TestFixedRegionsOp(OpView):
     OPERATION_NAME = "custom.test_op"
     _ODS_REGIONS = (2, True)
 
   class TestVariadicRegionsOp(OpView):
     OPERATION_NAME = "custom.test_any_regions_op"
     _ODS_REGIONS = (2, False)
 
   with Context() as ctx, Location.unknown():
     ctx.allow_unregistered_dialects = True
     m = Module.create()
     with InsertionPoint(m.body):
       op = TestFixedRegionsOp.build_generic(results=[], operands=[])
       # CHECK: NUM_REGIONS: 2
       print(f"NUM_REGIONS: {len(op.regions)}")
       # Including a regions= that matches should be fine.
       op = TestFixedRegionsOp.build_generic(results=[], operands=[], regions=2)
       print(f"NUM_REGIONS: {len(op.regions)}")
       # Reject greater than.
       try:
         op = TestFixedRegionsOp.build_generic(results=[], operands=[], regions=3)
       except ValueError as e:
         # CHECK: ERROR:Operation "custom.test_op" requires a maximum of 2 regions but was built with regions=3
         print(f"ERROR:{e}")
       # Reject less than.
       try:
         op = TestFixedRegionsOp.build_generic(results=[], operands=[], regions=1)
       except ValueError as e:
         # CHECK: ERROR:Operation "custom.test_op" requires a minimum of 2 regions but was built with regions=1
         print(f"ERROR:{e}")
 
       # If no regions specified for a variadic region op, build the minimum.
       op = TestVariadicRegionsOp.build_generic(results=[], operands=[])
       # CHECK: DEFAULT_NUM_REGIONS: 2
       print(f"DEFAULT_NUM_REGIONS: {len(op.regions)}")
       # Should also accept an explicit regions= that matches the minimum.
       op = TestVariadicRegionsOp.build_generic(
           results=[], operands=[], regions=2)
       # CHECK: EQ_NUM_REGIONS: 2
       print(f"EQ_NUM_REGIONS: {len(op.regions)}")
       # And accept greater than minimum.
       # Should also accept an explicit regions= that matches the minimum.
       op = TestVariadicRegionsOp.build_generic(
           results=[], operands=[], regions=3)
       # CHECK: GT_NUM_REGIONS: 3
       print(f"GT_NUM_REGIONS: {len(op.regions)}")
       # Should reject less than minimum.
       try:
         op = TestVariadicRegionsOp.build_generic(results=[], operands=[], regions=1)
       except ValueError as e:
         # CHECK: ERROR:Operation "custom.test_any_regions_op" requires a minimum of 2 regions but was built with regions=1
         print(f"ERROR:{e}")
 
 
 
 run(testOdsBuildDefaultImplicitRegions)
 
 
 def testOdsBuildDefaultNonVariadic():
 
   class TestOp(OpView):
     OPERATION_NAME = "custom.test_op"
 
   with Context() as ctx, Location.unknown():
     ctx.allow_unregistered_dialects = True
     m = Module.create()
     with InsertionPoint(m.body):
       v0 = add_dummy_value()
       v1 = add_dummy_value()
       t0 = IntegerType.get_signless(8)
       t1 = IntegerType.get_signless(16)
       op = TestOp.build_generic(results=[t0, t1], operands=[v0, v1])
       # CHECK: %[[V0:.+]] = "custom.value"
       # CHECK: %[[V1:.+]] = "custom.value"
       # CHECK: "custom.test_op"(%[[V0]], %[[V1]])
       # CHECK-NOT: operand_segment_sizes
       # CHECK-NOT: result_segment_sizes
       # CHECK-SAME: : (i32, i32) -> (i8, i16)
       print(m)
 
 run(testOdsBuildDefaultNonVariadic)
 
 
 def testOdsBuildDefaultSizedVariadic():
 
   class TestOp(OpView):
     OPERATION_NAME = "custom.test_op"
     _ODS_OPERAND_SEGMENTS = [1, -1, 0]
     _ODS_RESULT_SEGMENTS = [-1, 0, 1]
 
   with Context() as ctx, Location.unknown():
     ctx.allow_unregistered_dialects = True
     m = Module.create()
     with InsertionPoint(m.body):
       v0 = add_dummy_value()
       v1 = add_dummy_value()
       v2 = add_dummy_value()
       v3 = add_dummy_value()
       t0 = IntegerType.get_signless(8)
       t1 = IntegerType.get_signless(16)
       t2 = IntegerType.get_signless(32)
       t3 = IntegerType.get_signless(64)
       # CHECK: %[[V0:.+]] = "custom.value"
       # CHECK: %[[V1:.+]] = "custom.value"
       # CHECK: %[[V2:.+]] = "custom.value"
       # CHECK: %[[V3:.+]] = "custom.value"
       # CHECK: "custom.test_op"(%[[V0]], %[[V1]], %[[V2]], %[[V3]])
       # CHECK-SAME: operand_segment_sizes = dense<[1, 2, 1]> : vector<3xi32>
       # CHECK-SAME: result_segment_sizes = dense<[2, 1, 1]> : vector<3xi32>
       # CHECK-SAME: : (i32, i32, i32, i32) -> (i8, i16, i32, i64)
       op = TestOp.build_generic(
           results=[[t0, t1], t2, t3],
           operands=[v0, [v1, v2], v3])
 
       # Now test with optional omitted.
       # CHECK: "custom.test_op"(%[[V0]])
       # CHECK-SAME: operand_segment_sizes = dense<[1, 0, 0]>
       # CHECK-SAME: result_segment_sizes = dense<[0, 0, 1]>
       # CHECK-SAME: (i32) -> i64
       op = TestOp.build_generic(
           results=[None, None, t3],
           operands=[v0, None, None])
       print(m)
 
       # And verify that errors are raised for None in a required operand.
       try:
         op = TestOp.build_generic(
             results=[None, None, t3],
             operands=[None, None, None])
       except ValueError as e:
         # CHECK: OPERAND_CAST_ERROR:Operand 0 of operation "custom.test_op" must be a Value (was None and operand is not optional)
         print(f"OPERAND_CAST_ERROR:{e}")
 
       # And verify that errors are raised for None in a required result.
       try:
         op = TestOp.build_generic(
             results=[None, None, None],
             operands=[v0, None, None])
       except ValueError as e:
         # CHECK: RESULT_CAST_ERROR:Result 2 of operation "custom.test_op" must be a Type (was None and result is not optional)
         print(f"RESULT_CAST_ERROR:{e}")
 
       # Variadic lists with None elements should reject.
       try:
         op = TestOp.build_generic(
             results=[None, None, t3],
             operands=[v0, [None], None])
       except ValueError as e:
         # CHECK: OPERAND_LIST_CAST_ERROR:Operand 1 of operation "custom.test_op" must be a Sequence of Values (contained a None item)
         print(f"OPERAND_LIST_CAST_ERROR:{e}")
       try:
         op = TestOp.build_generic(
             results=[[None], None, t3],
             operands=[v0, None, None])
       except ValueError as e:
         # CHECK: RESULT_LIST_CAST_ERROR:Result 0 of operation "custom.test_op" must be a Sequence of Types (contained a None item)
         print(f"RESULT_LIST_CAST_ERROR:{e}")
 
 run(testOdsBuildDefaultSizedVariadic)
 
 
 def testOdsBuildDefaultCastError():
 
   class TestOp(OpView):
     OPERATION_NAME = "custom.test_op"
 
   with Context() as ctx, Location.unknown():
     ctx.allow_unregistered_dialects = True
     m = Module.create()
     with InsertionPoint(m.body):
       v0 = add_dummy_value()
       v1 = add_dummy_value()
       t0 = IntegerType.get_signless(8)
       t1 = IntegerType.get_signless(16)
       try:
         op = TestOp.build_generic(
             results=[t0, t1],
             operands=[None, v1])
       except ValueError as e:
         # CHECK: ERROR: Operand 0 of operation "custom.test_op" must be a Value
         print(f"ERROR: {e}")
       try:
         op = TestOp.build_generic(
             results=[t0, None],
             operands=[v0, v1])
       except ValueError as e:
         # CHECK: Result 1 of operation "custom.test_op" must be a Type
         print(f"ERROR: {e}")
 
 run(testOdsBuildDefaultCastError)

mlir/test/python/dialects/python_test.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 import mlir.dialects.python_test as test
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
 
 # CHECK-LABEL: TEST: testAttributes
 def testAttributes():
   with Context(), Location.unknown():
     #
     # Check op construction with attributes.
     #
 
     i32 = IntegerType.get_signless(32)
     one = IntegerAttr.get(i32, 1)
     two = IntegerAttr.get(i32, 2)
     unit = UnitAttr.get()
 
     # CHECK: "python_test.attributed_op"() {
     # CHECK-DAG: mandatory_i32 = 1 : i32
     # CHECK-DAG: optional_i32 = 2 : i32
     # CHECK-DAG: unit
     # CHECK: }
     op = test.AttributedOp(one, two, unit)
     print(f"{op}")
 
     # CHECK: "python_test.attributed_op"() {
     # CHECK: mandatory_i32 = 2 : i32
     # CHECK: }
     op2 = test.AttributedOp(two, None, None)
     print(f"{op2}")
 
     #
     # Check generic "attributes" access and mutation.
     #
 
     assert "additional" not in op.attributes
 
     # CHECK: "python_test.attributed_op"() {
     # CHECK-DAG: additional = 1 : i32
     # CHECK-DAG: mandatory_i32 = 2 : i32
     # CHECK: }
     op2.attributes["additional"] = one
     print(f"{op2}")
 
     # CHECK: "python_test.attributed_op"() {
     # CHECK-DAG: additional = 2 : i32
     # CHECK-DAG: mandatory_i32 = 2 : i32
     # CHECK: }
     op2.attributes["additional"] = two
     print(f"{op2}")
 
     # CHECK: "python_test.attributed_op"() {
     # CHECK-NOT: additional = 2 : i32
     # CHECK:     mandatory_i32 = 2 : i32
     # CHECK: }
     del op2.attributes["additional"]
     print(f"{op2}")
 
     try:
       print(op.attributes["additional"])
     except KeyError:
       pass
     else:
       assert False, "expected KeyError on unknown attribute key"
 
     #
     # Check accessors to defined attributes.
     #
 
     # CHECK: Mandatory: 1
     # CHECK: Optional: 2
     # CHECK: Unit: True
     print(f"Mandatory: {op.mandatory_i32.value}")
     print(f"Optional: {op.optional_i32.value}")
     print(f"Unit: {op.unit}")
 
     # CHECK: Mandatory: 2
     # CHECK: Optional: None
     # CHECK: Unit: False
     print(f"Mandatory: {op2.mandatory_i32.value}")
     print(f"Optional: {op2.optional_i32}")
     print(f"Unit: {op2.unit}")
 
     # CHECK: Mandatory: 2
     # CHECK: Optional: None
     # CHECK: Unit: False
     op.mandatory_i32 = two
     op.optional_i32 = None
     op.unit = False
     print(f"Mandatory: {op.mandatory_i32.value}")
     print(f"Optional: {op.optional_i32}")
     print(f"Unit: {op.unit}")
     assert "optional_i32" not in op.attributes
     assert "unit" not in op.attributes
 
     try:
       op.mandatory_i32 = None
     except ValueError:
       pass
     else:
       assert False, "expected ValueError on setting a mandatory attribute to None"
 
     # CHECK: Optional: 2
     op.optional_i32 = two
     print(f"Optional: {op.optional_i32.value}")
 
     # CHECK: Optional: None
     del op.optional_i32
     print(f"Optional: {op.optional_i32}")
 
     # CHECK: Unit: False
     op.unit = None
     print(f"Unit: {op.unit}")
     assert "unit" not in op.attributes
 
     # CHECK: Unit: True
     op.unit = True
     print(f"Unit: {op.unit}")
 
     # CHECK: Unit: False
     del op.unit
     print(f"Unit: {op.unit}")
 
 run(testAttributes)

mlir/test/python/execution_engine.py

 # RUN: %PYTHON %s 2>&1 | FileCheck %s
 
 import gc, sys
 from mlir.ir import *
 from mlir.passmanager import *
 from mlir.execution_engine import *
 from mlir.runtime import *
 
 # Log everything to stderr and flush so that we have a unified stream to match
 # errors/info emitted by MLIR to stderr.
 def log(*args):
   print(*args, file=sys.stderr)
   sys.stderr.flush()
 
 def run(f):
   log("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
 
 # Verify capsule interop.
 # CHECK-LABEL: TEST: testCapsule
 def testCapsule():
   with Context():
     module = Module.parse(r"""
 llvm.func @none() {
   llvm.return
 }
     """)
     execution_engine = ExecutionEngine(module)
     execution_engine_capsule = execution_engine._CAPIPtr
     # CHECK: mlir.execution_engine.ExecutionEngine._CAPIPtr
     log(repr(execution_engine_capsule))
     execution_engine._testing_release()
     execution_engine1 = ExecutionEngine._CAPICreate(execution_engine_capsule)
     # CHECK: _mlir.execution_engine.ExecutionEngine
     log(repr(execution_engine1))
 
 run(testCapsule)
 
 # Test invalid ExecutionEngine creation
 # CHECK-LABEL: TEST: testInvalidModule
 def testInvalidModule():
   with Context():
     # Builtin function
     module = Module.parse(r"""
     func @foo() { return }
     """)
     # CHECK: Got RuntimeError:  Failure while creating the ExecutionEngine.
     try:
       execution_engine = ExecutionEngine(module)
     except RuntimeError as e:
       log("Got RuntimeError: ", e)
 
 run(testInvalidModule)
 
 def lowerToLLVM(module):
   import mlir.conversions
   pm = PassManager.parse("convert-std-to-llvm")
   pm.run(module)
   return module
 
 # Test simple ExecutionEngine execution
 # CHECK-LABEL: TEST: testInvokeVoid
 def testInvokeVoid():
   with Context():
     module = Module.parse(r"""
 func @void() attributes { llvm.emit_c_interface } {
   return
 }
     """)    
     execution_engine = ExecutionEngine(lowerToLLVM(module))
     # Nothing to check other than no exception thrown here.
     execution_engine.invoke("void")
 
 run(testInvokeVoid)
 
 
 # Test argument passing and result with a simple float addition.
 # CHECK-LABEL: TEST: testInvokeFloatAdd
 def testInvokeFloatAdd():
   with Context():
     module = Module.parse(r"""
 func @add(%arg0: f32, %arg1: f32) -> f32 attributes { llvm.emit_c_interface } {
   %add = std.addf %arg0, %arg1 : f32
   return %add : f32
 }
     """)
     execution_engine = ExecutionEngine(lowerToLLVM(module))
     # Prepare arguments: two input floats and one result.
     # Arguments must be passed as pointers.
     c_float_p = ctypes.c_float * 1
     arg0 = c_float_p(42.)
     arg1 = c_float_p(2.)
     res = c_float_p(-1.)
     execution_engine.invoke("add", arg0, arg1, res)
     # CHECK: 42.0 + 2.0 = 44.0
     log("{0} + {1} = {2}".format(arg0[0], arg1[0], res[0]))
 
 run(testInvokeFloatAdd)
 
 
 # Test callback
 # CHECK-LABEL: TEST: testBasicCallback
 def testBasicCallback():
   # Define a callback function that takes a float and an integer and returns a float.
   @ctypes.CFUNCTYPE(ctypes.c_float, ctypes.c_float, ctypes.c_int)
   def callback(a, b):
     return a/2 + b/2
 
   with Context():
     # The module just forwards to a runtime function known as "some_callback_into_python".
     module = Module.parse(r"""
 func @add(%arg0: f32, %arg1: i32) -> f32 attributes { llvm.emit_c_interface } {
   %resf = call @some_callback_into_python(%arg0, %arg1) : (f32, i32) -> (f32)
   return %resf : f32
 }
 func private @some_callback_into_python(f32, i32) -> f32 attributes { llvm.emit_c_interface }
     """)
     execution_engine = ExecutionEngine(lowerToLLVM(module))
     execution_engine.register_runtime("some_callback_into_python", callback)
 
     # Prepare arguments: two input floats and one result.
     # Arguments must be passed as pointers.
     c_float_p = ctypes.c_float * 1
     c_int_p = ctypes.c_int * 1
     arg0 = c_float_p(42.)
     arg1 = c_int_p(2)
     res = c_float_p(-1.)
     execution_engine.invoke("add", arg0, arg1, res)
     # CHECK: 42.0 + 2 = 44.0
     log("{0} + {1} = {2}".format(arg0[0], arg1[0], res[0]*2))
 
 run(testBasicCallback)
 
 # Test callback with an unranked memref
 # CHECK-LABEL: TEST: testUnrankedMemRefCallback
 def testUnrankedMemRefCallback():
     # Define a callback function that takes an unranked memref, converts it to a numpy array and prints it.
     @ctypes.CFUNCTYPE(None, ctypes.POINTER(UnrankedMemRefDescriptor))
     def callback(a):
         arr = unranked_memref_to_numpy(a, np.float32)
         log("Inside callback: ")
         log(arr)
 
     with Context():
         # The module just forwards to a runtime function known as "some_callback_into_python".
         module = Module.parse(
             r"""
 func @callback_memref(%arg0: memref<*xf32>) attributes { llvm.emit_c_interface } {
   call @some_callback_into_python(%arg0) : (memref<*xf32>) -> ()
   return
 }
 func private @some_callback_into_python(memref<*xf32>) -> () attributes { llvm.emit_c_interface }
 """
         )
         execution_engine = ExecutionEngine(lowerToLLVM(module))
         execution_engine.register_runtime("some_callback_into_python", callback)
         inp_arr = np.array([[1.0, 2.0], [3.0, 4.0]], np.float32)
         # CHECK: Inside callback: 
         # CHECK{LITERAL}: [[1. 2.]
         # CHECK{LITERAL}:  [3. 4.]]
         execution_engine.invoke(
             "callback_memref",
             ctypes.pointer(ctypes.pointer(get_unranked_memref_descriptor(inp_arr))),
         )
         inp_arr_1 = np.array([5, 6, 7], dtype=np.float32)
         strided_arr = np.lib.stride_tricks.as_strided(
             inp_arr_1, strides=(4, 0), shape=(3, 4)
         )
         # CHECK: Inside callback: 
         # CHECK{LITERAL}: [[5. 5. 5. 5.]
         # CHECK{LITERAL}:  [6. 6. 6. 6.]
         # CHECK{LITERAL}:  [7. 7. 7. 7.]]
         execution_engine.invoke(
             "callback_memref",
             ctypes.pointer(
                 ctypes.pointer(get_unranked_memref_descriptor(strided_arr))
             ),
         )
 
 run(testUnrankedMemRefCallback)
 
 # Test callback with a ranked memref.
 # CHECK-LABEL: TEST: testRankedMemRefCallback
 def testRankedMemRefCallback():
     # Define a callback function that takes a ranked memref, converts it to a numpy array and prints it.
     @ctypes.CFUNCTYPE(
         None,
         ctypes.POINTER(
             make_nd_memref_descriptor(2, np.ctypeslib.as_ctypes_type(np.float32))
         ),
     )
     def callback(a):
         arr = ranked_memref_to_numpy(a)
         log("Inside Callback: ")
         log(arr)
 
     with Context():
         # The module just forwards to a runtime function known as "some_callback_into_python".
         module = Module.parse(
             r"""
 func @callback_memref(%arg0: memref<2x2xf32>) attributes { llvm.emit_c_interface } {
   call @some_callback_into_python(%arg0) : (memref<2x2xf32>) -> ()
   return
 }
 func private @some_callback_into_python(memref<2x2xf32>) -> () attributes { llvm.emit_c_interface }
 """
         )
         execution_engine = ExecutionEngine(lowerToLLVM(module))
         execution_engine.register_runtime("some_callback_into_python", callback)
         inp_arr = np.array([[1.0, 5.0], [6.0, 7.0]], np.float32)
         # CHECK: Inside Callback:
         # CHECK{LITERAL}: [[1. 5.]
         # CHECK{LITERAL}:  [6. 7.]]
         execution_engine.invoke(
             "callback_memref", ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(inp_arr)))
         )
 
 run(testRankedMemRefCallback)
 
 #  Test addition of two memref
 # CHECK-LABEL: TEST: testMemrefAdd
 def testMemrefAdd():
     with Context():
         module = Module.parse(
             """
       module  {
       func @main(%arg0: memref<1xf32>, %arg1: memref<f32>, %arg2: memref<1xf32>) attributes { llvm.emit_c_interface } {
         %0 = constant 0 : index
         %1 = memref.load %arg0[%0] : memref<1xf32>
         %2 = memref.load %arg1[] : memref<f32>
         %3 = addf %1, %2 : f32
         memref.store %3, %arg2[%0] : memref<1xf32>
         return
       }
      } """
         )
         arg1 = np.array([32.5]).astype(np.float32)
         arg2 = np.array(6).astype(np.float32)
         res = np.array([0]).astype(np.float32)
 
         arg1_memref_ptr = ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(arg1)))
         arg2_memref_ptr = ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(arg2)))
         res_memref_ptr = ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(res)))
 
         execution_engine = ExecutionEngine(lowerToLLVM(module))
         execution_engine.invoke(
             "main", arg1_memref_ptr, arg2_memref_ptr, res_memref_ptr
         )
         # CHECK: [32.5] + 6.0 = [38.5]
         log("{0} + {1} = {2}".format(arg1, arg2, res))
 
 run(testMemrefAdd)
 
 #  Test addition of two 2d_memref
 # CHECK-LABEL: TEST: testDynamicMemrefAdd2D
 def testDynamicMemrefAdd2D():
     with Context():
         module = Module.parse(
       """
       module  {
         func @memref_add_2d(%arg0: memref<2x2xf32>, %arg1: memref<?x?xf32>, %arg2: memref<2x2xf32>) attributes {llvm.emit_c_interface} {
           %c0 = constant 0 : index
           %c2 = constant 2 : index
           %c1 = constant 1 : index
           br ^bb1(%c0 : index)
         ^bb1(%0: index):  // 2 preds: ^bb0, ^bb5
           %1 = cmpi slt, %0, %c2 : index
           cond_br %1, ^bb2, ^bb6
         ^bb2:  // pred: ^bb1
           %c0_0 = constant 0 : index
           %c2_1 = constant 2 : index
           %c1_2 = constant 1 : index
           br ^bb3(%c0_0 : index)
         ^bb3(%2: index):  // 2 preds: ^bb2, ^bb4
           %3 = cmpi slt, %2, %c2_1 : index
           cond_br %3, ^bb4, ^bb5
         ^bb4:  // pred: ^bb3
           %4 = memref.load %arg0[%0, %2] : memref<2x2xf32>
           %5 = memref.load %arg1[%0, %2] : memref<?x?xf32>
           %6 = addf %4, %5 : f32
           memref.store %6, %arg2[%0, %2] : memref<2x2xf32>
           %7 = addi %2, %c1_2 : index
           br ^bb3(%7 : index)
         ^bb5:  // pred: ^bb3
           %8 = addi %0, %c1 : index
           br ^bb1(%8 : index)
         ^bb6:  // pred: ^bb1
           return
         }
       }
         """
         )
         arg1 = np.random.randn(2,2).astype(np.float32)
         arg2 = np.random.randn(2,2).astype(np.float32)
         res = np.random.randn(2,2).astype(np.float32)
 
         arg1_memref_ptr = ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(arg1)))
         arg2_memref_ptr = ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(arg2)))
         res_memref_ptr = ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(res)))
 
         execution_engine = ExecutionEngine(lowerToLLVM(module))
         execution_engine.invoke(
             "memref_add_2d", arg1_memref_ptr, arg2_memref_ptr, res_memref_ptr
         )
         # CHECK: True
         log(np.allclose(arg1+arg2, res))
 
 run(testDynamicMemrefAdd2D)

mlir/test/python/ir/affine_expr.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
 
 
 # CHECK-LABEL: TEST: testAffineExprCapsule
 def testAffineExprCapsule():
   with Context() as ctx:
     affine_expr = AffineExpr.get_constant(42)
 
   affine_expr_capsule = affine_expr._CAPIPtr
   # CHECK: capsule object
   # CHECK: mlir.ir.AffineExpr._CAPIPtr
   print(affine_expr_capsule)
 
   affine_expr_2 = AffineExpr._CAPICreate(affine_expr_capsule)
   assert affine_expr == affine_expr_2
   assert affine_expr_2.context == ctx
 
 run(testAffineExprCapsule)
 
 
 # CHECK-LABEL: TEST: testAffineExprEq
 def testAffineExprEq():
   with Context():
     a1 = AffineExpr.get_constant(42)
     a2 = AffineExpr.get_constant(42)
     a3 = AffineExpr.get_constant(43)
     # CHECK: True
     print(a1 == a1)
     # CHECK: True
     print(a1 == a2)
     # CHECK: False
     print(a1 == a3)
     # CHECK: False
     print(a1 == None)
     # CHECK: False
     print(a1 == "foo")
 
 run(testAffineExprEq)
 
 
 # CHECK-LABEL: TEST: testAffineExprContext
 def testAffineExprContext():
   with Context():
     a1 = AffineExpr.get_constant(42)
   with Context():
     a2 = AffineExpr.get_constant(42)
 
   # CHECK: False
   print(a1 == a2)
 
 run(testAffineExprContext)
 
 
 # CHECK-LABEL: TEST: testAffineExprConstant
 def testAffineExprConstant():
   with Context():
     a1 = AffineExpr.get_constant(42)
     # CHECK: 42
     print(a1.value)
     # CHECK: 42
     print(a1)
 
     a2 = AffineConstantExpr.get(42)
     # CHECK: 42
     print(a2.value)
     # CHECK: 42
     print(a2)
 
     assert a1 == a2
 
 run(testAffineExprConstant)
 
 
 # CHECK-LABEL: TEST: testAffineExprDim
 def testAffineExprDim():
   with Context():
     d1 = AffineExpr.get_dim(1)
     d11 = AffineDimExpr.get(1)
     d2 = AffineDimExpr.get(2)
 
     # CHECK: 1
     print(d1.position)
     # CHECK: d1
     print(d1)
 
     # CHECK: 2
     print(d2.position)
     # CHECK: d2
     print(d2)
 
     assert d1 == d11
     assert d1 != d2
 
 run(testAffineExprDim)
 
 
 # CHECK-LABEL: TEST: testAffineExprSymbol
 def testAffineExprSymbol():
   with Context():
     s1 = AffineExpr.get_symbol(1)
     s11 = AffineSymbolExpr.get(1)
     s2 = AffineSymbolExpr.get(2)
 
     # CHECK: 1
     print(s1.position)
     # CHECK: s1
     print(s1)
 
     # CHECK: 2
     print(s2.position)
     # CHEKC: s2
     print(s2)
 
     assert s1 == s11
     assert s1 != s2
 
 run(testAffineExprSymbol)
 
 
 # CHECK-LABEL: TEST: testAffineAddExpr
 def testAffineAddExpr():
   with Context():
     d1 = AffineDimExpr.get(1)
     d2 = AffineDimExpr.get(2)
     d12 = AffineExpr.get_add(d1, d2)
     # CHECK: d1 + d2
     print(d12)
 
     d12op = d1 + d2
     # CHECK: d1 + d2
     print(d12op)
 
     assert d12 == d12op
     assert d12.lhs == d1
     assert d12.rhs == d2
 
 run(testAffineAddExpr)
 
 
 # CHECK-LABEL: TEST: testAffineMulExpr
 def testAffineMulExpr():
   with Context():
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     expr = AffineExpr.get_mul(d1, c2)
     # CHECK: d1 * 2
     print(expr)
 
     # CHECK: d1 * 2
     op = d1 * c2
     print(op)
 
     assert expr == op
     assert expr.lhs == d1
     assert expr.rhs == c2
 
 run(testAffineMulExpr)
 
 
 # CHECK-LABEL: TEST: testAffineModExpr
 def testAffineModExpr():
   with Context():
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     expr = AffineExpr.get_mod(d1, c2)
     # CHECK: d1 mod 2
     print(expr)
 
     # CHECK: d1 mod 2
     op = d1 % c2
     print(op)
 
     assert expr == op
     assert expr.lhs == d1
     assert expr.rhs == c2
 
 run(testAffineModExpr)
 
 
 # CHECK-LABEL: TEST: testAffineFloorDivExpr
 def testAffineFloorDivExpr():
   with Context():
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     expr = AffineExpr.get_floor_div(d1, c2)
     # CHECK: d1 floordiv 2
     print(expr)
 
     assert expr.lhs == d1
     assert expr.rhs == c2
 
 run(testAffineFloorDivExpr)
 
 
 # CHECK-LABEL: TEST: testAffineCeilDivExpr
 def testAffineCeilDivExpr():
   with Context():
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     expr = AffineExpr.get_ceil_div(d1, c2)
     # CHECK: d1 ceildiv 2
     print(expr)
 
     assert expr.lhs == d1
     assert expr.rhs == c2
 
 run(testAffineCeilDivExpr)
 
 
 # CHECK-LABEL: TEST: testAffineExprSub
 def testAffineExprSub():
   with Context():
     d1 = AffineDimExpr.get(1)
     d2 = AffineDimExpr.get(2)
     expr = d1 - d2
     # CHECK: d1 - d2
     print(expr)
 
     assert expr.lhs == d1
     rhs = AffineMulExpr(expr.rhs)
     # CHECK: d2
     print(rhs.lhs)
     # CHECK: -1
     print(rhs.rhs)
 
 run(testAffineExprSub)
 
 
 def testClassHierarchy():
   with Context():
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     add = AffineAddExpr.get(d1, c2)
     mul = AffineMulExpr.get(d1, c2)
     mod = AffineModExpr.get(d1, c2)
     floor_div = AffineFloorDivExpr.get(d1, c2)
     ceil_div = AffineCeilDivExpr.get(d1, c2)
 
     # CHECK: False
     print(isinstance(d1, AffineBinaryExpr))
     # CHECK: False
     print(isinstance(c2, AffineBinaryExpr))
     # CHECK: True
     print(isinstance(add, AffineBinaryExpr))
     # CHECK: True
     print(isinstance(mul, AffineBinaryExpr))
     # CHECK: True
     print(isinstance(mod, AffineBinaryExpr))
     # CHECK: True
     print(isinstance(floor_div, AffineBinaryExpr))
     # CHECK: True
     print(isinstance(ceil_div, AffineBinaryExpr))
 
     try:
       AffineBinaryExpr(d1)
     except ValueError as e:
       # CHECK: Cannot cast affine expression to AffineBinaryExpr
       print(e)
 
     try:
       AffineBinaryExpr(c2)
     except ValueError as e:
       # CHECK: Cannot cast affine expression to AffineBinaryExpr
       print(e)
 
 run(testClassHierarchy)

mlir/test/python/ir/affine_map.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
 
 
 # CHECK-LABEL: TEST: testAffineMapCapsule
 def testAffineMapCapsule():
   with Context() as ctx:
     am1 = AffineMap.get_empty(ctx)
   # CHECK: mlir.ir.AffineMap._CAPIPtr
   affine_map_capsule = am1._CAPIPtr
   print(affine_map_capsule)
   am2 = AffineMap._CAPICreate(affine_map_capsule)
   assert am2 == am1
   assert am2.context is ctx
 
 run(testAffineMapCapsule)
 
 
 # CHECK-LABEL: TEST: testAffineMapGet
 def testAffineMapGet():
   with Context() as ctx:
     d0 = AffineDimExpr.get(0)
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
 
     # CHECK: (d0, d1)[s0, s1, s2] -> ()
     map0 = AffineMap.get(2, 3, [])
     print(map0)
 
     # CHECK: (d0, d1)[s0, s1, s2] -> (d1, 2)
     map1 = AffineMap.get(2, 3, [d1, c2])
     print(map1)
 
     # CHECK: () -> (2)
     map2 = AffineMap.get(0, 0, [c2])
     print(map2)
 
     # CHECK: (d0, d1) -> (d0, d1)
     map3 = AffineMap.get(2, 0, [d0, d1])
     print(map3)
 
     # CHECK: (d0, d1) -> (d1)
     map4 = AffineMap.get(2, 0, [d1])
     print(map4)
 
     # CHECK: (d0, d1, d2) -> (d2, d0, d1)
     map5 = AffineMap.get_permutation([2, 0, 1])
     print(map5)
 
     assert map1 == AffineMap.get(2, 3, [d1, c2])
     assert AffineMap.get(0, 0, []) == AffineMap.get_empty()
     assert map2 == AffineMap.get_constant(2)
     assert map3 == AffineMap.get_identity(2)
     assert map4 == AffineMap.get_minor_identity(2, 1)
 
     try:
       AffineMap.get(1, 1, [1])
     except RuntimeError as e:
       # CHECK: Invalid expression when attempting to create an AffineMap
       print(e)
 
     try:
       AffineMap.get(1, 1, [None])
     except RuntimeError as e:
       # CHECK: Invalid expression (None?) when attempting to create an AffineMap
       print(e)
 
     try:
       AffineMap.get_permutation([1, 0, 1])
     except RuntimeError as e:
       # CHECK: Invalid permutation when attempting to create an AffineMap
       print(e)
 
     try:
       map3.get_submap([42])
     except ValueError as e:
       # CHECK: result position out of bounds
       print(e)
 
     try:
       map3.get_minor_submap(42)
     except ValueError as e:
       # CHECK: number of results out of bounds
       print(e)
 
     try:
       map3.get_major_submap(42)
     except ValueError as e:
       # CHECK: number of results out of bounds
       print(e)
 
 run(testAffineMapGet)
 
 
 # CHECK-LABEL: TEST: testAffineMapDerive
 def testAffineMapDerive():
   with Context() as ctx:
     map5 = AffineMap.get_identity(5)
 
     # CHECK: (d0, d1, d2, d3, d4) -> (d1, d2, d3)
     map123 = map5.get_submap([1, 2, 3])
     print(map123)
 
     # CHECK: (d0, d1, d2, d3, d4) -> (d0, d1)
     map01 = map5.get_major_submap(2)
     print(map01)
 
     # CHECK: (d0, d1, d2, d3, d4) -> (d3, d4)
     map34 = map5.get_minor_submap(2)
     print(map34)
 
 run(testAffineMapDerive)
 
 
 # CHECK-LABEL: TEST: testAffineMapProperties
 def testAffineMapProperties():
   with Context():
     d0 = AffineDimExpr.get(0)
     d1 = AffineDimExpr.get(1)
     d2 = AffineDimExpr.get(2)
     map1 = AffineMap.get(3, 0, [d2, d0])
     map2 = AffineMap.get(3, 0, [d2, d0, d1])
     map3 = AffineMap.get(3, 1, [d2, d0, d1])
     # CHECK: False
     print(map1.is_permutation)
     # CHECK: True
     print(map1.is_projected_permutation)
     # CHECK: True
     print(map2.is_permutation)
     # CHECK: True
     print(map2.is_projected_permutation)
     # CHECK: False
     print(map3.is_permutation)
     # CHECK: False
     print(map3.is_projected_permutation)
 
 run(testAffineMapProperties)
 
 
 # CHECK-LABEL: TEST: testAffineMapExprs
 def testAffineMapExprs():
   with Context():
     d0 = AffineDimExpr.get(0)
     d1 = AffineDimExpr.get(1)
     d2 = AffineDimExpr.get(2)
     map3 = AffineMap.get(3, 1, [d2, d0, d1])
 
     # CHECK: 3
     print(map3.n_dims)
     # CHECK: 4
     print(map3.n_inputs)
     # CHECK: 1
     print(map3.n_symbols)
     assert map3.n_inputs == map3.n_dims + map3.n_symbols
 
     # CHECK: 3
     print(len(map3.results))
     for expr in map3.results:
       # CHECK: d2
       # CHECK: d0
       # CHECK: d1
       print(expr)
     for expr in map3.results[-1:-4:-1]:
       # CHECK: d1
       # CHECK: d0
       # CHECK: d2
       print(expr)
     assert list(map3.results) == [d2, d0, d1]
 
 run(testAffineMapExprs)
 
 # CHECK-LABEL: TEST: testCompressUnusedSymbols
 def testCompressUnusedSymbols():
   with Context() as ctx:
     d0, d1, d2 = (
       AffineDimExpr.get(0), 
       AffineDimExpr.get(1), 
       AffineDimExpr.get(2))
     s0, s1, s2 = (
       AffineSymbolExpr.get(0), 
       AffineSymbolExpr.get(1), 
       AffineSymbolExpr.get(2))
     maps = [
         AffineMap.get(3, 3, [d2, d0, d1]),
         AffineMap.get(3, 3, [d2, d0 + s2, d1]),
         AffineMap.get(3, 3, [d1, d2, d0])]
 
     compressed_maps = AffineMap.compress_unused_symbols(maps, ctx)
 
     #      CHECK: AffineMap((d0, d1, d2)[s0, s1, s2] -> (d2, d0, d1))
     # CHECK-SAME: AffineMap((d0, d1, d2)[s0, s1, s2] -> (d2, d0 + s2, d1))
     # CHECK-SAME: AffineMap((d0, d1, d2)[s0, s1, s2] -> (d1, d2, d0))
     print(maps)
 
     #      CHECK: AffineMap((d0, d1, d2)[s0] -> (d2, d0, d1))
     # CHECK-SAME: AffineMap((d0, d1, d2)[s0] -> (d2, d0 + s0, d1))
     # CHECK-SAME: AffineMap((d0, d1, d2)[s0] -> (d1, d2, d0))
     print(compressed_maps)
 
 
 run(testCompressUnusedSymbols)

mlir/test/python/ir/array_attributes.py

 # RUN: %PYTHON %s | FileCheck %s
 # Note that this is separate from ir_attributes.py since it depends on numpy,
 # and we may want to disable if not available.
 
 import gc
 from mlir.ir import *
 import numpy as np
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
 
 ################################################################################
 # Tests of the array/buffer .get() factory method on unsupported dtype.
 ################################################################################
 
 def testGetDenseElementsUnsupported():
   with Context():
     array = np.array([["hello", "goodbye"]])
     try:
       attr = DenseElementsAttr.get(array)
     except ValueError as e:
       # CHECK: unimplemented array format conversion from format:
       print(e)
 
 run(testGetDenseElementsUnsupported)
 
 ################################################################################
 # Splats.
 ################################################################################
 
 # CHECK-LABEL: TEST: testGetDenseElementsSplatInt
 def testGetDenseElementsSplatInt():
   with Context(), Location.unknown():
     t = IntegerType.get_signless(32)
     element = IntegerAttr.get(t, 555)
     shaped_type = RankedTensorType.get((2, 3, 4), t)
     attr = DenseElementsAttr.get_splat(shaped_type, element)
     # CHECK: dense<555> : tensor<2x3x4xi32>
     print(attr)
     # CHECK: is_splat: True
     print("is_splat:", attr.is_splat)
 
 run(testGetDenseElementsSplatInt)
 
 
 # CHECK-LABEL: TEST: testGetDenseElementsSplatFloat
 def testGetDenseElementsSplatFloat():
   with Context(), Location.unknown():
     t = F32Type.get()
     element = FloatAttr.get(t, 1.2)
     shaped_type = RankedTensorType.get((2, 3, 4), t)
     attr = DenseElementsAttr.get_splat(shaped_type, element)
     # CHECK: dense<1.200000e+00> : tensor<2x3x4xf32>
     print(attr)
 
 run(testGetDenseElementsSplatFloat)
 
 
 # CHECK-LABEL: TEST: testGetDenseElementsSplatErrors
 def testGetDenseElementsSplatErrors():
   with Context(), Location.unknown():
     t = F32Type.get()
     other_t = F64Type.get()
     element = FloatAttr.get(t, 1.2)
     other_element = FloatAttr.get(other_t, 1.2)
     shaped_type = RankedTensorType.get((2, 3, 4), t)
     dynamic_shaped_type = UnrankedTensorType.get(t)
     non_shaped_type = t
 
     try:
       attr = DenseElementsAttr.get_splat(non_shaped_type, element)
     except ValueError as e:
       # CHECK: Expected a static ShapedType for the shaped_type parameter: Type(f32)
       print(e)
 
     try:
       attr = DenseElementsAttr.get_splat(dynamic_shaped_type, element)
     except ValueError as e:
       # CHECK: Expected a static ShapedType for the shaped_type parameter: Type(tensor<*xf32>)
       print(e)
 
     try:
       attr = DenseElementsAttr.get_splat(shaped_type, other_element)
     except ValueError as e:
       # CHECK: Shaped element type and attribute type must be equal: shaped=Type(tensor<2x3x4xf32>), element=Attribute(1.200000e+00 : f64)
       print(e)
 
 run(testGetDenseElementsSplatErrors)
 
 
 ################################################################################
 # Tests of the array/buffer .get() factory method, in all of its permutations.
 ################################################################################
 
 ### float and double arrays.
 
 # CHECK-LABEL: TEST: testGetDenseElementsF32
 def testGetDenseElementsF32():
   with Context():
     array = np.array([[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]], dtype=np.float32)
     attr = DenseElementsAttr.get(array)
     # CHECK: dense<{{\[}}[1.100000e+00, 2.200000e+00, 3.300000e+00], [4.400000e+00, 5.500000e+00, 6.600000e+00]]> : tensor<2x3xf32>
     print(attr)
     # CHECK: is_splat: False
     print("is_splat:", attr.is_splat)
     # CHECK: {{\[}}[1.1 2.2 3.3]
     # CHECK: {{\[}}4.4 5.5 6.6]]
     print(np.array(attr))
 
 run(testGetDenseElementsF32)
 
 
 # CHECK-LABEL: TEST: testGetDenseElementsF64
 def testGetDenseElementsF64():
   with Context():
     array = np.array([[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]], dtype=np.float64)
     attr = DenseElementsAttr.get(array)
     # CHECK: dense<{{\[}}[1.100000e+00, 2.200000e+00, 3.300000e+00], [4.400000e+00, 5.500000e+00, 6.600000e+00]]> : tensor<2x3xf64>
     print(attr)
     # CHECK: {{\[}}[1.1 2.2 3.3]
     # CHECK: {{\[}}4.4 5.5 6.6]]
     print(np.array(attr))
 
 run(testGetDenseElementsF64)
 
 
 ### 32 bit integer arrays
 # CHECK-LABEL: TEST: testGetDenseElementsI32Signless
 def testGetDenseElementsI32Signless():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32)
     attr = DenseElementsAttr.get(array)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32>
     print(attr)
     # CHECK: {{\[}}[1 2 3]
     # CHECK: {{\[}}4 5 6]]
     print(np.array(attr))
 
 run(testGetDenseElementsI32Signless)
 
 
 # CHECK-LABEL: TEST: testGetDenseElementsUI32Signless
 def testGetDenseElementsUI32Signless():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint32)
     attr = DenseElementsAttr.get(array)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32>
     print(attr)
     # CHECK: {{\[}}[1 2 3]
     # CHECK: {{\[}}4 5 6]]
     print(np.array(attr))
 
 run(testGetDenseElementsUI32Signless)
 
 # CHECK-LABEL: TEST: testGetDenseElementsI32
 def testGetDenseElementsI32():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32)
     attr = DenseElementsAttr.get(array, signless=False)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xsi32>
     print(attr)
     # CHECK: {{\[}}[1 2 3]
     # CHECK: {{\[}}4 5 6]]
     print(np.array(attr))
 
 run(testGetDenseElementsI32)
 
 
 # CHECK-LABEL: TEST: testGetDenseElementsUI32
 def testGetDenseElementsUI32():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint32)
     attr = DenseElementsAttr.get(array, signless=False)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xui32>
     print(attr)
     # CHECK: {{\[}}[1 2 3]
     # CHECK: {{\[}}4 5 6]]
     print(np.array(attr))
 
 run(testGetDenseElementsUI32)
 
 
 ## 64bit integer arrays
 # CHECK-LABEL: TEST: testGetDenseElementsI64Signless
 def testGetDenseElementsI64Signless():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64)
     attr = DenseElementsAttr.get(array)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi64>
     print(attr)
     # CHECK: {{\[}}[1 2 3]
     # CHECK: {{\[}}4 5 6]]
     print(np.array(attr))
 
 run(testGetDenseElementsI64Signless)
 
 
 # CHECK-LABEL: TEST: testGetDenseElementsUI64Signless
 def testGetDenseElementsUI64Signless():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint64)
     attr = DenseElementsAttr.get(array)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi64>
     print(attr)
     # CHECK: {{\[}}[1 2 3]
     # CHECK: {{\[}}4 5 6]]
     print(np.array(attr))
 
 run(testGetDenseElementsUI64Signless)
 
 # CHECK-LABEL: TEST: testGetDenseElementsI64
 def testGetDenseElementsI64():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64)
     attr = DenseElementsAttr.get(array, signless=False)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xsi64>
     print(attr)
     # CHECK: {{\[}}[1 2 3]
     # CHECK: {{\[}}4 5 6]]
     print(np.array(attr))
 
 run(testGetDenseElementsI64)
 
 
 # CHECK-LABEL: TEST: testGetDenseElementsUI64
 def testGetDenseElementsUI64():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint64)
     attr = DenseElementsAttr.get(array, signless=False)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xui64>
     print(attr)
     # CHECK: {{\[}}[1 2 3]
     # CHECK: {{\[}}4 5 6]]
     print(np.array(attr))
 
 run(testGetDenseElementsUI64)
 

mlir/test/python/ir/attributes.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
 
 
 # CHECK-LABEL: TEST: testParsePrint
 def testParsePrint():
   with Context() as ctx:
     t = Attribute.parse('"hello"')
   assert t.context is ctx
   ctx = None
   gc.collect()
   # CHECK: "hello"
   print(str(t))
   # CHECK: Attribute("hello")
   print(repr(t))
 
 run(testParsePrint)
 
 
 # CHECK-LABEL: TEST: testParseError
 # TODO: Hook the diagnostic manager to capture a more meaningful error
 # message.
 def testParseError():
   with Context():
     try:
       t = Attribute.parse("BAD_ATTR_DOES_NOT_EXIST")
     except ValueError as e:
       # CHECK: Unable to parse attribute: 'BAD_ATTR_DOES_NOT_EXIST'
       print("testParseError:", e)
     else:
       print("Exception not produced")
 
 run(testParseError)
 
 
 # CHECK-LABEL: TEST: testAttrEq
 def testAttrEq():
   with Context():
     a1 = Attribute.parse('"attr1"')
     a2 = Attribute.parse('"attr2"')
     a3 = Attribute.parse('"attr1"')
     # CHECK: a1 == a1: True
     print("a1 == a1:", a1 == a1)
     # CHECK: a1 == a2: False
     print("a1 == a2:", a1 == a2)
     # CHECK: a1 == a3: True
     print("a1 == a3:", a1 == a3)
     # CHECK: a1 == None: False
     print("a1 == None:", a1 == None)
 
 run(testAttrEq)
 
 
 # CHECK-LABEL: TEST: testAttrEqDoesNotRaise
 def testAttrEqDoesNotRaise():
   with Context():
     a1 = Attribute.parse('"attr1"')
     not_an_attr = "foo"
     # CHECK: False
     print(a1 == not_an_attr)
     # CHECK: False
     print(a1 == None)
     # CHECK: True
     print(a1 != None)
 
 run(testAttrEqDoesNotRaise)
 
 
 # CHECK-LABEL: TEST: testAttrCapsule
 def testAttrCapsule():
   with Context() as ctx:
     a1 = Attribute.parse('"attr1"')
   # CHECK: mlir.ir.Attribute._CAPIPtr
   attr_capsule = a1._CAPIPtr
   print(attr_capsule)
   a2 = Attribute._CAPICreate(attr_capsule)
   assert a2 == a1
   assert a2.context is ctx
 
 run(testAttrCapsule)
 
 
 # CHECK-LABEL: TEST: testStandardAttrCasts
 def testStandardAttrCasts():
   with Context():
     a1 = Attribute.parse('"attr1"')
     astr = StringAttr(a1)
     aself = StringAttr(astr)
     # CHECK: Attribute("attr1")
     print(repr(astr))
     try:
       tillegal = StringAttr(Attribute.parse("1.0"))
     except ValueError as e:
       # CHECK: ValueError: Cannot cast attribute to StringAttr (from Attribute(1.000000e+00 : f64))
       print("ValueError:", e)
     else:
       print("Exception not produced")
 
 run(testStandardAttrCasts)
 
 
 # CHECK-LABEL: TEST: testAffineMapAttr
 def testAffineMapAttr():
   with Context() as ctx:
     d0 = AffineDimExpr.get(0)
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     map0 = AffineMap.get(2, 3, [])
 
     # CHECK: affine_map<(d0, d1)[s0, s1, s2] -> ()>
     attr_built = AffineMapAttr.get(map0)
     print(str(attr_built))
 
     attr_parsed = Attribute.parse(str(attr_built))
     assert attr_built == attr_parsed
 
 run(testAffineMapAttr)
 
 
 # CHECK-LABEL: TEST: testFloatAttr
 def testFloatAttr():
   with Context(), Location.unknown():
     fattr = FloatAttr(Attribute.parse("42.0 : f32"))
     # CHECK: fattr value: 42.0
     print("fattr value:", fattr.value)
 
     # Test factory methods.
     # CHECK: default_get: 4.200000e+01 : f32
     print("default_get:", FloatAttr.get(
         F32Type.get(), 42.0))
     # CHECK: f32_get: 4.200000e+01 : f32
     print("f32_get:", FloatAttr.get_f32(42.0))
     # CHECK: f64_get: 4.200000e+01 : f64
     print("f64_get:", FloatAttr.get_f64(42.0))
     try:
       fattr_invalid = FloatAttr.get(
           IntegerType.get_signless(32), 42)
     except ValueError as e:
       # CHECK: invalid 'Type(i32)' and expected floating point type.
       print(e)
     else:
       print("Exception not produced")
 
 run(testFloatAttr)
 
 
 # CHECK-LABEL: TEST: testIntegerAttr
 def testIntegerAttr():
   with Context() as ctx:
     iattr = IntegerAttr(Attribute.parse("42"))
     # CHECK: iattr value: 42
     print("iattr value:", iattr.value)
     # CHECK: iattr type: i64
     print("iattr type:", iattr.type)
 
     # Test factory methods.
     # CHECK: default_get: 42 : i32
     print("default_get:", IntegerAttr.get(
         IntegerType.get_signless(32), 42))
 
 run(testIntegerAttr)
 
 
 # CHECK-LABEL: TEST: testBoolAttr
 def testBoolAttr():
   with Context() as ctx:
     battr = BoolAttr(Attribute.parse("true"))
     # CHECK: iattr value: True
     print("iattr value:", battr.value)
 
     # Test factory methods.
     # CHECK: default_get: true
     print("default_get:", BoolAttr.get(True))
 
 run(testBoolAttr)
 
 
 # CHECK-LABEL: TEST: testFlatSymbolRefAttr
 def testFlatSymbolRefAttr():
   with Context() as ctx:
     sattr = FlatSymbolRefAttr(Attribute.parse('@symbol'))
     # CHECK: symattr value: symbol
     print("symattr value:", sattr.value)
 
     # Test factory methods.
     # CHECK: default_get: @foobar
     print("default_get:", FlatSymbolRefAttr.get("foobar"))
 
 run(testFlatSymbolRefAttr)
 
 
 # CHECK-LABEL: TEST: testStringAttr
 def testStringAttr():
   with Context() as ctx:
     sattr = StringAttr(Attribute.parse('"stringattr"'))
     # CHECK: sattr value: stringattr
     print("sattr value:", sattr.value)
 
     # Test factory methods.
     # CHECK: default_get: "foobar"
     print("default_get:", StringAttr.get("foobar"))
     # CHECK: typed_get: "12345" : i32
     print("typed_get:", StringAttr.get_typed(
         IntegerType.get_signless(32), "12345"))
 
 run(testStringAttr)
 
 
 # CHECK-LABEL: TEST: testNamedAttr
 def testNamedAttr():
   with Context():
     a = Attribute.parse('"stringattr"')
     named = a.get_named("foobar")  # Note: under the small object threshold
     # CHECK: attr: "stringattr"
     print("attr:", named.attr)
     # CHECK: name: foobar
     print("name:", named.name)
     # CHECK: named: NamedAttribute(foobar="stringattr")
     print("named:", named)
 
 run(testNamedAttr)
 
 
 # CHECK-LABEL: TEST: testDenseIntAttr
 def testDenseIntAttr():
   with Context():
     raw = Attribute.parse("dense<[[0,1,2],[3,4,5]]> : vector<2x3xi32>")
     # CHECK: attr: dense<[{{\[}}0, 1, 2], [3, 4, 5]]>
     print("attr:", raw)
 
     a = DenseIntElementsAttr(raw)
     assert len(a) == 6
 
     # CHECK: 0 1 2 3 4 5
     for value in a:
       print(value, end=" ")
     print()
 
     # CHECK: i32
     print(ShapedType(a.type).element_type)
 
     raw = Attribute.parse("dense<[true,false,true,false]> : vector<4xi1>")
     # CHECK: attr: dense<[true, false, true, false]>
     print("attr:", raw)
 
     a = DenseIntElementsAttr(raw)
     assert len(a) == 4
 
     # CHECK: 1 0 1 0
     for value in a:
       print(value, end=" ")
     print()
 
     # CHECK: i1
     print(ShapedType(a.type).element_type)
 
 
 run(testDenseIntAttr)
 
 
 # CHECK-LABEL: TEST: testDenseFPAttr
 def testDenseFPAttr():
   with Context():
     raw = Attribute.parse("dense<[0.0, 1.0, 2.0, 3.0]> : vector<4xf32>")
     # CHECK: attr: dense<[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00]>
 
     print("attr:", raw)
 
     a = DenseFPElementsAttr(raw)
     assert len(a) == 4
 
     # CHECK: 0.0 1.0 2.0 3.0
     for value in a:
       print(value, end=" ")
     print()
 
     # CHECK: f32
     print(ShapedType(a.type).element_type)
 
 
 run(testDenseFPAttr)
 
 
 # CHECK-LABEL: TEST: testDictAttr
 def testDictAttr():
   with Context():
     dict_attr = {
       'stringattr':  StringAttr.get('string'),
       'integerattr' : IntegerAttr.get(
         IntegerType.get_signless(32), 42)
     }
 
     a = DictAttr.get(dict_attr)
 
     # CHECK attr: {integerattr = 42 : i32, stringattr = "string"}
     print("attr:", a)
 
     assert len(a) == 2
 
     # CHECK: 42 : i32
     print(a['integerattr'])
 
     # CHECK: "string"
     print(a['stringattr'])
 
     # Check that exceptions are raised as expected.
     try:
       _ = a['does_not_exist']
     except KeyError:
       pass
     else:
       assert False, "Exception not produced"
 
     try:
       _ = a[42]
     except IndexError:
       pass
     else:
       assert False, "expected IndexError on accessing an out-of-bounds attribute"
 
 
 
 run(testDictAttr)
 
 # CHECK-LABEL: TEST: testTypeAttr
 def testTypeAttr():
   with Context():
     raw = Attribute.parse("vector<4xf32>")
     # CHECK: attr: vector<4xf32>
     print("attr:", raw)
     type_attr = TypeAttr(raw)
     # CHECK: f32
     print(ShapedType(type_attr.value).element_type)
 
 
 run(testTypeAttr)
 
 
 # CHECK-LABEL: TEST: testArrayAttr
 def testArrayAttr():
   with Context():
     raw = Attribute.parse("[42, true, vector<4xf32>]")
   # CHECK: attr: [42, true, vector<4xf32>]
   print("raw attr:", raw)
   # CHECK: - 42
   # CHECK: - true
   # CHECK: - vector<4xf32>
   for attr in ArrayAttr(raw):
     print("- ", attr)
 
   with Context():
     intAttr = Attribute.parse("42")
     vecAttr = Attribute.parse("vector<4xf32>")
     boolAttr = BoolAttr.get(True)
     raw = ArrayAttr.get([vecAttr, boolAttr, intAttr])
   # CHECK: attr: [vector<4xf32>, true, 42]
   print("raw attr:", raw)
   # CHECK: - vector<4xf32>
   # CHECK: - true
   # CHECK: - 42
   arr = ArrayAttr(raw)
   for attr in arr:
     print("- ", attr)
   # CHECK: attr[0]: vector<4xf32>
   print("attr[0]:", arr[0])
   # CHECK: attr[1]: true
   print("attr[1]:", arr[1])
   # CHECK: attr[2]: 42
   print("attr[2]:", arr[2])
   try:
     print("attr[3]:", arr[3])
   except IndexError as e:
     # CHECK: Error: ArrayAttribute index out of range
     print("Error: ", e)
   with Context():
     try:
       ArrayAttr.get([None])
     except RuntimeError as e:
       # CHECK: Error: Invalid attribute (None?) when attempting to create an ArrayAttribute
       print("Error: ", e)
     try:
       ArrayAttr.get([42])
     except RuntimeError as e:
       # CHECK: Error: Invalid attribute when attempting to create an ArrayAttribute
       print("Error: ", e)
 run(testArrayAttr)
 

mlir/test/python/ir/context_lifecycle.py

 # RUN: %PYTHON %s
 # Standalone sanity check of context life-cycle.
 import gc
 import mlir.ir
 
 assert mlir.ir.Context._get_live_count() == 0
 
 # Create first context.
 print("CREATE C1")
 c1 = mlir.ir.Context()
 assert mlir.ir.Context._get_live_count() == 1
 c1_repr = repr(c1)
 print("C1 = ", c1_repr)
 
 print("GETTING AGAIN...")
 c2 = c1._get_context_again()
 c2_repr = repr(c2)
 assert mlir.ir.Context._get_live_count() == 1
 assert c1_repr == c2_repr
 
 print("C2 =", c2)
 
 # Make sure new contexts on constructor.
 print("CREATE C3")
 c3 = mlir.ir.Context()
 assert mlir.ir.Context._get_live_count() == 2
 c3_repr = repr(c3)
 print("C3 =", c3)
 assert c3_repr != c1_repr
 print("FREE C3")
 c3 = None
 gc.collect()
 assert mlir.ir.Context._get_live_count() == 1
 
 print("Free C1")
 c1 = None
 gc.collect()
 assert mlir.ir.Context._get_live_count() == 1
 print("Free C2")
 c2 = None
 gc.collect()
 assert mlir.ir.Context._get_live_count() == 0
 
 # Create a context, get its capsule and create from capsule.
 c4 = mlir.ir.Context()
 c4_capsule = c4._CAPIPtr
 assert '"mlir.ir.Context._CAPIPtr"' in repr(c4_capsule)
 c5 = mlir.ir.Context._CAPICreate(c4_capsule)
 assert c4 is c5

mlir/test/python/ir/context_managers.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
 
 
 # CHECK-LABEL: TEST: testContextEnterExit
 def testContextEnterExit():
   with Context() as ctx:
     assert Context.current is ctx
   try:
     _ = Context.current
   except ValueError as e:
     # CHECK: No current Context
     print(e)
   else: assert False, "Expected exception"
 
 run(testContextEnterExit)
 
 
 # CHECK-LABEL: TEST: testLocationEnterExit
 def testLocationEnterExit():
   ctx1 = Context()
   with Location.unknown(ctx1) as loc1:
     assert Context.current is ctx1
     assert Location.current is loc1
 
     # Re-asserting the same context should not change the location.
     with ctx1:
       assert Context.current is ctx1
       assert Location.current is loc1
       # Asserting a different context should clear it.
       with Context() as ctx2:
         assert Context.current is ctx2
         try:
           _ = Location.current
         except ValueError: pass
         else: assert False, "Expected exception"
 
       # And should restore.
       assert Context.current is ctx1
       assert Location.current is loc1
 
   # All should clear.
   try:
     _ = Location.current
   except ValueError as e:
     # CHECK: No current Location
     print(e)
   else: assert False, "Expected exception"
 
 run(testLocationEnterExit)
 
 
 # CHECK-LABEL: TEST: testInsertionPointEnterExit
 def testInsertionPointEnterExit():
   ctx1 = Context()
   m = Module.create(Location.unknown(ctx1))
   ip = InsertionPoint(m.body)
 
   with ip:
     assert InsertionPoint.current is ip
     # Asserting a location from the same context should preserve.
     with Location.unknown(ctx1) as loc1:
       assert InsertionPoint.current is ip
       assert Location.current is loc1
     # Location should clear.
     try:
       _ = Location.current
     except ValueError: pass
     else: assert False, "Expected exception"
 
     # Asserting the same Context should preserve.
     with ctx1:
       assert InsertionPoint.current is ip
 
     # Asserting a different context should clear it.
     with Context() as ctx2:
       assert Context.current is ctx2
       try:
         _ = InsertionPoint.current
       except ValueError: pass
       else: assert False, "Expected exception"
 
   # All should clear.
   try:
     _ = InsertionPoint.current
   except ValueError as e:
     # CHECK: No current InsertionPoint
     print(e)
   else: assert False, "Expected exception"
 
 run(testInsertionPointEnterExit)

mlir/test/python/ir/debug.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
 
 
 # CHECK-LABEL: TEST: testNameIsPrivate
 def testNameIsPrivate():
   # `import *` ignores private names starting with an understore, so the debug
   # flag shouldn't be visible unless explicitly imported.
   try:
     _GlobalDebug.flag = True
   except NameError:
     pass
   else:
     assert False, "_GlobalDebug must not be available by default"
 
 run(testNameIsPrivate)
 
 
 # CHECK-LABEL: TEST: testDebugDlag
 def testDebugDlag():
   # Private names must be imported expilcitly.
   from mlir.ir import _GlobalDebug
 
   # CHECK: False
   print(_GlobalDebug.flag)
   _GlobalDebug.flag = True
   # CHECK: True
   print(_GlobalDebug.flag)
   _GlobalDebug.flag = False
   # CHECK: False
   print(_GlobalDebug.flag)
 
 run(testDebugDlag)
 

mlir/test/python/ir/dialects.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
   return f
 
 
 # CHECK-LABEL: TEST: testDialectDescriptor
 @run
 def testDialectDescriptor():
   ctx = Context()
   d = ctx.get_dialect_descriptor("std")
   # CHECK: <DialectDescriptor std>
   print(d)
   # CHECK: std
   print(d.namespace)
   try:
     _ = ctx.get_dialect_descriptor("not_existing")
   except ValueError:
     pass
   else:
     assert False, "Expected exception"
 
 
 # CHECK-LABEL: TEST: testUserDialectClass
 @run
 def testUserDialectClass():
   ctx = Context()
   # Access using attribute.
   d = ctx.dialects.std
   # Note that the standard dialect namespace prints as ''. Others will print
   # as "<Dialect %namespace (..."
   # CHECK: <Dialect (class mlir.dialects._std_ops_gen._Dialect)>
   print(d)
   try:
     _ = ctx.dialects.not_existing
   except AttributeError:
     pass
   else:
     assert False, "Expected exception"
 
   # Access using index.
   d = ctx.dialects["std"]
   # CHECK: <Dialect (class mlir.dialects._std_ops_gen._Dialect)>
   print(d)
   try:
     _ = ctx.dialects["not_existing"]
   except IndexError:
     pass
   else:
     assert False, "Expected exception"
 
   # Using the 'd' alias.
   d = ctx.d["std"]
   # CHECK: <Dialect (class mlir.dialects._std_ops_gen._Dialect)>
   print(d)
 
 
 # CHECK-LABEL: TEST: testCustomOpView
 # This test uses the standard dialect AddFOp as an example of a user op.
 # TODO: Op creation and access is still quite verbose: simplify this test as
 # additional capabilities come online.
 @run
 def testCustomOpView():
 
   def createInput():
     op = Operation.create("pytest_dummy.intinput", results=[f32])
     # TODO: Auto result cast from operation
     return op.results[0]
 
   with Context() as ctx, Location.unknown():
     ctx.allow_unregistered_dialects = True
     m = Module.create()
 
     with InsertionPoint(m.body):
       f32 = F32Type.get()
       # Create via dialects context collection.
       input1 = createInput()
       input2 = createInput()
       op1 = ctx.dialects.std.AddFOp(input1.type, input1, input2)
 
       # Create via an import
       from mlir.dialects.std import AddFOp
       AddFOp(input1.type, input1, op1.result)
 
   # CHECK: %[[INPUT0:.*]] = "pytest_dummy.intinput"
   # CHECK: %[[INPUT1:.*]] = "pytest_dummy.intinput"
   # CHECK: %[[R0:.*]] = addf %[[INPUT0]], %[[INPUT1]] : f32
   # CHECK: %[[R1:.*]] = addf %[[INPUT0]], %[[R0]] : f32
   m.operation.print()
 
 
 # CHECK-LABEL: TEST: testIsRegisteredOperation
 @run
 def testIsRegisteredOperation():
   ctx = Context()
 
   # CHECK: std.cond_br: True
   print(f"std.cond_br: {ctx.is_registered_operation('std.cond_br')}")
   # CHECK: std.not_existing: False
   print(f"std.not_existing: {ctx.is_registered_operation('std.not_existing')}")

mlir/test/python/ir/insertion_point.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 import io
 import itertools
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
 
 
 # CHECK-LABEL: TEST: test_insert_at_block_end
 def test_insert_at_block_end():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     module = Module.parse(r"""
       func @foo() -> () {
         "custom.op1"() : () -> ()
       }
     """)
     entry_block = module.body.operations[0].regions[0].blocks[0]
     ip = InsertionPoint(entry_block)
     ip.insert(Operation.create("custom.op2"))
     # CHECK: "custom.op1"
     # CHECK: "custom.op2"
     module.operation.print()
 
 run(test_insert_at_block_end)
 
 
 # CHECK-LABEL: TEST: test_insert_before_operation
 def test_insert_before_operation():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     module = Module.parse(r"""
       func @foo() -> () {
         "custom.op1"() : () -> ()
         "custom.op2"() : () -> ()
       }
     """)
     entry_block = module.body.operations[0].regions[0].blocks[0]
     ip = InsertionPoint(entry_block.operations[1])
     ip.insert(Operation.create("custom.op3"))
     # CHECK: "custom.op1"
     # CHECK: "custom.op3"
     # CHECK: "custom.op2"
     module.operation.print()
 
 run(test_insert_before_operation)
 
 
 # CHECK-LABEL: TEST: test_insert_at_block_begin
 def test_insert_at_block_begin():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     module = Module.parse(r"""
       func @foo() -> () {
         "custom.op2"() : () -> ()
       }
     """)
     entry_block = module.body.operations[0].regions[0].blocks[0]
     ip = InsertionPoint.at_block_begin(entry_block)
     ip.insert(Operation.create("custom.op1"))
     # CHECK: "custom.op1"
     # CHECK: "custom.op2"
     module.operation.print()
 
 run(test_insert_at_block_begin)
 
 
 # CHECK-LABEL: TEST: test_insert_at_block_begin_empty
 def test_insert_at_block_begin_empty():
   # TODO: Write this test case when we can create such a situation.
   pass
 
 run(test_insert_at_block_begin_empty)
 
 
 # CHECK-LABEL: TEST: test_insert_at_terminator
 def test_insert_at_terminator():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     module = Module.parse(r"""
       func @foo() -> () {
         "custom.op1"() : () -> ()
         return
       }
     """)
     entry_block = module.body.operations[0].regions[0].blocks[0]
     ip = InsertionPoint.at_block_terminator(entry_block)
     ip.insert(Operation.create("custom.op2"))
     # CHECK: "custom.op1"
     # CHECK: "custom.op2"
     module.operation.print()
 
 run(test_insert_at_terminator)
 
 
 # CHECK-LABEL: TEST: test_insert_at_block_terminator_missing
 def test_insert_at_block_terminator_missing():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with ctx:
     module = Module.parse(r"""
       func @foo() -> () {
         "custom.op1"() : () -> ()
       }
     """)
     entry_block = module.body.operations[0].regions[0].blocks[0]
     try:
       ip = InsertionPoint.at_block_terminator(entry_block)
     except ValueError as e:
       # CHECK: Block has no terminator
       print(e)
     else:
       assert False, "Expected exception"
 
 run(test_insert_at_block_terminator_missing)
 
 
 # CHECK-LABEL: TEST: test_insert_at_end_with_terminator_errors
 def test_insert_at_end_with_terminator_errors():
   with Context() as ctx, Location.unknown():
     ctx.allow_unregistered_dialects = True
     module = Module.parse(r"""
       func @foo() -> () {
         return
       }
     """)
     entry_block = module.body.operations[0].regions[0].blocks[0]
     with InsertionPoint(entry_block):
       try:
         Operation.create("custom.op1", results=[], operands=[])
       except IndexError as e:
         # CHECK: ERROR: Cannot insert operation at the end of a block that already has a terminator.
         print(f"ERROR: {e}")
 
 run(test_insert_at_end_with_terminator_errors)
 
 
 # CHECK-LABEL: TEST: test_insertion_point_context
 def test_insertion_point_context():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     module = Module.parse(r"""
       func @foo() -> () {
         "custom.op1"() : () -> ()
       }
     """)
     entry_block = module.body.operations[0].regions[0].blocks[0]
     with InsertionPoint(entry_block):
       Operation.create("custom.op2")
       with InsertionPoint.at_block_begin(entry_block):
         Operation.create("custom.opa")
         Operation.create("custom.opb")
       Operation.create("custom.op3")
     # CHECK: "custom.opa"
     # CHECK: "custom.opb"
     # CHECK: "custom.op1"
     # CHECK: "custom.op2"
     # CHECK: "custom.op3"
     module.operation.print()
 
 run(test_insertion_point_context)

mlir/test/python/ir/integer_set.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
 
 
 # CHECK-LABEL: TEST: testIntegerSetCapsule
 def testIntegerSetCapsule():
   with Context() as ctx:
     is1 = IntegerSet.get_empty(1, 1, ctx)
   capsule = is1._CAPIPtr
   # CHECK: mlir.ir.IntegerSet._CAPIPtr
   print(capsule)
   is2 = IntegerSet._CAPICreate(capsule)
   assert is1 == is2
   assert is2.context is ctx
 
 run(testIntegerSetCapsule)
 
 
 # CHECK-LABEL: TEST: testIntegerSetGet
 def testIntegerSetGet():
   with Context():
     d0 = AffineDimExpr.get(0)
     d1 = AffineDimExpr.get(1)
     s0 = AffineSymbolExpr.get(0)
     c42 = AffineConstantExpr.get(42)
 
     # CHECK: (d0, d1)[s0] : (d0 - d1 == 0, s0 - 42 >= 0)
     set0 = IntegerSet.get(2, 1, [d0 - d1, s0 - c42], [True, False])
     print(set0)
 
     # CHECK: (d0)[s0] : (1 == 0)
     set1 = IntegerSet.get_empty(1, 1)
     print(set1)
 
     # CHECK: (d0)[s0, s1] : (d0 - s1 == 0, s0 - 42 >= 0)
     set2 = set0.get_replaced([d0, AffineSymbolExpr.get(1)], [s0], 1, 2)
     print(set2)
 
     try:
       IntegerSet.get(2, 1, [], [])
     except ValueError as e:
       # CHECK: Expected non-empty list of constraints
       print(e)
 
     try:
       IntegerSet.get(2, 1, [d0 - d1], [True, False])
     except ValueError as e:
       # CHECK: Expected the number of constraints to match that of equality flags
       print(e)
 
     try:
       IntegerSet.get(2, 1, [0], [True])
     except RuntimeError as e:
       # CHECK: Invalid expression when attempting to create an IntegerSet
       print(e)
 
     try:
       IntegerSet.get(2, 1, [None], [True])
     except RuntimeError as e:
       # CHECK: Invalid expression (None?) when attempting to create an IntegerSet
       print(e)
 
     try:
       set0.get_replaced([d0], [s0], 1, 1)
     except ValueError as e:
       # CHECK: Expected the number of dimension replacement expressions to match that of dimensions
       print(e)
 
     try:
       set0.get_replaced([d0, d1], [s0, s0], 1, 1)
     except ValueError as e:
       # CHECK: Expected the number of symbol replacement expressions to match that of symbols
       print(e)
 
     try:
       set0.get_replaced([d0, 1], [s0], 1, 1)
     except RuntimeError as e:
       # CHECK: Invalid expression when attempting to create an IntegerSet by replacing dimensions
       print(e)
 
     try:
       set0.get_replaced([d0, d1], [None], 1, 1)
     except RuntimeError as e:
       # CHECK: Invalid expression (None?) when attempting to create an IntegerSet by replacing symbols
       print(e)
 
 run(testIntegerSetGet)
 
 
 # CHECK-LABEL: TEST: testIntegerSetProperties
 def testIntegerSetProperties():
   with Context():
     d0 = AffineDimExpr.get(0)
     d1 = AffineDimExpr.get(1)
     s0 = AffineSymbolExpr.get(0)
     c42 = AffineConstantExpr.get(42)
 
     set0 = IntegerSet.get(2, 1, [d0 - d1, s0 - c42, s0 - d0], [True, False, False])
     # CHECK: 2
     print(set0.n_dims)
     # CHECK: 1
     print(set0.n_symbols)
     # CHECK: 3
     print(set0.n_inputs)
     # CHECK: 1
     print(set0.n_equalities)
     # CHECK: 2
     print(set0.n_inequalities)
 
     # CHECK: 3
     print(len(set0.constraints))
 
     # CHECK-DAG: d0 - d1 == 0
     # CHECK-DAG: s0 - 42 >= 0
     # CHECK-DAG: -d0 + s0 >= 0
     for cstr in set0.constraints:
       print(cstr.expr, end='')
       print(" == 0" if cstr.is_eq else " >= 0")
 
 run(testIntegerSetProperties)

mlir/test/python/ir/location.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
 
 
 # CHECK-LABEL: TEST: testUnknown
 def testUnknown():
   with Context() as ctx:
     loc = Location.unknown()
   assert loc.context is ctx
   ctx = None
   gc.collect()
   # CHECK: unknown str: loc(unknown)
   print("unknown str:", str(loc))
   # CHECK: unknown repr: loc(unknown)
   print("unknown repr:", repr(loc))
 
 run(testUnknown)
 
 
 # CHECK-LABEL: TEST: testFileLineCol
 def testFileLineCol():
   with Context() as ctx:
     loc = Location.file("foo.txt", 123, 56)
   ctx = None
   gc.collect()
   # CHECK: file str: loc("foo.txt":123:56)
   print("file str:", str(loc))
   # CHECK: file repr: loc("foo.txt":123:56)
   print("file repr:", repr(loc))
 
 run(testFileLineCol)
 
 
 # CHECK-LABEL: TEST: testLocationCapsule
 def testLocationCapsule():
   with Context() as ctx:
     loc1 = Location.file("foo.txt", 123, 56)
   # CHECK: mlir.ir.Location._CAPIPtr
   loc_capsule = loc1._CAPIPtr
   print(loc_capsule)
   loc2 = Location._CAPICreate(loc_capsule)
   assert loc2 == loc1
   assert loc2.context is ctx
 
 run(testLocationCapsule)

mlir/test/python/ir/module.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
 
 
 # Verify successful parse.
 # CHECK-LABEL: TEST: testParseSuccess
 # CHECK: module @successfulParse
 def testParseSuccess():
   ctx = Context()
   module = Module.parse(r"""module @successfulParse {}""", ctx)
   assert module.context is ctx
   print("CLEAR CONTEXT")
   ctx = None  # Ensure that module captures the context.
   gc.collect()
   module.dump()  # Just outputs to stderr. Verifies that it functions.
   print(str(module))
 
 run(testParseSuccess)
 
 
 # Verify parse error.
 # CHECK-LABEL: TEST: testParseError
 # CHECK: testParseError: Unable to parse module assembly (see diagnostics)
 def testParseError():
   ctx = Context()
   try:
     module = Module.parse(r"""}SYNTAX ERROR{""", ctx)
   except ValueError as e:
     print("testParseError:", e)
   else:
     print("Exception not produced")
 
 run(testParseError)
 
 
 # Verify successful parse.
 # CHECK-LABEL: TEST: testCreateEmpty
 # CHECK: module {
 def testCreateEmpty():
   ctx = Context()
   loc = Location.unknown(ctx)
   module = Module.create(loc)
   print("CLEAR CONTEXT")
   ctx = None  # Ensure that module captures the context.
   gc.collect()
   print(str(module))
 
 run(testCreateEmpty)
 
 
 # Verify round-trip of ASM that contains unicode.
 # Note that this does not test that the print path converts unicode properly
 # because MLIR asm always normalizes it to the hex encoding.
 # CHECK-LABEL: TEST: testRoundtripUnicode
 # CHECK: func private @roundtripUnicode()
 # CHECK: foo = "\F0\9F\98\8A"
 def testRoundtripUnicode():
   ctx = Context()
   module = Module.parse(r"""
     func private @roundtripUnicode() attributes { foo = "😊" }
   """, ctx)
   print(str(module))
 
 run(testRoundtripUnicode)
 
 
 # Tests that module.operation works and correctly interns instances.
 # CHECK-LABEL: TEST: testModuleOperation
 def testModuleOperation():
   ctx = Context()
   module = Module.parse(r"""module @successfulParse {}""", ctx)
   assert ctx._get_live_module_count() == 1
   op1 = module.operation
   assert ctx._get_live_operation_count() == 1
   # CHECK: module @successfulParse
   print(op1)
 
   # Ensure that operations are the same on multiple calls.
   op2 = module.operation
   assert ctx._get_live_operation_count() == 1
   assert op1 is op2
 
   # Ensure that if module is de-referenced, the operations are still valid.
   module = None
   gc.collect()
   print(op1)
 
   # Collect and verify lifetime.
   op1 = None
   op2 = None
   gc.collect()
   print("LIVE OPERATIONS:", ctx._get_live_operation_count())
   assert ctx._get_live_operation_count() == 0
   assert ctx._get_live_module_count() == 0
 
 run(testModuleOperation)
 
 
 # CHECK-LABEL: TEST: testModuleCapsule
 def testModuleCapsule():
   ctx = Context()
   module = Module.parse(r"""module @successfulParse {}""", ctx)
   assert ctx._get_live_module_count() == 1
   # CHECK: "mlir.ir.Module._CAPIPtr"
   module_capsule = module._CAPIPtr
   print(module_capsule)
   module_dup = Module._CAPICreate(module_capsule)
   assert module is module_dup
   assert module_dup.context is ctx
   # Gc and verify destructed.
   module = None
   module_capsule = None
   module_dup = None
   gc.collect()
   assert ctx._get_live_module_count() == 0
 
 
 run(testModuleCapsule)

mlir/test/python/ir/operation.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 import io
 import itertools
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
 
 
 # Verify iterator based traversal of the op/region/block hierarchy.
 # CHECK-LABEL: TEST: testTraverseOpRegionBlockIterators
 def testTraverseOpRegionBlockIterators():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   module = Module.parse(r"""
     func @f1(%arg0: i32) -> i32 {
       %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
       return %1 : i32
     }
   """, ctx)
   op = module.operation
   assert op.context is ctx
   # Get the block using iterators off of the named collections.
   regions = list(op.regions)
   blocks = list(regions[0].blocks)
   # CHECK: MODULE REGIONS=1 BLOCKS=1
   print(f"MODULE REGIONS={len(regions)} BLOCKS={len(blocks)}")
 
   # Should verify.
   # CHECK: .verify = True
   print(f".verify = {module.operation.verify()}")
 
   # Get the regions and blocks from the default collections.
   default_regions = list(op)
   default_blocks = list(default_regions[0])
   # They should compare equal regardless of how obtained.
   assert default_regions == regions
   assert default_blocks == blocks
 
   # Should be able to get the operations from either the named collection
   # or the block.
   operations = list(blocks[0].operations)
   default_operations = list(blocks[0])
   assert default_operations == operations
 
   def walk_operations(indent, op):
     for i, region in enumerate(op):
       print(f"{indent}REGION {i}:")
       for j, block in enumerate(region):
         print(f"{indent}  BLOCK {j}:")
         for k, child_op in enumerate(block):
           print(f"{indent}    OP {k}: {child_op}")
           walk_operations(indent + "      ", child_op)
 
   # CHECK: REGION 0:
   # CHECK:   BLOCK 0:
   # CHECK:     OP 0: func
   # CHECK:       REGION 0:
   # CHECK:         BLOCK 0:
   # CHECK:           OP 0: %0 = "custom.addi"
   # CHECK:           OP 1: return
   walk_operations("", op)
 
 run(testTraverseOpRegionBlockIterators)
 
 
 # Verify index based traversal of the op/region/block hierarchy.
 # CHECK-LABEL: TEST: testTraverseOpRegionBlockIndices
 def testTraverseOpRegionBlockIndices():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   module = Module.parse(r"""
     func @f1(%arg0: i32) -> i32 {
       %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
       return %1 : i32
     }
   """, ctx)
 
   def walk_operations(indent, op):
     for i in range(len(op.regions)):
       region = op.regions[i]
       print(f"{indent}REGION {i}:")
       for j in range(len(region.blocks)):
         block = region.blocks[j]
         print(f"{indent}  BLOCK {j}:")
         for k in range(len(block.operations)):
           child_op = block.operations[k]
           print(f"{indent}    OP {k}: {child_op}")
           walk_operations(indent + "      ", child_op)
 
   # CHECK: REGION 0:
   # CHECK:   BLOCK 0:
   # CHECK:     OP 0: func
   # CHECK:       REGION 0:
   # CHECK:         BLOCK 0:
   # CHECK:           OP 0: %0 = "custom.addi"
   # CHECK:           OP 1: return
   walk_operations("", module.operation)
 
 run(testTraverseOpRegionBlockIndices)
 
 
 # CHECK-LABEL: TEST: testBlockArgumentList
 def testBlockArgumentList():
   with Context() as ctx:
     module = Module.parse(r"""
       func @f1(%arg0: i32, %arg1: f64, %arg2: index) {
         return
       }
     """, ctx)
     func = module.body.operations[0]
     entry_block = func.regions[0].blocks[0]
     assert len(entry_block.arguments) == 3
     # CHECK: Argument 0, type i32
     # CHECK: Argument 1, type f64
     # CHECK: Argument 2, type index
     for arg in entry_block.arguments:
       print(f"Argument {arg.arg_number}, type {arg.type}")
       new_type = IntegerType.get_signless(8 * (arg.arg_number + 1))
       arg.set_type(new_type)
 
     # CHECK: Argument 0, type i8
     # CHECK: Argument 1, type i16
     # CHECK: Argument 2, type i24
     for arg in entry_block.arguments:
       print(f"Argument {arg.arg_number}, type {arg.type}")
 
 
 run(testBlockArgumentList)
 
 
 # CHECK-LABEL: TEST: testOperationOperands
 def testOperationOperands():
   with Context() as ctx:
     ctx.allow_unregistered_dialects = True
     module = Module.parse(r"""
       func @f1(%arg0: i32) {
         %0 = "test.producer"() : () -> i64
         "test.consumer"(%arg0, %0) : (i32, i64) -> ()
         return
       }""")
     func = module.body.operations[0]
     entry_block = func.regions[0].blocks[0]
     consumer = entry_block.operations[1]
     assert len(consumer.operands) == 2
     # CHECK: Operand 0, type i32
     # CHECK: Operand 1, type i64
     for i, operand in enumerate(consumer.operands):
       print(f"Operand {i}, type {operand.type}")
 
 
 run(testOperationOperands)
 
 
 # CHECK-LABEL: TEST: testOperationOperandsSlice
 def testOperationOperandsSlice():
   with Context() as ctx:
     ctx.allow_unregistered_dialects = True
     module = Module.parse(r"""
       func @f1() {
         %0 = "test.producer0"() : () -> i64
         %1 = "test.producer1"() : () -> i64
         %2 = "test.producer2"() : () -> i64
         %3 = "test.producer3"() : () -> i64
         %4 = "test.producer4"() : () -> i64
         "test.consumer"(%0, %1, %2, %3, %4) : (i64, i64, i64, i64, i64) -> ()
         return
       }""")
     func = module.body.operations[0]
     entry_block = func.regions[0].blocks[0]
     consumer = entry_block.operations[5]
     assert len(consumer.operands) == 5
     for left, right in zip(consumer.operands, consumer.operands[::-1][::-1]):
       assert left == right
 
     # CHECK: test.producer0
     # CHECK: test.producer1
     # CHECK: test.producer2
     # CHECK: test.producer3
     # CHECK: test.producer4
     full_slice = consumer.operands[:]
     for operand in full_slice:
       print(operand)
 
     # CHECK: test.producer0
     # CHECK: test.producer1
     first_two = consumer.operands[0:2]
     for operand in first_two:
       print(operand)
 
     # CHECK: test.producer3
     # CHECK: test.producer4
     last_two = consumer.operands[3:]
     for operand in last_two:
       print(operand)
 
     # CHECK: test.producer0
     # CHECK: test.producer2
     # CHECK: test.producer4
     even = consumer.operands[::2]
     for operand in even:
       print(operand)
 
     # CHECK: test.producer2
     fourth = consumer.operands[::2][1::2]
     for operand in fourth:
       print(operand)
 
 
 run(testOperationOperandsSlice)
 
 
 # CHECK-LABEL: TEST: testOperationOperandsSet
 def testOperationOperandsSet():
   with Context() as ctx, Location.unknown(ctx):
     ctx.allow_unregistered_dialects = True
     module = Module.parse(r"""
       func @f1() {
         %0 = "test.producer0"() : () -> i64
         %1 = "test.producer1"() : () -> i64
         %2 = "test.producer2"() : () -> i64
         "test.consumer"(%0) : (i64) -> ()
         return
       }""")
     func = module.body.operations[0]
     entry_block = func.regions[0].blocks[0]
     producer1 = entry_block.operations[1]
     producer2 = entry_block.operations[2]
     consumer = entry_block.operations[3]
     assert len(consumer.operands) == 1
     type = consumer.operands[0].type
 
     # CHECK: test.producer1
     consumer.operands[0] = producer1.result
     print(consumer.operands[0])
 
     # CHECK: test.producer2
     consumer.operands[-1] = producer2.result
     print(consumer.operands[0])
 
 
 run(testOperationOperandsSet)
 
 
 # CHECK-LABEL: TEST: testDetachedOperation
 def testDetachedOperation():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     i32 = IntegerType.get_signed(32)
     op1 = Operation.create(
         "custom.op1", results=[i32, i32], regions=1, attributes={
             "foo": StringAttr.get("foo_value"),
             "bar": StringAttr.get("bar_value"),
         })
     # CHECK: %0:2 = "custom.op1"() ( {
     # CHECK: }) {bar = "bar_value", foo = "foo_value"} : () -> (si32, si32)
     print(op1)
 
   # TODO: Check successors once enough infra exists to do it properly.
 
 run(testDetachedOperation)
 
 
 # CHECK-LABEL: TEST: testOperationInsertionPoint
 def testOperationInsertionPoint():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   module = Module.parse(r"""
     func @f1(%arg0: i32) -> i32 {
       %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
       return %1 : i32
     }
   """, ctx)
 
   # Create test op.
   with Location.unknown(ctx):
     op1 = Operation.create("custom.op1")
     op2 = Operation.create("custom.op2")
 
     func = module.body.operations[0]
     entry_block = func.regions[0].blocks[0]
     ip = InsertionPoint.at_block_begin(entry_block)
     ip.insert(op1)
     ip.insert(op2)
     # CHECK: func @f1
     # CHECK: "custom.op1"()
     # CHECK: "custom.op2"()
     # CHECK: %0 = "custom.addi"
     print(module)
 
   # Trying to add a previously added op should raise.
   try:
     ip.insert(op1)
   except ValueError:
     pass
   else:
     assert False, "expected insert of attached op to raise"
 
 run(testOperationInsertionPoint)
 
 
 # CHECK-LABEL: TEST: testOperationWithRegion
 def testOperationWithRegion():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     i32 = IntegerType.get_signed(32)
     op1 = Operation.create("custom.op1", regions=1)
     block = op1.regions[0].blocks.append(i32, i32)
     # CHECK: "custom.op1"() ( {
     # CHECK: ^bb0(%arg0: si32, %arg1: si32):  // no predecessors
     # CHECK:   "custom.terminator"() : () -> ()
     # CHECK: }) : () -> ()
     terminator = Operation.create("custom.terminator")
     ip = InsertionPoint(block)
     ip.insert(terminator)
     print(op1)
 
     # Now add the whole operation to another op.
     # TODO: Verify lifetime hazard by nulling out the new owning module and
     # accessing op1.
     # TODO: Also verify accessing the terminator once both parents are nulled
     # out.
     module = Module.parse(r"""
       func @f1(%arg0: i32) -> i32 {
         %1 = "custom.addi"(%arg0, %arg0) : (i32, i32) -> i32
         return %1 : i32
       }
     """)
     func = module.body.operations[0]
     entry_block = func.regions[0].blocks[0]
     ip = InsertionPoint.at_block_begin(entry_block)
     ip.insert(op1)
     # CHECK: func @f1
     # CHECK: "custom.op1"()
     # CHECK:   "custom.terminator"
     # CHECK: %0 = "custom.addi"
     print(module)
 
 run(testOperationWithRegion)
 
 
 # CHECK-LABEL: TEST: testOperationResultList
 def testOperationResultList():
   ctx = Context()
   module = Module.parse(r"""
     func @f1() {
       %0:3 = call @f2() : () -> (i32, f64, index)
       return
     }
     func private @f2() -> (i32, f64, index)
   """, ctx)
   caller = module.body.operations[0]
   call = caller.regions[0].blocks[0].operations[0]
   assert len(call.results) == 3
   # CHECK: Result 0, type i32
   # CHECK: Result 1, type f64
   # CHECK: Result 2, type index
   for res in call.results:
     print(f"Result {res.result_number}, type {res.type}")
 
 
 run(testOperationResultList)
 
 
 # CHECK-LABEL: TEST: testOperationResultListSlice
 def testOperationResultListSlice():
   with Context() as ctx:
     ctx.allow_unregistered_dialects = True
     module = Module.parse(r"""
       func @f1() {
         "some.op"() : () -> (i1, i2, i3, i4, i5)
         return
       }
     """)
     func = module.body.operations[0]
     entry_block = func.regions[0].blocks[0]
     producer = entry_block.operations[0]
 
     assert len(producer.results) == 5
     for left, right in zip(producer.results, producer.results[::-1][::-1]):
       assert left == right
       assert left.result_number == right.result_number
 
     # CHECK: Result 0, type i1
     # CHECK: Result 1, type i2
     # CHECK: Result 2, type i3
     # CHECK: Result 3, type i4
     # CHECK: Result 4, type i5
     full_slice = producer.results[:]
     for res in full_slice:
       print(f"Result {res.result_number}, type {res.type}")
 
     # CHECK: Result 1, type i2
     # CHECK: Result 2, type i3
     # CHECK: Result 3, type i4
     middle = producer.results[1:4]
     for res in middle:
       print(f"Result {res.result_number}, type {res.type}")
 
     # CHECK: Result 1, type i2
     # CHECK: Result 3, type i4
     odd = producer.results[1::2]
     for res in odd:
       print(f"Result {res.result_number}, type {res.type}")
 
     # CHECK: Result 3, type i4
     # CHECK: Result 1, type i2
     inverted_middle = producer.results[-2:0:-2]
     for res in inverted_middle:
       print(f"Result {res.result_number}, type {res.type}")
 
 
 run(testOperationResultListSlice)
 
 
 # CHECK-LABEL: TEST: testOperationAttributes
 def testOperationAttributes():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   module = Module.parse(r"""
     "some.op"() { some.attribute = 1 : i8,
                   other.attribute = 3.0,
                   dependent = "text" } : () -> ()
   """, ctx)
   op = module.body.operations[0]
   assert len(op.attributes) == 3
   iattr = IntegerAttr(op.attributes["some.attribute"])
   fattr = FloatAttr(op.attributes["other.attribute"])
   sattr = StringAttr(op.attributes["dependent"])
   # CHECK: Attribute type i8, value 1
   print(f"Attribute type {iattr.type}, value {iattr.value}")
   # CHECK: Attribute type f64, value 3.0
   print(f"Attribute type {fattr.type}, value {fattr.value}")
   # CHECK: Attribute value text
   print(f"Attribute value {sattr.value}")
 
   # We don't know in which order the attributes are stored.
   # CHECK-DAG: NamedAttribute(dependent="text")
   # CHECK-DAG: NamedAttribute(other.attribute=3.000000e+00 : f64)
   # CHECK-DAG: NamedAttribute(some.attribute=1 : i8)
   for attr in op.attributes:
     print(str(attr))
 
   # Check that exceptions are raised as expected.
   try:
     op.attributes["does_not_exist"]
   except KeyError:
     pass
   else:
     assert False, "expected KeyError on accessing a non-existent attribute"
 
   try:
     op.attributes[42]
   except IndexError:
     pass
   else:
     assert False, "expected IndexError on accessing an out-of-bounds attribute"
 
 
 run(testOperationAttributes)
 
 
 # CHECK-LABEL: TEST: testOperationPrint
 def testOperationPrint():
   ctx = Context()
   module = Module.parse(r"""
     func @f1(%arg0: i32) -> i32 {
       %0 = constant dense<[1, 2, 3, 4]> : tensor<4xi32>
       return %arg0 : i32
     }
   """, ctx)
 
   # Test print to stdout.
   # CHECK: return %arg0 : i32
   module.operation.print()
 
   # Test print to text file.
   f = io.StringIO()
   # CHECK: <class 'str'>
   # CHECK: return %arg0 : i32
   module.operation.print(file=f)
   str_value = f.getvalue()
   print(str_value.__class__)
   print(f.getvalue())
 
   # Test print to binary file.
   f = io.BytesIO()
   # CHECK: <class 'bytes'>
   # CHECK: return %arg0 : i32
   module.operation.print(file=f, binary=True)
   bytes_value = f.getvalue()
   print(bytes_value.__class__)
   print(bytes_value)
 
   # Test get_asm with options.
   # CHECK: value = opaque<"_", "0xDEADBEEF"> : tensor<4xi32>
   # CHECK: "std.return"(%arg0) : (i32) -> () -:4:7
   module.operation.print(large_elements_limit=2, enable_debug_info=True,
       pretty_debug_info=True, print_generic_op_form=True, use_local_scope=True)
 
 run(testOperationPrint)
 
 
 # CHECK-LABEL: TEST: testKnownOpView
 def testKnownOpView():
   with Context(), Location.unknown():
     Context.current.allow_unregistered_dialects = True
     module = Module.parse(r"""
       %1 = "custom.f32"() : () -> f32
       %2 = "custom.f32"() : () -> f32
       %3 = addf %1, %2 : f32
     """)
     print(module)
 
     # 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_ops_gen._AddFOp object
     print(repr(addf))
     # CHECK: "custom.f32"()
     print(addf.lhs)
 
     # One of the custom ops should resolve to the default OpView.
     custom = module.body.operations[0]
     # CHECK: <_mlir.ir.OpView object
     print(repr(custom))
 
     # Check again to make sure negative caching works.
     custom = module.body.operations[0]
     # CHECK: <_mlir.ir.OpView object
     print(repr(custom))
 
 run(testKnownOpView)
 
 
 # CHECK-LABEL: TEST: testSingleResultProperty
 def testSingleResultProperty():
   with Context(), Location.unknown():
     Context.current.allow_unregistered_dialects = True
     module = Module.parse(r"""
       "custom.no_result"() : () -> ()
       %0:2 = "custom.two_result"() : () -> (f32, f32)
       %1 = "custom.one_result"() : () -> f32
     """)
     print(module)
 
   try:
     module.body.operations[0].result
   except ValueError as e:
     # CHECK: Cannot call .result on operation custom.no_result which has 0 results
     print(e)
   else:
     assert False, "Expected exception"
 
   try:
     module.body.operations[1].result
   except ValueError as e:
     # CHECK: Cannot call .result on operation custom.two_result which has 2 results
     print(e)
   else:
     assert False, "Expected exception"
 
   # CHECK: %1 = "custom.one_result"() : () -> f32
   print(module.body.operations[2])
 
 run(testSingleResultProperty)
 
 # CHECK-LABEL: TEST: testPrintInvalidOperation
 def testPrintInvalidOperation():
   ctx = Context()
   with Location.unknown(ctx):
     module = Operation.create("module", regions=2)
     # This module has two region and is invalid verify that we fallback
     # to the generic printer for safety.
     block = module.regions[0].blocks.append()
     # CHECK: // Verification failed, printing generic form
     # CHECK: "module"() ( {
     # CHECK: }) : () -> ()
     print(module)
     # CHECK: .verify = False
     print(f".verify = {module.operation.verify()}")
 run(testPrintInvalidOperation)
 
 
 # CHECK-LABEL: TEST: testCreateWithInvalidAttributes
 def testCreateWithInvalidAttributes():
   ctx = Context()
   with Location.unknown(ctx):
     try:
       Operation.create("module", attributes={None:StringAttr.get("name")})
     except Exception as e:
       # CHECK: Invalid attribute key (not a string) when attempting to create the operation "module"
       print(e)
     try:
       Operation.create("module", attributes={42:StringAttr.get("name")})
     except Exception as e:
       # CHECK: Invalid attribute key (not a string) when attempting to create the operation "module"
       print(e)
     try:
       Operation.create("module", attributes={"some_key":ctx})
     except Exception as e:
       # CHECK: Invalid attribute value for the key "some_key" when attempting to create the operation "module"
       print(e)
     try:
       Operation.create("module", attributes={"some_key":None})
     except Exception as e:
       # CHECK: Found an invalid (`None`?) attribute value for the key "some_key" when attempting to create the operation "module"
       print(e)
 run(testCreateWithInvalidAttributes)
 
 
 # CHECK-LABEL: TEST: testOperationName
 def testOperationName():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   module = Module.parse(r"""
     %0 = "custom.op1"() : () -> f32
     %1 = "custom.op2"() : () -> i32
     %2 = "custom.op1"() : () -> f32
   """, ctx)
 
   # CHECK: custom.op1
   # CHECK: custom.op2
   # CHECK: custom.op1
   for op in module.body.operations:
     print(op.operation.name)
 
 run(testOperationName)
 
 # CHECK-LABEL: TEST: testCapsuleConversions
 def testCapsuleConversions():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     m = Operation.create("custom.op1").operation
     m_capsule = m._CAPIPtr
     assert '"mlir.ir.Operation._CAPIPtr"' in repr(m_capsule)
     m2 = Operation._CAPICreate(m_capsule)
     assert m2 is m
 
 run(testCapsuleConversions)

mlir/test/python/ir/types.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
 
 
 # CHECK-LABEL: TEST: testParsePrint
 def testParsePrint():
   ctx = Context()
   t = Type.parse("i32", ctx)
   assert t.context is ctx
   ctx = None
   gc.collect()
   # CHECK: i32
   print(str(t))
   # CHECK: Type(i32)
   print(repr(t))
 
 run(testParsePrint)
 
 
 # CHECK-LABEL: TEST: testParseError
 # TODO: Hook the diagnostic manager to capture a more meaningful error
 # message.
 def testParseError():
   ctx = Context()
   try:
     t = Type.parse("BAD_TYPE_DOES_NOT_EXIST", ctx)
   except ValueError as e:
     # CHECK: Unable to parse type: 'BAD_TYPE_DOES_NOT_EXIST'
     print("testParseError:", e)
   else:
     print("Exception not produced")
 
 run(testParseError)
 
 
 # CHECK-LABEL: TEST: testTypeEq
 def testTypeEq():
   ctx = Context()
   t1 = Type.parse("i32", ctx)
   t2 = Type.parse("f32", ctx)
   t3 = Type.parse("i32", ctx)
   # CHECK: t1 == t1: True
   print("t1 == t1:", t1 == t1)
   # CHECK: t1 == t2: False
   print("t1 == t2:", t1 == t2)
   # CHECK: t1 == t3: True
   print("t1 == t3:", t1 == t3)
   # CHECK: t1 == None: False
   print("t1 == None:", t1 == None)
 
 run(testTypeEq)
 
 
 # CHECK-LABEL: TEST: testTypeIsInstance
 def testTypeIsInstance():
   ctx = Context()
   t1 = Type.parse("i32", ctx)
   t2 = Type.parse("f32", ctx)
   # CHECK: True
   print(IntegerType.isinstance(t1))
   # CHECK: False
   print(F32Type.isinstance(t1))
   # CHECK: True
   print(F32Type.isinstance(t2))
 
 run(testTypeIsInstance)
 
 
 # CHECK-LABEL: TEST: testTypeEqDoesNotRaise
 def testTypeEqDoesNotRaise():
   ctx = Context()
   t1 = Type.parse("i32", ctx)
   not_a_type = "foo"
   # CHECK: False
   print(t1 == not_a_type)
   # CHECK: False
   print(t1 == None)
   # CHECK: True
   print(t1 != None)
 
 run(testTypeEqDoesNotRaise)
 
 
 # CHECK-LABEL: TEST: testTypeCapsule
 def testTypeCapsule():
   with Context() as ctx:
     t1 = Type.parse("i32", ctx)
   # CHECK: mlir.ir.Type._CAPIPtr
   type_capsule = t1._CAPIPtr
   print(type_capsule)
   t2 = Type._CAPICreate(type_capsule)
   assert t2 == t1
   assert t2.context is ctx
 
 run(testTypeCapsule)
 
 
 # CHECK-LABEL: TEST: testStandardTypeCasts
 def testStandardTypeCasts():
   ctx = Context()
   t1 = Type.parse("i32", ctx)
   tint = IntegerType(t1)
   tself = IntegerType(tint)
   # CHECK: Type(i32)
   print(repr(tint))
   try:
     tillegal = IntegerType(Type.parse("f32", ctx))
   except ValueError as e:
     # CHECK: ValueError: Cannot cast type to IntegerType (from Type(f32))
     print("ValueError:", e)
   else:
     print("Exception not produced")
 
 run(testStandardTypeCasts)
 
 
 # CHECK-LABEL: TEST: testIntegerType
 def testIntegerType():
   with Context() as ctx:
     i32 = IntegerType(Type.parse("i32"))
     # CHECK: i32 width: 32
     print("i32 width:", i32.width)
     # CHECK: i32 signless: True
     print("i32 signless:", i32.is_signless)
     # CHECK: i32 signed: False
     print("i32 signed:", i32.is_signed)
     # CHECK: i32 unsigned: False
     print("i32 unsigned:", i32.is_unsigned)
 
     s32 = IntegerType(Type.parse("si32"))
     # CHECK: s32 signless: False
     print("s32 signless:", s32.is_signless)
     # CHECK: s32 signed: True
     print("s32 signed:", s32.is_signed)
     # CHECK: s32 unsigned: False
     print("s32 unsigned:", s32.is_unsigned)
 
     u32 = IntegerType(Type.parse("ui32"))
     # CHECK: u32 signless: False
     print("u32 signless:", u32.is_signless)
     # CHECK: u32 signed: False
     print("u32 signed:", u32.is_signed)
     # CHECK: u32 unsigned: True
     print("u32 unsigned:", u32.is_unsigned)
 
     # CHECK: signless: i16
     print("signless:", IntegerType.get_signless(16))
     # CHECK: signed: si8
     print("signed:", IntegerType.get_signed(8))
     # CHECK: unsigned: ui64
     print("unsigned:", IntegerType.get_unsigned(64))
 
 run(testIntegerType)
 
 # CHECK-LABEL: TEST: testIndexType
 def testIndexType():
   with Context() as ctx:
     # CHECK: index type: index
     print("index type:", IndexType.get())
 
 run(testIndexType)
 
 # CHECK-LABEL: TEST: testFloatType
 def testFloatType():
   with Context():
     # CHECK: float: bf16
     print("float:", BF16Type.get())
     # CHECK: float: f16
     print("float:", F16Type.get())
     # CHECK: float: f32
     print("float:", F32Type.get())
     # CHECK: float: f64
     print("float:", F64Type.get())
 
 run(testFloatType)
 
 # CHECK-LABEL: TEST: testNoneType
 def testNoneType():
   with Context():
     # CHECK: none type: none
     print("none type:", NoneType.get())
 
 run(testNoneType)
 
 # CHECK-LABEL: TEST: testComplexType
 def testComplexType():
   with Context() as ctx:
     complex_i32 = ComplexType(Type.parse("complex<i32>"))
     # CHECK: complex type element: i32
     print("complex type element:", complex_i32.element_type)
 
     f32 = F32Type.get()
     # CHECK: complex type: complex<f32>
     print("complex type:", ComplexType.get(f32))
 
     index = IndexType.get()
     try:
       complex_invalid = ComplexType.get(index)
     except ValueError as e:
       # CHECK: invalid 'Type(index)' and expected floating point or integer type.
       print(e)
     else:
       print("Exception not produced")
 
 run(testComplexType)
 
 
 # CHECK-LABEL: TEST: testConcreteShapedType
 # Shaped type is not a kind of builtin types, it is the base class for vectors,
 # memrefs and tensors, so this test case uses an instance of vector to test the
 # shaped type. The class hierarchy is preserved on the python side.
 def testConcreteShapedType():
   with Context() as ctx:
     vector = VectorType(Type.parse("vector<2x3xf32>"))
     # CHECK: element type: f32
     print("element type:", vector.element_type)
     # CHECK: whether the given shaped type is ranked: True
     print("whether the given shaped type is ranked:", vector.has_rank)
     # CHECK: rank: 2
     print("rank:", vector.rank)
     # CHECK: whether the shaped type has a static shape: True
     print("whether the shaped type has a static shape:", vector.has_static_shape)
     # CHECK: whether the dim-th dimension is dynamic: False
     print("whether the dim-th dimension is dynamic:", vector.is_dynamic_dim(0))
     # CHECK: dim size: 3
     print("dim size:", vector.get_dim_size(1))
     # CHECK: is_dynamic_size: False
     print("is_dynamic_size:", vector.is_dynamic_size(3))
     # CHECK: is_dynamic_stride_or_offset: False
     print("is_dynamic_stride_or_offset:", vector.is_dynamic_stride_or_offset(1))
     # CHECK: isinstance(ShapedType): True
     print("isinstance(ShapedType):", isinstance(vector, ShapedType))
 
 run(testConcreteShapedType)
 
 # CHECK-LABEL: TEST: testAbstractShapedType
 # Tests that ShapedType operates as an abstract base class of a concrete
 # shaped type (using vector as an example).
 def testAbstractShapedType():
   ctx = Context()
   vector = ShapedType(Type.parse("vector<2x3xf32>", ctx))
   # CHECK: element type: f32
   print("element type:", vector.element_type)
 
 run(testAbstractShapedType)
 
 # CHECK-LABEL: TEST: testVectorType
 def testVectorType():
   with Context(), Location.unknown():
     f32 = F32Type.get()
     shape = [2, 3]
     # CHECK: vector type: vector<2x3xf32>
     print("vector type:", VectorType.get(shape, f32))
 
     none = NoneType.get()
     try:
       vector_invalid = VectorType.get(shape, none)
     except ValueError as e:
       # CHECK: invalid 'Type(none)' and expected floating point or integer type.
       print(e)
     else:
       print("Exception not produced")
 
 run(testVectorType)
 
 # CHECK-LABEL: TEST: testRankedTensorType
 def testRankedTensorType():
   with Context(), Location.unknown():
     f32 = F32Type.get()
     shape = [2, 3]
     loc = Location.unknown()
     # CHECK: ranked tensor type: tensor<2x3xf32>
     print("ranked tensor type:",
           RankedTensorType.get(shape, f32))
 
     none = NoneType.get()
     try:
       tensor_invalid = RankedTensorType.get(shape, none)
     except ValueError as e:
       # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
       # CHECK: or complex type.
       print(e)
     else:
       print("Exception not produced")
 
 run(testRankedTensorType)
 
 # CHECK-LABEL: TEST: testUnrankedTensorType
 def testUnrankedTensorType():
   with Context(), Location.unknown():
     f32 = F32Type.get()
     loc = Location.unknown()
     unranked_tensor = UnrankedTensorType.get(f32)
     # CHECK: unranked tensor type: tensor<*xf32>
     print("unranked tensor type:", unranked_tensor)
     try:
       invalid_rank = unranked_tensor.rank
     except ValueError as e:
       # CHECK: calling this method requires that the type has a rank.
       print(e)
     else:
       print("Exception not produced")
     try:
       invalid_is_dynamic_dim = unranked_tensor.is_dynamic_dim(0)
     except ValueError as e:
       # CHECK: calling this method requires that the type has a rank.
       print(e)
     else:
       print("Exception not produced")
     try:
       invalid_get_dim_size = unranked_tensor.get_dim_size(1)
     except ValueError as e:
       # CHECK: calling this method requires that the type has a rank.
       print(e)
     else:
       print("Exception not produced")
 
     none = NoneType.get()
     try:
       tensor_invalid = UnrankedTensorType.get(none)
     except ValueError as e:
       # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
       # CHECK: or complex type.
       print(e)
     else:
       print("Exception not produced")
 
 run(testUnrankedTensorType)
 
 # CHECK-LABEL: TEST: testMemRefType
 def testMemRefType():
   with Context(), Location.unknown():
     f32 = F32Type.get()
     shape = [2, 3]
     loc = Location.unknown()
     memref = MemRefType.get(shape, f32, memory_space=Attribute.parse("2"))
     # CHECK: memref type: memref<2x3xf32, 2>
     print("memref type:", memref)
     # CHECK: number of affine layout maps: 0
     print("number of affine layout maps:", len(memref.layout))
     # CHECK: memory space: 2
     print("memory space:", memref.memory_space)
 
     layout = AffineMap.get_permutation([1, 0])
     memref_layout = MemRefType.get(shape, f32, [layout])
     # CHECK: memref type: memref<2x3xf32, affine_map<(d0, d1) -> (d1, d0)>>
     print("memref type:", memref_layout)
     assert len(memref_layout.layout) == 1
     # CHECK: memref layout: (d0, d1) -> (d1, d0)
     print("memref layout:", memref_layout.layout[0])
     # CHECK: memory space: <<NULL ATTRIBUTE>>
     print("memory space:", memref_layout.memory_space)
 
     none = NoneType.get()
     try:
       memref_invalid = MemRefType.get(shape, none)
     except ValueError as e:
       # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
       # CHECK: or complex type.
       print(e)
     else:
       print("Exception not produced")
 
 run(testMemRefType)
 
 # CHECK-LABEL: TEST: testUnrankedMemRefType
 def testUnrankedMemRefType():
   with Context(), Location.unknown():
     f32 = F32Type.get()
     loc = Location.unknown()
     unranked_memref = UnrankedMemRefType.get(f32, Attribute.parse("2"))
     # CHECK: unranked memref type: memref<*xf32, 2>
     print("unranked memref type:", unranked_memref)
     try:
       invalid_rank = unranked_memref.rank
     except ValueError as e:
       # CHECK: calling this method requires that the type has a rank.
       print(e)
     else:
       print("Exception not produced")
     try:
       invalid_is_dynamic_dim = unranked_memref.is_dynamic_dim(0)
     except ValueError as e:
       # CHECK: calling this method requires that the type has a rank.
       print(e)
     else:
       print("Exception not produced")
     try:
       invalid_get_dim_size = unranked_memref.get_dim_size(1)
     except ValueError as e:
       # CHECK: calling this method requires that the type has a rank.
       print(e)
     else:
       print("Exception not produced")
 
     none = NoneType.get()
     try:
       memref_invalid = UnrankedMemRefType.get(none, Attribute.parse("2"))
     except ValueError as e:
       # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
       # CHECK: or complex type.
       print(e)
     else:
       print("Exception not produced")
 
 run(testUnrankedMemRefType)
 
 # CHECK-LABEL: TEST: testTupleType
 def testTupleType():
   with Context() as ctx:
     i32 = IntegerType(Type.parse("i32"))
     f32 = F32Type.get()
     vector = VectorType(Type.parse("vector<2x3xf32>"))
     l = [i32, f32, vector]
     tuple_type = TupleType.get_tuple(l)
     # CHECK: tuple type: tuple<i32, f32, vector<2x3xf32>>
     print("tuple type:", tuple_type)
     # CHECK: number of types: 3
     print("number of types:", tuple_type.num_types)
     # CHECK: pos-th type in the tuple type: f32
     print("pos-th type in the tuple type:", tuple_type.get_type(1))
 
 run(testTupleType)
 
 
 # CHECK-LABEL: TEST: testFunctionType
 def testFunctionType():
   with Context() as ctx:
     input_types = [IntegerType.get_signless(32),
                   IntegerType.get_signless(16)]
     result_types = [IndexType.get()]
     func = FunctionType.get(input_types, result_types)
     # CHECK: INPUTS: [Type(i32), Type(i16)]
     print("INPUTS:", func.inputs)
     # CHECK: RESULTS: [Type(index)]
     print("RESULTS:", func.results)
 
 
 run(testFunctionType)

mlir/test/python/ir/value.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
 
 
 # CHECK-LABEL: TEST: testCapsuleConversions
 def testCapsuleConversions():
   ctx = Context()
   ctx.allow_unregistered_dialects = True
   with Location.unknown(ctx):
     i32 = IntegerType.get_signless(32)
     value = Operation.create("custom.op1", results=[i32]).result
     value_capsule = value._CAPIPtr
     assert '"mlir.ir.Value._CAPIPtr"' in repr(value_capsule)
     value2 = Value._CAPICreate(value_capsule)
     assert value2 == value
 
 
 run(testCapsuleConversions)

mlir/test/python/lit.local.cfg

 if not config.enable_bindings_python:
   config.unsupported = True
 config.excludes.add('python_test_ops.td')

mlir/test/python/pass_manager.py

 # RUN: %PYTHON %s 2>&1 | FileCheck %s
 
 import gc, sys
 from mlir.ir import *
 from mlir.passmanager import *
 
 # Log everything to stderr and flush so that we have a unified stream to match
 # errors/info emitted by MLIR to stderr.
 def log(*args):
   print(*args, file=sys.stderr)
   sys.stderr.flush()
 
 def run(f):
   log("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
 
 # Verify capsule interop.
 # CHECK-LABEL: TEST: testCapsule
 def testCapsule():
   with Context():
     pm = PassManager()
     pm_capsule = pm._CAPIPtr
     assert '"mlir.passmanager.PassManager._CAPIPtr"' in repr(pm_capsule)
     pm._testing_release()
     pm1 = PassManager._CAPICreate(pm_capsule)
     assert pm1 is not None  # And does not crash.
 run(testCapsule)
 
 
 # Verify successful round-trip.
 # CHECK-LABEL: TEST: testParseSuccess
 def testParseSuccess():
   with Context():
     # A first import is expected to fail because the pass isn't registered
     # until we import mlir.transforms
     try:
       pm = PassManager.parse("module(func(print-op-stats))")
       # TODO: this error should be propagate to Python but the C API does not help right now.
       # CHECK: error: 'print-op-stats' does not refer to a registered pass or pass pipeline
     except ValueError as e:
       # CHECK: ValueError exception: invalid pass pipeline 'module(func(print-op-stats))'.
       log("ValueError exception:", e)
     else:
       log("Exception not produced")
 
     # This will register the pass and round-trip should be possible now.
     import mlir.transforms
     pm = PassManager.parse("module(func(print-op-stats))")
     # CHECK: Roundtrip: module(func(print-op-stats))
     log("Roundtrip: ", pm)
 run(testParseSuccess)
 
 # Verify failure on unregistered pass.
 # CHECK-LABEL: TEST: testParseFail
 def testParseFail():
   with Context():
     try:
       pm = PassManager.parse("unknown-pass")
     except ValueError as e:
       # CHECK: ValueError exception: invalid pass pipeline 'unknown-pass'.
       log("ValueError exception:", e)
     else:
       log("Exception not produced")
 run(testParseFail)
 
 
 # Verify failure on incorrect level of nesting.
 # CHECK-LABEL: TEST: testInvalidNesting
 def testInvalidNesting():
   with Context():
     try:
       pm = PassManager.parse("func(print-op-graph)")
     except ValueError as e:
       # CHECK: Can't add pass 'ViewOpGraphPass' restricted to 'module' on a PassManager intended to run on 'func', did you intend to nest?
       # CHECK: ValueError exception: invalid pass pipeline 'func(print-op-graph)'.
       log("ValueError exception:", e)
     else:
       log("Exception not produced")
 run(testInvalidNesting)
 
 
 # Verify that a pass manager can execute on IR
 # CHECK-LABEL: TEST: testRun
 def testRunPipeline():
   with Context():
     pm = PassManager.parse("print-op-stats")
     module = Module.parse(r"""func @successfulParse() { return }""")
     pm.run(module)
 # CHECK: Operations encountered:
 # CHECK: func              , 1
 # CHECK: module            , 1
 # CHECK: std.return        , 1
 run(testRunPipeline)

mlir/test/python/python_test_ops.td

 //===-- python_test_ops.td - Python test Op definitions ----*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef PYTHON_TEST_OPS
 #define PYTHON_TEST_OPS
 
 include "mlir/Bindings/Python/Attributes.td"
 include "mlir/IR/OpBase.td"
 
 def Python_Test_Dialect : Dialect {
   let name = "python_test";
   let cppNamespace = "PythonTest";
 }
 class TestOp<string mnemonic, list<OpTrait> traits = []>
     : Op<Python_Test_Dialect, mnemonic, traits>;
 
 def AttributedOp : TestOp<"attributed_op"> {
   let arguments = (ins I32Attr:$mandatory_i32,
                    OptionalAttr<I32Attr>:$optional_i32,
                    UnitAttr:$unit);
 }
 
 #endif // PYTHON_TEST_OPS