diff --git a/.ci/all_requirements.txt b/.ci/all_requirements.txt index dea9646f648ad..ac9682a09bec1 100644 --- a/.ci/all_requirements.txt +++ b/.ci/all_requirements.txt @@ -194,9 +194,9 @@ ml-dtypes==0.5.1 ; python_version < "3.13" \ --hash=sha256:d13755f8e8445b3870114e5b6240facaa7cb0c3361e54beba3e07fa912a6e12b \ --hash=sha256:fd918d4e6a4e0c110e2e05be7a7814d10dc1b95872accbf6512b80a109b71ae1 # via -r mlir/python/requirements.txt -nanobind==2.7.0 \ - --hash=sha256:73b12d0e751d140d6c1bf4b215e18818a8debfdb374f08dc3776ad208d808e74 \ - --hash=sha256:f9f1b160580c50dcf37b6495a0fd5ec61dc0d95dae5f8004f87dd9ad7eb46b34 +nanobind==2.9.2 \ + --hash=sha256:c37957ffd5eac7eda349cff3622ecd32e5ee1244ecc912c99b5bc8188bafd16e \ + --hash=sha256:e7608472de99d375759814cab3e2c94aba3f9ec80e62cfef8ced495ca5c27d6e # via -r mlir/python/requirements.txt numpy==2.0.2 \ --hash=sha256:0123ffdaa88fa4ab64835dcbde75dcdf89c453c922f18dced6e27c90d1d0ec5a \ @@ -383,6 +383,10 @@ swig==4.3.1 \ --hash=sha256:efec16327029f682f649a26da726bb0305be8800bd0f1fa3e81bf0769cf5b476 \ --hash=sha256:fc496c0d600cf1bb2d91e28d3d6eae9c4301e5ea7a0dec5a4281b5efed4245a8 # via -r lldb/test/requirements.txt +typing-extensions==4.15.0 \ + --hash=sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466 \ + --hash=sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548 + # via -r mlir/python/requirements.txt urllib3==2.5.0 \ --hash=sha256:3fc47733c7e419d4bc3f6b3dc2b4f890bb743906a30d56ba4a5bfa4bbff92760 \ --hash=sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake index 2b883558d33c6..f25595116edca 100644 --- a/mlir/cmake/modules/AddMLIRPython.cmake +++ b/mlir/cmake/modules/AddMLIRPython.cmake @@ -99,6 +99,89 @@ function(declare_mlir_python_sources name) endif() endfunction() +# Function: mlir_generate_type_stubs +# Turns on automatic type stub generation for extension modules. +# Specifically, performs add_custom_command to run nanobind's stubgen on an extension module. +# +# Arguments: +# MODULE_NAME: The fully-qualified name of the extension module (used for importing in python). +# DEPENDS_TARGETS: List of targets these type stubs depend on being built; usually corresponding to the +# specific extension module (e.g., something like StandalonePythonModules.extension._standaloneDialectsNanobind.dso) +# and the core bindings extension module (e.g., something like StandalonePythonModules.extension._mlir.dso). +# OUTPUT_DIR: The root output directory to emit the type stubs into. +# OUTPUTS: List of expected outputs. +# DEPENDS_TARGET_SRC_DEPS: List of cpp sources for extension library (for generating a DEPFILE). +# IMPORT_PATHS: List of paths to add to PYTHONPATH for stubgen. +# PATTERN_FILE: (Optional) Pattern file (see https://nanobind.readthedocs.io/en/latest/typing.html#pattern-files). +# Outputs: +# NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on. +function(mlir_generate_type_stubs) + cmake_parse_arguments(ARG + "" + "MODULE_NAME;OUTPUT_DIR;PATTERN_FILE" + "IMPORT_PATHS;DEPENDS_TARGETS;OUTPUTS;DEPENDS_TARGET_SRC_DEPS" + ${ARGN}) + + # for people doing find_package(nanobind) + if(EXISTS ${nanobind_DIR}/../src/stubgen.py) + set(NB_STUBGEN "${nanobind_DIR}/../src/stubgen.py") + elseif(EXISTS ${nanobind_DIR}/../stubgen.py) + set(NB_STUBGEN "${nanobind_DIR}/../stubgen.py") + # for people using FetchContent_Declare and FetchContent_MakeAvailable + elseif(EXISTS ${nanobind_SOURCE_DIR}/src/stubgen.py) + set(NB_STUBGEN "${nanobind_SOURCE_DIR}/src/stubgen.py") + elseif(EXISTS ${nanobind_SOURCE_DIR}/stubgen.py) + set(NB_STUBGEN "${nanobind_SOURCE_DIR}/stubgen.py") + else() + message(FATAL_ERROR "mlir_generate_type_stubs(): could not locate 'stubgen.py'!") + endif() + + file(REAL_PATH "${NB_STUBGEN}" NB_STUBGEN) + set(_import_paths) + foreach(_import_path IN LISTS ARG_IMPORT_PATHS) + file(REAL_PATH "${_import_path}" _import_path) + list(APPEND _import_paths "-i;${_import_path}") + endforeach() + set(_nb_stubgen_cmd + "${Python_EXECUTABLE}" + "${NB_STUBGEN}" + --module + "${ARG_MODULE_NAME}" + "${_import_paths}" + --recursive + --include-private + --output-dir + "${ARG_OUTPUT_DIR}") + if(ARG_PATTERN_FILE) + list(APPEND _nb_stubgen_cmd "-p;${ARG_PATTERN_FILE}") + list(APPEND ARG_DEPENDS_TARGETS "${ARG_PATTERN_FILE}") + endif() + + list(TRANSFORM ARG_OUTPUTS PREPEND "${ARG_OUTPUT_DIR}/" OUTPUT_VARIABLE _generated_type_stubs) + set(_depfile "${ARG_OUTPUT_DIR}/${ARG_MODULE_NAME}.d") + if ((NOT EXISTS ${_depfile}) AND ARG_DEPENDS_TARGET_SRC_DEPS) + list(JOIN ARG_DEPENDS_TARGET_SRC_DEPS " " _depfiles) + list(TRANSFORM _generated_type_stubs APPEND ": ${_depfiles}" OUTPUT_VARIABLE _depfiles) + list(JOIN _depfiles "\n" _depfiles) + file(GENERATE OUTPUT "${_depfile}" CONTENT "${_depfiles}") + endif() + + message(DEBUG "Generating type-stubs outputs ${_generated_type_stubs}") + add_custom_command( + OUTPUT ${_generated_type_stubs} + COMMAND ${_nb_stubgen_cmd} + WORKING_DIRECTORY "${CMAKE_CURRENT_FUNCTION_LIST_DIR}" + DEPENDS "${ARG_DEPENDS_TARGETS}" + DEPFILE "${_depfile}" + COMMENT "Generating type stubs using: ${_nb_stubgen_cmd}" + ) + + list(JOIN ARG_DEPENDS_TARGETS "." _name) + set(_name "${_name}.${ARG_MODULE_NAME}.type_stubs") + add_custom_target("${_name}" DEPENDS ${_generated_type_stubs}) + set(NB_STUBGEN_CUSTOM_TARGET "${_name}" PARENT_SCOPE) +endfunction() + # Function: declare_mlir_python_extension # Declares a buildable python extension from C++ source files. The built # module is considered a python source file and included as everything else. @@ -678,26 +761,28 @@ function(add_mlir_python_extension libname extname) # the super project handle compile options as it wishes. get_property(NB_LIBRARY_TARGET_NAME TARGET ${libname} PROPERTY LINK_LIBRARIES) target_compile_options(${NB_LIBRARY_TARGET_NAME} - PRIVATE - -Wall -Wextra -Wpedantic - -Wno-c++98-compat-extra-semi - -Wno-cast-qual - -Wno-covered-switch-default - -Wno-nested-anon-types - -Wno-unused-parameter - -Wno-zero-length-array - ${eh_rtti_enable}) + PRIVATE + -Wall -Wextra -Wpedantic + -Wno-c++98-compat-extra-semi + -Wno-cast-qual + -Wno-covered-switch-default + -Wno-deprecated-literal-operator + -Wno-nested-anon-types + -Wno-unused-parameter + -Wno-zero-length-array + ${eh_rtti_enable}) target_compile_options(${libname} - PRIVATE - -Wall -Wextra -Wpedantic - -Wno-c++98-compat-extra-semi - -Wno-cast-qual - -Wno-covered-switch-default - -Wno-nested-anon-types - -Wno-unused-parameter - -Wno-zero-length-array - ${eh_rtti_enable}) + PRIVATE + -Wall -Wextra -Wpedantic + -Wno-c++98-compat-extra-semi + -Wno-cast-qual + -Wno-covered-switch-default + -Wno-deprecated-literal-operator + -Wno-nested-anon-types + -Wno-unused-parameter + -Wno-zero-length-array + ${eh_rtti_enable}) endif() if(APPLE) diff --git a/mlir/cmake/modules/MLIRDetectPythonEnv.cmake b/mlir/cmake/modules/MLIRDetectPythonEnv.cmake index f7a6fa6248440..d18f8c06158b2 100644 --- a/mlir/cmake/modules/MLIRDetectPythonEnv.cmake +++ b/mlir/cmake/modules/MLIRDetectPythonEnv.cmake @@ -54,7 +54,7 @@ macro(mlir_configure_python_dev_packages) "extension = '${PYTHON_MODULE_EXTENSION}") mlir_detect_nanobind_install() - find_package(nanobind 2.4 CONFIG REQUIRED) + find_package(nanobind 2.9 CONFIG REQUIRED) message(STATUS "Found nanobind v${nanobind_VERSION}: ${nanobind_INCLUDE_DIR}") message(STATUS "Python prefix = '${PYTHON_MODULE_PREFIX}', " "suffix = '${PYTHON_MODULE_SUFFIX}', " diff --git a/mlir/examples/standalone/CMakeLists.txt b/mlir/examples/standalone/CMakeLists.txt index 03627c0c10496..e2bcda7fa6f0b 100644 --- a/mlir/examples/standalone/CMakeLists.txt +++ b/mlir/examples/standalone/CMakeLists.txt @@ -6,9 +6,9 @@ set(CMAKE_BUILD_WITH_INSTALL_NAME_DIR ON) set(CMAKE_CXX_STANDARD 17 CACHE STRING "C++ standard to conform to") if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR) + set(EXTERNAL_PROJECT_BUILD FALSE) find_package(MLIR REQUIRED CONFIG) - # Define the default argument to use by `lit` when testing. set(LLVM_LIT_ARGS "-sv" CACHE STRING "Default options for lit") @@ -27,11 +27,19 @@ if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR) include(AddMLIR) include(HandleLLVMOptions) else() + set(EXTERNAL_PROJECT_BUILD TRUE) # Build via external projects mechanism + set(LLVM_RUNTIME_OUTPUT_INTDIR ${CMAKE_BINARY_DIR}/bin) + set(LLVM_LIBRARY_OUTPUT_INTDIR ${CMAKE_BINARY_DIR}/tools/Standalone/lib) + + set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR}) set(MLIR_MAIN_SRC_DIR ${LLVM_MAIN_SRC_DIR}/../mlir) set(MLIR_INCLUDE_DIR ${MLIR_MAIN_SRC_DIR}/include) set(MLIR_GENERATED_INCLUDE_DIR ${LLVM_BINARY_DIR}/tools/mlir/include) set(MLIR_INCLUDE_DIRS "${MLIR_INCLUDE_DIR};${MLIR_GENERATED_INCLUDE_DIR}") + + list(APPEND CMAKE_MODULE_PATH "${MLIR_MAIN_SRC_DIR}/cmake/modules") + list(APPEND CMAKE_MODULE_PATH "${LLVM_MAIN_SRC_DIR}/cmake/modules") endif() if(MLIR_ENABLE_BINDINGS_PYTHON) @@ -54,8 +62,9 @@ if(MLIR_ENABLE_BINDINGS_PYTHON) message(STATUS "Enabling Python API") include(MLIRDetectPythonEnv) mlir_configure_python_dev_packages() + # Note: for EXTERNAL_PROJECT_BUILD this must be set from the command line. set(MLIR_PYTHON_PACKAGE_PREFIX "mlir_standalone" CACHE STRING "" FORCE) - set(MLIR_BINDINGS_PYTHON_INSTALL_PREFIX "python_packages/standalone/mlir_standalone" CACHE STRING "" FORCE) + set(MLIR_BINDINGS_PYTHON_INSTALL_PREFIX "python_packages/standalone/${MLIR_PYTHON_PACKAGE_PREFIX}" CACHE STRING "" FORCE) add_subdirectory(python) endif() add_subdirectory(test) diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt index 1ab27ce3b533a..d48c5bcdde137 100644 --- a/mlir/examples/standalone/python/CMakeLists.txt +++ b/mlir/examples/standalone/python/CMakeLists.txt @@ -18,6 +18,7 @@ declare_mlir_dialect_python_bindings( SOURCES dialects/standalone_pybind11.py dialects/standalone_nanobind.py + _mlir_libs/_standaloneDialectsNanobind/py.typed DIALECT_NAME standalone) @@ -68,16 +69,85 @@ add_mlir_python_common_capi_library(StandalonePythonCAPI # Instantiation of all Python modules ################################################################################ +set(StandalonePythonModules_ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}") + +# Everything here is very tightly coupled. See the ample descriptions at the bottom of +# mlir/python/CMakeLists.txt. + +# For a non-external projects build (e.g., installed distro) the type gen targets for the core _mlir module +# need to be re-declared. On the contrary, for an external projects build, the MLIRPythonExtension.Core.type_stub_gen +# target already exists and can just be added to DECLARED_SOURCES (see below). +if(NOT EXTERNAL_PROJECT_BUILD) + set(_core_type_stub_sources + _mlir/__init__.pyi + _mlir/ir.pyi + _mlir/passmanager.pyi + _mlir/rewrite.pyi + ) + get_target_property(_core_extension_srcs MLIRPythonExtension.Core INTERFACE_SOURCES) + mlir_generate_type_stubs( + MODULE_NAME _mlir + DEPENDS_TARGETS StandalonePythonModules.extension._mlir.dso + OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs" + OUTPUTS "${_core_type_stub_sources}" + DEPENDS_TARGET_SRC_DEPS "${_core_extension_srcs}" + IMPORT_PATHS "${StandalonePythonModules_ROOT_PREFIX}/_mlir_libs" + ) + set(_mlir_typestub_gen_target "${NB_STUBGEN_CUSTOM_TARGET}") + + list(TRANSFORM _core_type_stub_sources PREPEND "_mlir_libs/") + declare_mlir_python_sources( + StandalonePythonExtension.Core.type_stub_gen + ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs" + ADD_TO_PARENT StandalonePythonSources + SOURCES "${_core_type_stub_sources}" + ) +endif() + +get_target_property(_standalone_extension_srcs StandalonePythonSources.NanobindExtension INTERFACE_SOURCES) +mlir_generate_type_stubs( + MODULE_NAME mlir_standalone._mlir_libs._standaloneDialectsNanobind + DEPENDS_TARGETS + StandalonePythonModules.extension._mlir.dso + StandalonePythonModules.extension._standaloneDialectsNanobind.dso + OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs" + OUTPUTS + _standaloneDialectsNanobind/__init__.pyi + _standaloneDialectsNanobind/standalone.pyi + DEPENDS_TARGET_SRC_DEPS "${_standalone_extension_srcs}" + IMPORT_PATHS "${StandalonePythonModules_ROOT_PREFIX}/.." +) +set(_standaloneDialectsNanobind_typestub_gen_target "${NB_STUBGEN_CUSTOM_TARGET}") + +declare_mlir_python_sources( + StandalonePythonSources.type_stub_gen + ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs" + ADD_TO_PARENT StandalonePythonSources + SOURCES + _mlir_libs/_standaloneDialectsNanobind/__init__.pyi + _mlir_libs/_standaloneDialectsNanobind/standalone.pyi +) +set(_declared_sources + StandalonePythonSources + # TODO: Remove this in favor of showing fine grained registration once + # available. + MLIRPythonExtension.RegisterEverything + MLIRPythonSources.Core + MLIRPythonSources.Dialects.builtin +) +# For an external projects build, the MLIRPythonExtension.Core.type_stub_gen +# target already exists and can just be added to DECLARED_SOURCES. +if(EXTERNAL_PROJECT_BUILD) + list(APPEND _declared_sources MLIRPythonExtension.Core.type_stub_gen) +endif() add_mlir_python_modules(StandalonePythonModules - ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}" + ROOT_PREFIX "${StandalonePythonModules_ROOT_PREFIX}" INSTALL_PREFIX "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}" - DECLARED_SOURCES - StandalonePythonSources - # TODO: Remove this in favor of showing fine grained registration once - # available. - MLIRPythonExtension.RegisterEverything - MLIRPythonSources.Core - MLIRPythonSources.Dialects.builtin + DECLARED_SOURCES "${_declared_sources}" COMMON_CAPI_LINK_LIBS StandalonePythonCAPI - ) +) +if(NOT EXTERNAL_PROJECT_BUILD) + add_dependencies(StandalonePythonModules "${_mlir_typestub_gen_target}") +endif() +add_dependencies(StandalonePythonModules "${_standaloneDialectsNanobind_typestub_gen_target}") diff --git a/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp b/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp index 189ebac368bf5..e06ec3b6472b8 100644 --- a/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp +++ b/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp @@ -30,5 +30,9 @@ NB_MODULE(_standaloneDialectsNanobind, m) { mlirDialectHandleLoadDialect(handle, context); } }, - nb::arg("context").none() = nb::none(), nb::arg("load") = true); + nb::arg("context").none() = nb::none(), nb::arg("load") = true, + // clang-format off + nb::sig("def register_dialect(context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") ", load: bool = True) -> None") + // clang-format on + ); } diff --git a/mlir/examples/standalone/python/mlir_standalone/_mlir_libs/_standaloneDialectsNanobind/py.typed b/mlir/examples/standalone/python/mlir_standalone/_mlir_libs/_standaloneDialectsNanobind/py.typed new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h index 1428d5ccf00f4..8744d8d0e4bca 100644 --- a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h +++ b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h @@ -512,8 +512,13 @@ class mlir_attribute_subclass : public pure_subclass { .attr("replace")(superCls.attr("__name__"), captureTypeName); }); if (getTypeIDFunction) { - def_staticmethod("get_static_typeid", - [getTypeIDFunction]() { return getTypeIDFunction(); }); + def_staticmethod( + "get_static_typeid", + [getTypeIDFunction]() { return getTypeIDFunction(); }, + // clang-format off + nanobind::sig("def get_static_typeid() -> " MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID")) + // clang-format on + ); nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir")) .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)( getTypeIDFunction())(nanobind::cpp_function( @@ -582,8 +587,9 @@ class mlir_type_subclass : public pure_subclass { // 'isinstance' method. static const char kIsinstanceSig[] = - "def isinstance(other_type: " MAKE_MLIR_PYTHON_QUALNAME( - "ir") ".Type) -> bool"; + // clang-format off + "def isinstance(other_type: " MAKE_MLIR_PYTHON_QUALNAME("ir.Type") ") -> bool"; + // clang-format on def_staticmethod( "isinstance", [isaFunction](MlirType other) { return isaFunction(other); }, @@ -599,8 +605,13 @@ class mlir_type_subclass : public pure_subclass { // `def_property_readonly_static` is not available in `pure_subclass` and // we do not want to introduce the complexity that pybind uses to // implement it. - def_staticmethod("get_static_typeid", - [getTypeIDFunction]() { return getTypeIDFunction(); }); + def_staticmethod( + "get_static_typeid", + [getTypeIDFunction]() { return getTypeIDFunction(); }, + // clang-format off + nanobind::sig("def get_static_typeid() -> " MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID")) + // clang-format on + ); nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir")) .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)( getTypeIDFunction())(nanobind::cpp_function( @@ -665,8 +676,9 @@ class mlir_value_subclass : public pure_subclass { // 'isinstance' method. static const char kIsinstanceSig[] = - "def isinstance(other_value: " MAKE_MLIR_PYTHON_QUALNAME( - "ir") ".Value) -> bool"; + // clang-format off + "def isinstance(other_value: " MAKE_MLIR_PYTHON_QUALNAME("ir.Value") ") -> bool"; + // clang-format on def_staticmethod( "isinstance", [isaFunction](MlirValue other) { return isaFunction(other); }, diff --git a/mlir/lib/Bindings/Python/DialectPDL.cpp b/mlir/lib/Bindings/Python/DialectPDL.cpp index ac163b59a0fb5..1acb41080f711 100644 --- a/mlir/lib/Bindings/Python/DialectPDL.cpp +++ b/mlir/lib/Bindings/Python/DialectPDL.cpp @@ -68,6 +68,8 @@ static void populateDialectPDLSubmodule(const nanobind::module_ &m) { rangeType.def_property_readonly( "element_type", [](MlirType type) { return mlirPDLRangeTypeGetElementType(type); }, + nb::sig( + "def element_type(self) -> " MAKE_MLIR_PYTHON_QUALNAME("ir.Type")), "Get the element type."); //===-------------------------------------------------------------------===// diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp index 4e0ade41fb708..7818caf2e8a55 100644 --- a/mlir/lib/Bindings/Python/IRAttributes.cpp +++ b/mlir/lib/Bindings/Python/IRAttributes.cpp @@ -167,7 +167,7 @@ struct nb_buffer_info { }; class nb_buffer : public nb::object { - NB_OBJECT_DEFAULT(nb_buffer, object, "buffer", PyObject_CheckBuffer); + NB_OBJECT_DEFAULT(nb_buffer, object, "Buffer", PyObject_CheckBuffer); nb_buffer_info request() const { int flags = PyBUF_STRIDES | PyBUF_FORMAT; @@ -252,8 +252,13 @@ class PyAffineMapAttribute : public PyConcreteAttribute { return PyAffineMapAttribute(affineMap.getContext(), attr); }, nb::arg("affine_map"), "Gets an attribute wrapping an AffineMap."); - c.def_prop_ro("value", mlirAffineMapAttrGetValue, - "Returns the value of the AffineMap attribute"); + c.def_prop_ro( + "value", + [](PyAffineMapAttribute &self) { + return PyAffineMap(self.getContext(), + mlirAffineMapAttrGetValue(self)); + }, + "Returns the value of the AffineMap attribute"); } }; @@ -480,11 +485,13 @@ class PyArrayAttribute : public PyConcreteAttribute { PyArrayAttributeIterator &dunderIter() { return *this; } - MlirAttribute dunderNext() { + nb::typed dunderNext() { // TODO: Throw is an inefficient way to stop iteration. if (nextIndex >= mlirArrayAttrGetNumElements(attr.get())) throw nb::stop_iteration(); - return mlirArrayAttrGetElement(attr.get(), nextIndex++); + return PyAttribute(this->attr.getContext(), + mlirArrayAttrGetElement(attr.get(), nextIndex++)) + .maybeDownCast(); } static void bind(nb::module_ &m) { @@ -517,12 +524,13 @@ class PyArrayAttribute : public PyConcreteAttribute { }, nb::arg("attributes"), nb::arg("context") = nb::none(), "Gets a uniqued Array attribute"); - c.def("__getitem__", - [](PyArrayAttribute &arr, intptr_t i) { - if (i >= mlirArrayAttrGetNumElements(arr)) - throw nb::index_error("ArrayAttribute index out of range"); - return arr.getItem(i); - }) + c.def( + "__getitem__", + [](PyArrayAttribute &arr, intptr_t i) { + if (i >= mlirArrayAttrGetNumElements(arr)) + throw nb::index_error("ArrayAttribute index out of range"); + return PyAttribute(arr.getContext(), arr.getItem(i)).maybeDownCast(); + }) .def("__len__", [](const PyArrayAttribute &arr) { return mlirArrayAttrGetNumElements(arr); @@ -611,10 +619,12 @@ class PyIntegerAttribute : public PyConcreteAttribute { "Returns the value of the integer attribute"); c.def("__int__", toPyInt, "Converts the value of the integer attribute to a Python int"); - c.def_prop_ro_static("static_typeid", - [](nb::object & /*class*/) -> MlirTypeID { - return mlirIntegerAttrGetTypeID(); - }); + c.def_prop_ro_static( + "static_typeid", + [](nb::object & /*class*/) { + return PyTypeID(mlirIntegerAttrGetTypeID()); + }, + nanobind::sig("def static_typeid(/) -> TypeID")); } private: @@ -657,8 +667,8 @@ class PySymbolRefAttribute : public PyConcreteAttribute { static constexpr const char *pyClassName = "SymbolRefAttr"; using PyConcreteAttribute::PyConcreteAttribute; - static MlirAttribute fromList(const std::vector &symbols, - PyMlirContext &context) { + static PySymbolRefAttribute fromList(const std::vector &symbols, + PyMlirContext &context) { if (symbols.empty()) throw std::runtime_error("SymbolRefAttr must be composed of at least " "one symbol."); @@ -668,8 +678,10 @@ class PySymbolRefAttribute : public PyConcreteAttribute { referenceAttrs.push_back( mlirFlatSymbolRefAttrGet(context.get(), toMlirStringRef(symbols[i]))); } - return mlirSymbolRefAttrGet(context.get(), rootSymbol, - referenceAttrs.size(), referenceAttrs.data()); + return PySymbolRefAttribute(context.getRef(), + mlirSymbolRefAttrGet(context.get(), rootSymbol, + referenceAttrs.size(), + referenceAttrs.data())); } static void bindDerived(ClassTy &c) { @@ -746,7 +758,11 @@ class PyOpaqueAttribute : public PyConcreteAttribute { return PyOpaqueAttribute(context->getRef(), attr); }, nb::arg("dialect_namespace"), nb::arg("buffer"), nb::arg("type"), - nb::arg("context") = nb::none(), "Gets an Opaque attribute."); + nb::arg("context") = nb::none(), + // clang-format off + nb::sig("def get(dialect_namespace: str, buffer: typing_extensions.Buffer, type: Type, context: Context | None = None) -> OpaqueAttr"), + // clang-format on + "Gets an Opaque attribute."); c.def_prop_ro( "dialect_namespace", [](PyOpaqueAttribute &self) { @@ -764,59 +780,6 @@ class PyOpaqueAttribute : public PyConcreteAttribute { } }; -class PyStringAttribute : public PyConcreteAttribute { -public: - static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAString; - static constexpr const char *pyClassName = "StringAttr"; - using PyConcreteAttribute::PyConcreteAttribute; - static constexpr GetTypeIDFunctionTy getTypeIdFunction = - mlirStringAttrGetTypeID; - - static void bindDerived(ClassTy &c) { - c.def_static( - "get", - [](const std::string &value, DefaultingPyMlirContext context) { - MlirAttribute attr = - mlirStringAttrGet(context->get(), toMlirStringRef(value)); - return PyStringAttribute(context->getRef(), attr); - }, - nb::arg("value"), nb::arg("context") = nb::none(), - "Gets a uniqued string attribute"); - c.def_static( - "get", - [](const nb::bytes &value, DefaultingPyMlirContext context) { - MlirAttribute attr = - mlirStringAttrGet(context->get(), toMlirStringRef(value)); - return PyStringAttribute(context->getRef(), attr); - }, - nb::arg("value"), nb::arg("context") = nb::none(), - "Gets a uniqued string attribute"); - c.def_static( - "get_typed", - [](PyType &type, const std::string &value) { - MlirAttribute attr = - mlirStringAttrTypedGet(type, toMlirStringRef(value)); - return PyStringAttribute(type.getContext(), attr); - }, - nb::arg("type"), nb::arg("value"), - "Gets a uniqued string attribute associated to a type"); - c.def_prop_ro( - "value", - [](PyStringAttribute &self) { - MlirStringRef stringRef = mlirStringAttrGetValue(self); - return nb::str(stringRef.data, stringRef.length); - }, - "Returns the value of the string attribute"); - c.def_prop_ro( - "value_bytes", - [](PyStringAttribute &self) { - MlirStringRef stringRef = mlirStringAttrGetValue(self); - return nb::bytes(stringRef.data, stringRef.length); - }, - "Returns the value of the string attribute as `bytes`"); - } -}; - // TODO: Support construction of string elements. class PyDenseElementsAttribute : public PyConcreteAttribute { @@ -1028,11 +991,14 @@ class PyDenseElementsAttribute PyDenseElementsAttribute::bf_releasebuffer; #endif c.def("__len__", &PyDenseElementsAttribute::dunderLen) - .def_static("get", PyDenseElementsAttribute::getFromBuffer, - nb::arg("array"), nb::arg("signless") = true, - nb::arg("type") = nb::none(), nb::arg("shape") = nb::none(), - nb::arg("context") = nb::none(), - kDenseElementsAttrGetDocstring) + .def_static( + "get", PyDenseElementsAttribute::getFromBuffer, nb::arg("array"), + nb::arg("signless") = true, nb::arg("type") = nb::none(), + nb::arg("shape") = nb::none(), nb::arg("context") = nb::none(), + // clang-format off + nb::sig("def get(array: typing_extensions.Buffer, signless: bool = True, type: Type | None = None, shape: Sequence[int] | None = None, context: Context | None = None) -> DenseElementsAttr"), + // clang-format on + kDenseElementsAttrGetDocstring) .def_static("get", PyDenseElementsAttribute::getFromList, nb::arg("attrs"), nb::arg("type") = nb::none(), nb::arg("context") = nb::none(), @@ -1048,7 +1014,9 @@ class PyDenseElementsAttribute if (!mlirDenseElementsAttrIsSplat(self)) throw nb::value_error( "get_splat_value called on a non-splat attribute"); - return mlirDenseElementsAttrGetSplatValue(self); + return PyAttribute(self.getContext(), + mlirDenseElementsAttrGetSplatValue(self)) + .maybeDownCast(); }); } @@ -1509,6 +1477,9 @@ class PyDenseResourceElementsAttribute nb::arg("array"), nb::arg("name"), nb::arg("type"), nb::arg("alignment") = nb::none(), nb::arg("is_mutable") = false, nb::arg("context") = nb::none(), + // clang-format off + nb::sig("def get_from_buffer(array: typing_extensions.Buffer, name: str, type: Type, alignment: int | None = None, is_mutable: bool = False, context: Context | None = None) -> DenseResourceElementsAttr"), + // clang-format on kDenseResourceElementsAttrGetFromBufferDocstring); } }; @@ -1556,7 +1527,7 @@ class PyDictAttribute : public PyConcreteAttribute { mlirDictionaryAttrGetElementByName(self, toMlirStringRef(name)); if (mlirAttributeIsNull(attr)) throw nb::key_error("attempt to access a non-existent attribute"); - return attr; + return PyAttribute(self.getContext(), attr).maybeDownCast(); }); c.def("__getitem__", [](PyDictAttribute &self, intptr_t index) { if (index < 0 || index >= self.dunderLen()) { @@ -1624,7 +1595,8 @@ class PyTypeAttribute : public PyConcreteAttribute { nb::arg("value"), nb::arg("context") = nb::none(), "Gets a uniqued Type attribute"); c.def_prop_ro("value", [](PyTypeAttribute &self) { - return mlirTypeAttrGetValue(self.get()); + return PyType(self.getContext(), mlirTypeAttrGetValue(self.get())) + .maybeDownCast(); }); } }; @@ -1761,6 +1733,50 @@ nb::object symbolRefOrFlatSymbolRefAttributeCaster(PyAttribute &pyAttribute) { } // namespace +void PyStringAttribute::bindDerived(ClassTy &c) { + c.def_static( + "get", + [](const std::string &value, DefaultingPyMlirContext context) { + MlirAttribute attr = + mlirStringAttrGet(context->get(), toMlirStringRef(value)); + return PyStringAttribute(context->getRef(), attr); + }, + nb::arg("value"), nb::arg("context") = nb::none(), + "Gets a uniqued string attribute"); + c.def_static( + "get", + [](const nb::bytes &value, DefaultingPyMlirContext context) { + MlirAttribute attr = + mlirStringAttrGet(context->get(), toMlirStringRef(value)); + return PyStringAttribute(context->getRef(), attr); + }, + nb::arg("value"), nb::arg("context") = nb::none(), + "Gets a uniqued string attribute"); + c.def_static( + "get_typed", + [](PyType &type, const std::string &value) { + MlirAttribute attr = + mlirStringAttrTypedGet(type, toMlirStringRef(value)); + return PyStringAttribute(type.getContext(), attr); + }, + nb::arg("type"), nb::arg("value"), + "Gets a uniqued string attribute associated to a type"); + c.def_prop_ro( + "value", + [](PyStringAttribute &self) { + MlirStringRef stringRef = mlirStringAttrGetValue(self); + return nb::str(stringRef.data, stringRef.length); + }, + "Returns the value of the string attribute"); + c.def_prop_ro( + "value_bytes", + [](PyStringAttribute &self) { + MlirStringRef stringRef = mlirStringAttrGetValue(self); + return nb::bytes(stringRef.data, stringRef.length); + }, + "Returns the value of the string attribute as `bytes`"); +} + void mlir::python::populateIRAttributes(nb::module_ &m) { PyAffineMapAttribute::bind(m); PyDenseBoolArrayAttribute::bind(m); diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp index 10360e448858c..609502041f4ae 100644 --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -513,7 +513,7 @@ class PyOperationIterator { PyOperationIterator &dunderIter() { return *this; } - nb::object dunderNext() { + nb::typed dunderNext() { parentOperation->checkValid(); if (mlirOperationIsNull(next)) { throw nb::stop_iteration(); @@ -562,7 +562,7 @@ class PyOperationList { return count; } - nb::object dunderGetItem(intptr_t index) { + nb::typed dunderGetItem(intptr_t index) { parentOperation->checkValid(); if (index < 0) { index += dunderLen(); @@ -598,7 +598,7 @@ class PyOpOperand { public: PyOpOperand(MlirOpOperand opOperand) : opOperand(opOperand) {} - nb::object getOwner() { + PyOpView getOwner() { MlirOperation owner = mlirOpOperandGetOwner(opOperand); PyMlirContextRef context = PyMlirContext::forContext(mlirOperationGetContext(owner)); @@ -1534,7 +1534,7 @@ nb::object PyOperation::create(std::string_view name, return created.getObject(); } -nb::object PyOperation::clone(const nb::object &maybeIp) { +nb::typed PyOperation::clone(const nb::object &maybeIp) { MlirOperation clonedOperation = mlirOperationClone(operation); PyOperationRef cloned = PyOperation::createDetached(getContext(), clonedOperation); @@ -1543,7 +1543,7 @@ nb::object PyOperation::clone(const nb::object &maybeIp) { return cloned->createOpView(); } -nb::object PyOperation::createOpView() { +nb::typed PyOperation::createOpView() { checkValid(); MlirIdentifier ident = mlirOperationGetName(get()); MlirStringRef identStr = mlirIdentifierStr(ident); @@ -1638,12 +1638,14 @@ class PyOpResult : public PyConcreteValue { /// Returns the list of types of the values held by container. template -static std::vector getValueTypes(Container &container, - PyMlirContextRef &context) { - std::vector result; +static std::vector> +getValueTypes(Container &container, PyMlirContextRef &context) { + std::vector> result; result.reserve(container.size()); for (int i = 0, e = container.size(); i < e; ++i) { - result.push_back(mlirValueGetType(container.getElement(i).get())); + result.push_back(PyType(context->getRef(), + mlirValueGetType(container.getElement(i).get())) + .maybeDownCast()); } return result; } @@ -2131,6 +2133,20 @@ PyAttribute PyAttribute::createFromCapsule(nb::object capsule) { PyMlirContext::forContext(mlirAttributeGetContext(rawAttr)), rawAttr); } +nb::typed PyAttribute::maybeDownCast() { + MlirTypeID mlirTypeID = mlirAttributeGetTypeID(this->get()); + assert(!mlirTypeIDIsNull(mlirTypeID) && + "mlirTypeID was expected to be non-null."); + std::optional typeCaster = PyGlobals::get().lookupTypeCaster( + mlirTypeID, mlirAttributeGetDialect(this->get())); + // nb::rv_policy::move means use std::move to move the return value + // contents into a new instance that will be owned by Python. + nb::object thisObj = nb::cast(this, nb::rv_policy::move); + if (!typeCaster) + return thisObj; + return typeCaster.value()(thisObj); +} + //------------------------------------------------------------------------------ // PyNamedAttribute. //------------------------------------------------------------------------------ @@ -2163,6 +2179,20 @@ PyType PyType::createFromCapsule(nb::object capsule) { rawType); } +nb::typed PyType::maybeDownCast() { + MlirTypeID mlirTypeID = mlirTypeGetTypeID(this->get()); + assert(!mlirTypeIDIsNull(mlirTypeID) && + "mlirTypeID was expected to be non-null."); + std::optional typeCaster = PyGlobals::get().lookupTypeCaster( + mlirTypeID, mlirTypeGetDialect(this->get())); + // nb::rv_policy::move means use std::move to move the return value + // contents into a new instance that will be owned by Python. + nb::object thisObj = nb::cast(this, nb::rv_policy::move); + if (!typeCaster) + return thisObj; + return typeCaster.value()(thisObj); +} + //------------------------------------------------------------------------------ // PyTypeID. //------------------------------------------------------------------------------ @@ -2189,7 +2219,7 @@ nb::object PyValue::getCapsule() { return nb::steal(mlirPythonValueToCapsule(get())); } -nb::object PyValue::maybeDownCast() { +nanobind::typed PyValue::maybeDownCast() { MlirType type = mlirValueGetType(get()); MlirTypeID mlirTypeID = mlirTypeGetTypeID(type); assert(!mlirTypeIDIsNull(mlirTypeID) && @@ -2233,7 +2263,8 @@ PySymbolTable::PySymbolTable(PyOperationBase &operation) } } -nb::object PySymbolTable::dunderGetItem(const std::string &name) { +nb::typed +PySymbolTable::dunderGetItem(const std::string &name) { operation->checkValid(); MlirOperation symbol = mlirSymbolTableLookup( symbolTable, mlirStringRefCreate(name.data(), name.length())); @@ -2261,17 +2292,19 @@ void PySymbolTable::dunderDel(const std::string &name) { erase(nb::cast(operation)); } -MlirAttribute PySymbolTable::insert(PyOperationBase &symbol) { +PyStringAttribute PySymbolTable::insert(PyOperationBase &symbol) { operation->checkValid(); symbol.getOperation().checkValid(); MlirAttribute symbolAttr = mlirOperationGetAttributeByName( symbol.getOperation().get(), mlirSymbolTableGetSymbolAttributeName()); if (mlirAttributeIsNull(symbolAttr)) throw nb::value_error("Expected operation to have a symbol name."); - return mlirSymbolTableInsert(symbolTable, symbol.getOperation().get()); + return PyStringAttribute( + symbol.getOperation().getContext(), + mlirSymbolTableInsert(symbolTable, symbol.getOperation().get())); } -MlirAttribute PySymbolTable::getSymbolName(PyOperationBase &symbol) { +PyStringAttribute PySymbolTable::getSymbolName(PyOperationBase &symbol) { // Op must already be a symbol. PyOperation &operation = symbol.getOperation(); operation.checkValid(); @@ -2280,7 +2313,8 @@ MlirAttribute PySymbolTable::getSymbolName(PyOperationBase &symbol) { mlirOperationGetAttributeByName(operation.get(), attrName); if (mlirAttributeIsNull(existingNameAttr)) throw nb::value_error("Expected operation to have a symbol name."); - return existingNameAttr; + return PyStringAttribute(symbol.getOperation().getContext(), + existingNameAttr); } void PySymbolTable::setSymbolName(PyOperationBase &symbol, @@ -2298,7 +2332,7 @@ void PySymbolTable::setSymbolName(PyOperationBase &symbol, mlirOperationSetAttributeByName(operation.get(), attrName, newNameAttr); } -MlirAttribute PySymbolTable::getVisibility(PyOperationBase &symbol) { +PyStringAttribute PySymbolTable::getVisibility(PyOperationBase &symbol) { PyOperation &operation = symbol.getOperation(); operation.checkValid(); MlirStringRef attrName = mlirSymbolTableGetVisibilityAttributeName(); @@ -2306,7 +2340,7 @@ MlirAttribute PySymbolTable::getVisibility(PyOperationBase &symbol) { mlirOperationGetAttributeByName(operation.get(), attrName); if (mlirAttributeIsNull(existingVisAttr)) throw nb::value_error("Expected operation to have a symbol visibility."); - return existingVisAttr; + return PyStringAttribute(symbol.getOperation().getContext(), existingVisAttr); } void PySymbolTable::setVisibility(PyOperationBase &symbol, @@ -2644,13 +2678,14 @@ class PyOpAttributeMap { PyOpAttributeMap(PyOperationRef operation) : operation(std::move(operation)) {} - MlirAttribute dunderGetItemNamed(const std::string &name) { + nb::typed + dunderGetItemNamed(const std::string &name) { MlirAttribute attr = mlirOperationGetAttributeByName(operation->get(), toMlirStringRef(name)); if (mlirAttributeIsNull(attr)) { throw nb::key_error("attempt to access a non-existent attribute"); } - return attr; + return PyAttribute(operation->getContext(), attr).maybeDownCast(); } PyNamedAttribute dunderGetItemIndexed(intptr_t index) { @@ -2915,20 +2950,13 @@ void mlir::python::populateIRCore(nb::module_ &m) { nb::arg("exc_type").none(), nb::arg("exc_value").none(), nb::arg("traceback").none()); - //---------------------------------------------------------------------------- - // Mapping of MlirContext. - // Note that this is exported as _BaseContext. The containing, Python level - // __init__.py will subclass it with site-specific functionality and set a - // "Context" attribute on this module. - //---------------------------------------------------------------------------- - // Expose DefaultThreadPool to python nb::class_(m, "ThreadPool") .def("__init__", [](PyThreadPool &self) { new (&self) PyThreadPool(); }) .def("get_max_concurrency", &PyThreadPool::getMaxConcurrency) .def("_mlir_thread_pool_ptr", &PyThreadPool::_mlir_thread_pool_ptr); - nb::class_(m, "_BaseContext") + nb::class_(m, "Context") .def("__init__", [](PyMlirContext &self) { MlirContext context = mlirContextCreateWithThreading(false); @@ -3024,7 +3052,8 @@ void mlir::python::populateIRCore(nb::module_ &m) { mlirContextAppendDialectRegistry(self.get(), registry); }, nb::arg("registry")) - .def_prop_rw("emit_error_diagnostics", nullptr, + .def_prop_rw("emit_error_diagnostics", + &PyMlirContext::getEmitErrorDiagnostics, &PyMlirContext::setEmitErrorDiagnostics, "Emit error diagnostics to diagnostic handlers. By default " "error diagnostics are captured and reported through " @@ -3042,13 +3071,16 @@ void mlir::python::populateIRCore(nb::module_ &m) { MlirStringRef ns = mlirDialectGetNamespace(self.get()); return nb::str(ns.data, ns.length); }) - .def("__repr__", [](PyDialectDescriptor &self) { - MlirStringRef ns = mlirDialectGetNamespace(self.get()); - std::string repr(""); - return repr; - }); + .def( + "__repr__", + [](PyDialectDescriptor &self) { + MlirStringRef ns = mlirDialectGetNamespace(self.get()); + std::string repr(""); + return repr; + }, + nb::sig("def __repr__(self) -> str")); //---------------------------------------------------------------------------- // Mapping of PyDialects @@ -3077,13 +3109,16 @@ void mlir::python::populateIRCore(nb::module_ &m) { .def(nb::init(), nb::arg("descriptor")) .def_prop_ro("descriptor", [](PyDialect &self) { return self.getDescriptor(); }) - .def("__repr__", [](nb::object self) { - auto clazz = self.attr("__class__"); - return nb::str(""); - }); + .def( + "__repr__", + [](const nb::object &self) { + auto clazz = self.attr("__class__"); + return nb::str(""); + }, + nb::sig("def __repr__(self) -> str")); //---------------------------------------------------------------------------- // Mapping of PyDialectRegistry @@ -3115,6 +3150,9 @@ void mlir::python::populateIRCore(nb::module_ &m) { return std::nullopt; return loc; }, + // clang-format off + nb::sig("def current(/) -> Location | None"), + // clang-format on "Gets the Location bound to the current thread or raises ValueError") .def_static( "unknown", @@ -3140,8 +3178,16 @@ void mlir::python::populateIRCore(nb::module_ &m) { nb::arg("callee"), nb::arg("frames"), nb::arg("context") = nb::none(), kContextGetCallSiteLocationDocstring) .def("is_a_callsite", mlirLocationIsACallSite) - .def_prop_ro("callee", mlirLocationCallSiteGetCallee) - .def_prop_ro("caller", mlirLocationCallSiteGetCaller) + .def_prop_ro("callee", + [](PyLocation &self) { + return PyLocation(self.getContext(), + mlirLocationCallSiteGetCallee(self)); + }) + .def_prop_ro("caller", + [](PyLocation &self) { + return PyLocation(self.getContext(), + mlirLocationCallSiteGetCaller(self)); + }) .def_static( "file", [](std::string filename, int line, int col, @@ -3166,11 +3212,13 @@ void mlir::python::populateIRCore(nb::module_ &m) { nb::arg("end_line"), nb::arg("end_col"), nb::arg("context") = nb::none(), kContextGetFileRangeDocstring) .def("is_a_file", mlirLocationIsAFileLineColRange) - .def_prop_ro("filename", - [](MlirLocation loc) { - return mlirIdentifierStr( - mlirLocationFileLineColRangeGetFilename(loc)); - }) + .def_prop_ro( + "filename", + [](MlirLocation loc) { + return mlirIdentifierStr( + mlirLocationFileLineColRangeGetFilename(loc)); + }, + nb::sig("def filename(self) -> str")) .def_prop_ro("start_line", mlirLocationFileLineColRangeGetStartLine) .def_prop_ro("start_col", mlirLocationFileLineColRangeGetStartColumn) .def_prop_ro("end_line", mlirLocationFileLineColRangeGetEndLine) @@ -3192,15 +3240,19 @@ void mlir::python::populateIRCore(nb::module_ &m) { nb::arg("locations"), nb::arg("metadata") = nb::none(), nb::arg("context") = nb::none(), kContextGetFusedLocationDocstring) .def("is_a_fused", mlirLocationIsAFused) - .def_prop_ro("locations", - [](MlirLocation loc) { - unsigned numLocations = - mlirLocationFusedGetNumLocations(loc); - std::vector locations(numLocations); - if (numLocations) - mlirLocationFusedGetLocations(loc, locations.data()); - return locations; - }) + .def_prop_ro( + "locations", + [](PyLocation &self) { + unsigned numLocations = mlirLocationFusedGetNumLocations(self); + std::vector locations(numLocations); + if (numLocations) + mlirLocationFusedGetLocations(self, locations.data()); + std::vector pyLocations{}; + pyLocations.reserve(numLocations); + for (unsigned i = 0; i < numLocations; ++i) + pyLocations.emplace_back(self.getContext(), locations[i]); + return pyLocations; + }) .def_static( "name", [](std::string name, std::optional childLoc, @@ -3215,11 +3267,17 @@ void mlir::python::populateIRCore(nb::module_ &m) { nb::arg("name"), nb::arg("childLoc") = nb::none(), nb::arg("context") = nb::none(), kContextGetNameLocationDocString) .def("is_a_name", mlirLocationIsAName) - .def_prop_ro("name_str", - [](MlirLocation loc) { - return mlirIdentifierStr(mlirLocationNameGetName(loc)); + .def_prop_ro( + "name_str", + [](MlirLocation loc) { + return mlirIdentifierStr(mlirLocationNameGetName(loc)); + }, + nb::sig("def name_str(self) -> str")) + .def_prop_ro("child_loc", + [](PyLocation &self) { + return PyLocation(self.getContext(), + mlirLocationNameGetChildLoc(self)); }) - .def_prop_ro("child_loc", mlirLocationNameGetChildLoc) .def_static( "from_attr", [](PyAttribute &attribute, DefaultingPyMlirContext context) { @@ -3234,7 +3292,10 @@ void mlir::python::populateIRCore(nb::module_ &m) { "Context that owns the Location") .def_prop_ro( "attr", - [](PyLocation &self) { return mlirLocationGetAttribute(self); }, + [](PyLocation &self) { + return PyAttribute(self.getContext(), + mlirLocationGetAttribute(self)); + }, "Get the underlying LocationAttr") .def( "emit_error", @@ -3332,11 +3393,11 @@ void mlir::python::populateIRCore(nb::module_ &m) { kDumpDocstring) .def( "__str__", - [](nb::object self) { + [](const nb::object &self) { // Defer to the operation's __str__. return self.attr("operation").attr("__str__")(); }, - kOperationStrDunderDocstring) + nb::sig("def __str__(self) -> str"), kOperationStrDunderDocstring) .def( "__eq__", [](PyModule &self, PyModule &other) { @@ -3377,13 +3438,15 @@ void mlir::python::populateIRCore(nb::module_ &m) { return concreteOperation.getContext().getObject(); }, "Context that owns the Operation") - .def_prop_ro("name", - [](PyOperationBase &self) { - auto &concreteOperation = self.getOperation(); - concreteOperation.checkValid(); - MlirOperation operation = concreteOperation.get(); - return mlirIdentifierStr(mlirOperationGetName(operation)); - }) + .def_prop_ro( + "name", + [](PyOperationBase &self) { + auto &concreteOperation = self.getOperation(); + concreteOperation.checkValid(); + MlirOperation operation = concreteOperation.get(); + return mlirIdentifierStr(mlirOperationGetName(operation)); + }, + nb::sig("def name(self) -> str")) .def_prop_ro("operands", [](PyOperationBase &self) { return PyOpOperandList(self.getOperation().getRef()); @@ -3437,6 +3500,7 @@ void mlir::python::populateIRCore(nb::module_ &m) { /*assumeVerified=*/false, /*skipRegions=*/false); }, + nb::sig("def __str__(self) -> str"), "Returns the assembly form of the operation.") .def("print", nb::overload_cast( @@ -3516,8 +3580,13 @@ void mlir::python::populateIRCore(nb::module_ &m) { }, "Reports if the operation is attached to its parent block.") .def("erase", [](PyOperationBase &self) { self.getOperation().erase(); }) - .def("walk", &PyOperationBase::walk, nb::arg("callback"), - nb::arg("walk_order") = MlirWalkPostOrder); + .def( + "walk", &PyOperationBase::walk, nb::arg("callback"), + nb::arg("walk_order") = MlirWalkPostOrder, + // clang-format off + nb::sig("def walk(self, callback: Callable[[Operation], WalkResult], walk_order: WalkOrder = " MAKE_MLIR_PYTHON_QUALNAME("ir.WalkOrder.POST_ORDER") ") -> None") + // clang-format on + ); nb::class_(m, "Operation") .def_static( @@ -3730,8 +3799,10 @@ void mlir::python::populateIRCore(nb::module_ &m) { .def( "add_argument", [](PyBlock &self, const PyType &type, const PyLocation &loc) { - return mlirBlockAddArgument(self.get(), type, loc); + return PyBlockArgument(self.getParentOperation(), + mlirBlockAddArgument(self.get(), type, loc)); }, + "type"_a, "loc"_a, "Append an argument of the specified type to the block and returns " "the newly added argument.") .def( @@ -3873,6 +3944,7 @@ void mlir::python::populateIRCore(nb::module_ &m) { throw nb::value_error("No current InsertionPoint"); return ip; }, + nb::sig("def current(/) -> InsertionPoint"), "Gets the InsertionPoint bound to the current thread or raises " "ValueError if none has been set") .def(nb::init(), nb::arg("beforeOperation"), @@ -3918,7 +3990,7 @@ void mlir::python::populateIRCore(nb::module_ &m) { context->get(), toMlirStringRef(attrSpec)); if (mlirAttributeIsNull(attr)) throw MLIRError("Unable to parse attribute", errors.take()); - return attr; + return PyAttribute(context.get()->getRef(), attr).maybeDownCast(); }, nb::arg("asm"), nb::arg("context") = nb::none(), "Parses an attribute from an assembly form. Raises an MLIRError on " @@ -3928,7 +4000,11 @@ void mlir::python::populateIRCore(nb::module_ &m) { [](PyAttribute &self) { return self.getContext().getObject(); }, "Context that owns the Attribute") .def_prop_ro("type", - [](PyAttribute &self) { return mlirAttributeGetType(self); }) + [](PyAttribute &self) { + return PyType(self.getContext(), + mlirAttributeGetType(self)) + .maybeDownCast(); + }) .def( "get_named", [](PyAttribute &self, std::string name) { @@ -3969,23 +4045,13 @@ void mlir::python::populateIRCore(nb::module_ &m) { return printAccum.join(); }) .def_prop_ro("typeid", - [](PyAttribute &self) -> MlirTypeID { + [](PyAttribute &self) { MlirTypeID mlirTypeID = mlirAttributeGetTypeID(self); assert(!mlirTypeIDIsNull(mlirTypeID) && "mlirTypeID was expected to be non-null."); - return mlirTypeID; + return PyTypeID(mlirTypeID); }) - .def(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR, [](PyAttribute &self) { - MlirTypeID mlirTypeID = mlirAttributeGetTypeID(self); - assert(!mlirTypeIDIsNull(mlirTypeID) && - "mlirTypeID was expected to be non-null."); - std::optional typeCaster = - PyGlobals::get().lookupTypeCaster(mlirTypeID, - mlirAttributeGetDialect(self)); - if (!typeCaster) - return nb::cast(self); - return typeCaster.value()(self); - }); + .def(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR, &PyAttribute::maybeDownCast); //---------------------------------------------------------------------------- // Mapping of PyNamedAttribute @@ -4010,11 +4076,12 @@ void mlir::python::populateIRCore(nb::module_ &m) { [](PyNamedAttribute &self) { return mlirIdentifierStr(self.namedAttr.name); }, + nb::sig("def name(self) -> str"), "The name of the NamedAttribute binding") .def_prop_ro( "attr", [](PyNamedAttribute &self) { return self.namedAttr.attribute; }, - nb::keep_alive<0, 1>(), + nb::keep_alive<0, 1>(), nb::sig("def attr(self) -> Attribute"), "The underlying generic attribute of the NamedAttribute binding"); //---------------------------------------------------------------------------- @@ -4035,7 +4102,7 @@ void mlir::python::populateIRCore(nb::module_ &m) { mlirTypeParseGet(context->get(), toMlirStringRef(typeSpec)); if (mlirTypeIsNull(type)) throw MLIRError("Unable to parse type", errors.take()); - return type; + return PyType(context.get()->getRef(), type).maybeDownCast(); }, nb::arg("asm"), nb::arg("context") = nb::none(), kContextParseTypeDocstring) @@ -4074,22 +4141,11 @@ void mlir::python::populateIRCore(nb::module_ &m) { printAccum.parts.append(")"); return printAccum.join(); }) - .def(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR, - [](PyType &self) { - MlirTypeID mlirTypeID = mlirTypeGetTypeID(self); - assert(!mlirTypeIDIsNull(mlirTypeID) && - "mlirTypeID was expected to be non-null."); - std::optional typeCaster = - PyGlobals::get().lookupTypeCaster(mlirTypeID, - mlirTypeGetDialect(self)); - if (!typeCaster) - return nb::cast(self); - return typeCaster.value()(self); - }) - .def_prop_ro("typeid", [](PyType &self) -> MlirTypeID { + .def(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR, &PyType::maybeDownCast) + .def_prop_ro("typeid", [](PyType &self) { MlirTypeID mlirTypeID = mlirTypeGetTypeID(self); if (!mlirTypeIDIsNull(mlirTypeID)) - return mlirTypeID; + return PyTypeID(mlirTypeID); auto origRepr = nb::cast(nb::repr(nb::cast(self))); throw nb::value_error( (origRepr + llvm::Twine(" has no typeid.")).str().c_str()); @@ -4124,7 +4180,9 @@ void mlir::python::populateIRCore(nb::module_ &m) { .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyValue::createFromCapsule) .def_prop_ro( "context", - [](PyValue &self) { return self.getParentOperation()->getContext(); }, + [](PyValue &self) { + return self.getParentOperation()->getContext().getObject(); + }, "Context in which the value lives.") .def( "dump", [](PyValue &self) { mlirValueDump(self.get()); }, @@ -4134,11 +4192,11 @@ void mlir::python::populateIRCore(nb::module_ &m) { [](PyValue &self) -> nb::object { MlirValue v = self.get(); if (mlirValueIsAOpResult(v)) { - assert( - mlirOperationEqual(self.getParentOperation()->get(), - mlirOpResultGetOwner(self.get())) && - "expected the owner of the value in Python to match that in " - "the IR"); + assert(mlirOperationEqual(self.getParentOperation()->get(), + mlirOpResultGetOwner(self.get())) && + "expected the owner of the value in Python to match " + "that in " + "the IR"); return self.getParentOperation().getObject(); } @@ -4149,7 +4207,10 @@ void mlir::python::populateIRCore(nb::module_ &m) { assert(false && "Value must be a block argument or an op result"); return nb::none(); - }) + }, + // clang-format off + nb::sig("def owner(self) -> Operation | Block | None")) + // clang-format on .def_prop_ro("uses", [](PyValue &self) { return PyOpOperandIterator( @@ -4207,7 +4268,11 @@ void mlir::python::populateIRCore(nb::module_ &m) { }, nb::arg("state"), kGetNameAsOperand) .def_prop_ro("type", - [](PyValue &self) { return mlirValueGetType(self.get()); }) + [](PyValue &self) { + return PyType(self.getParentOperation()->getContext(), + mlirValueGetType(self.get())) + .maybeDownCast(); + }) .def( "set_type", [](PyValue &self, const PyType &type) { @@ -4222,29 +4287,27 @@ void mlir::python::populateIRCore(nb::module_ &m) { kValueReplaceAllUsesWithDocstring) .def( "replace_all_uses_except", - [](MlirValue self, MlirValue with, PyOperation &exception) { + [](PyValue &self, PyValue &with, PyOperation &exception) { MlirOperation exceptedUser = exception.get(); mlirValueReplaceAllUsesExcept(self, with, 1, &exceptedUser); }, - nb::arg("with"), nb::arg("exceptions"), + nb::arg("with_"), nb::arg("exceptions"), kValueReplaceAllUsesExceptDocstring) .def( "replace_all_uses_except", - [](MlirValue self, MlirValue with, nb::list exceptions) { + [](PyValue &self, PyValue &with, + std::vector &exceptions) { // Convert Python list to a SmallVector of MlirOperations llvm::SmallVector exceptionOps; - for (nb::handle exception : exceptions) { - exceptionOps.push_back(nb::cast(exception).get()); - } - + for (PyOperation &exception : exceptions) + exceptionOps.push_back(exception); mlirValueReplaceAllUsesExcept( self, with, static_cast(exceptionOps.size()), exceptionOps.data()); }, - nb::arg("with"), nb::arg("exceptions"), + nb::arg("with_"), nb::arg("exceptions"), kValueReplaceAllUsesExceptDocstring) - .def(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR, - [](PyValue &self) { return self.maybeDownCast(); }) + .def(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR, &PyValue::maybeDownCast) .def_prop_ro( "location", [](MlirValue self) { diff --git a/mlir/lib/Bindings/Python/IRInterfaces.cpp b/mlir/lib/Bindings/Python/IRInterfaces.cpp index 6aa057c2a78aa..6c53289c5011e 100644 --- a/mlir/lib/Bindings/Python/IRInterfaces.cpp +++ b/mlir/lib/Bindings/Python/IRInterfaces.cpp @@ -223,7 +223,7 @@ class PyConcreteOpInterface { /// Returns the opview of the operation instance from which this object was /// constructed. Throws a type error if this object was constructed form a /// subclass of OpView. - nb::object getOpView() { + nb::typed getOpView() { if (operation == nullptr) { throw nb::type_error("Cannot get an opview from a static interface"); } @@ -331,6 +331,7 @@ class PyShapedTypeComponents { .def_prop_ro( "element_type", [](PyShapedTypeComponents &self) { return self.elementType; }, + nb::sig("def element_type(self) -> Type"), "Returns the element type of the shaped type components.") .def_static( "get", diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h index 28b885f136fe0..414f37cc97f2a 100644 --- a/mlir/lib/Bindings/Python/IRModule.h +++ b/mlir/lib/Bindings/Python/IRModule.h @@ -19,6 +19,7 @@ #include "NanobindUtils.h" #include "mlir-c/AffineExpr.h" #include "mlir-c/AffineMap.h" +#include "mlir-c/BuiltinAttributes.h" #include "mlir-c/Diagnostics.h" #include "mlir-c/IR.h" #include "mlir-c/IntegerSet.h" @@ -75,7 +76,7 @@ class PyObjectRef { /// Releases the object held by this instance, returning it. /// This is the proper thing to return from a function that wants to return /// the reference. Note that this does not work from initializers. - nanobind::object releaseObject() { + nanobind::typed releaseObject() { assert(referrent && object); referrent = nullptr; auto stolen = std::move(object); @@ -87,7 +88,7 @@ class PyObjectRef { assert(referrent && object); return referrent; } - nanobind::object getObject() { + nanobind::typed getObject() { assert(referrent && object); return object; } @@ -235,6 +236,7 @@ class PyMlirContext { /// Controls whether error diagnostics should be propagated to diagnostic /// handlers, instead of being captured by `ErrorCapture`. void setEmitErrorDiagnostics(bool value) { emitErrorDiagnostics = value; } + bool getEmitErrorDiagnostics() { return emitErrorDiagnostics; } struct ErrorCapture; private: @@ -269,7 +271,8 @@ class DefaultingPyMlirContext : public Defaulting { public: using Defaulting::Defaulting; - static constexpr const char kTypeDescription[] = "mlir.ir.Context"; + static constexpr const char kTypeDescription[] = + MAKE_MLIR_PYTHON_QUALNAME("ir.Context"); static PyMlirContext &resolve(); }; @@ -495,7 +498,8 @@ class DefaultingPyLocation : public Defaulting { public: using Defaulting::Defaulting; - static constexpr const char kTypeDescription[] = "mlir.ir.Location"; + static constexpr const char kTypeDescription[] = + MAKE_MLIR_PYTHON_QUALNAME("ir.Location"); static PyLocation &resolve(); operator MlirLocation() const { return *get(); } @@ -596,6 +600,7 @@ class PyOperationBase { /// drops to zero or it is attached to a parent, at which point its lifetime /// is bounded by its top-level parent reference. class PyOperation; +class PyOpView; using PyOperationRef = PyObjectRef; class PyOperation : public PyOperationBase, public BaseContextObject { public: @@ -675,7 +680,7 @@ class PyOperation : public PyOperationBase, public BaseContextObject { PyLocation &location, const nanobind::object &ip, bool inferType); /// Creates an OpView suitable for this operation. - nanobind::object createOpView(); + nanobind::typed createOpView(); /// Erases the underlying MlirOperation, removes its pointer from the /// parent context's live operations map, and sets the valid bit false. @@ -685,7 +690,7 @@ class PyOperation : public PyOperationBase, public BaseContextObject { void setInvalid() { valid = false; } /// Clones this operation. - nanobind::object clone(const nanobind::object &ip); + nanobind::typed clone(const nanobind::object &ip); PyOperation(PyMlirContextRef contextRef, MlirOperation operation); @@ -885,6 +890,8 @@ class PyType : public BaseContextObject { /// is taken by calling this function. static PyType createFromCapsule(nanobind::object capsule); + nanobind::typed maybeDownCast(); + private: MlirType type; }; @@ -958,16 +965,18 @@ class PyConcreteType : public BaseTy { }, nanobind::arg("other")); cls.def_prop_ro_static( - "static_typeid", [](nanobind::object & /*class*/) -> MlirTypeID { + "static_typeid", + [](nanobind::object & /*class*/) { if (DerivedTy::getTypeIdFunction) - return DerivedTy::getTypeIdFunction(); + return PyTypeID(DerivedTy::getTypeIdFunction()); throw nanobind::attribute_error( (DerivedTy::pyClassName + llvm::Twine(" has no typeid.")) .str() .c_str()); - }); + }, + nanobind::sig("def static_typeid(/) -> TypeID")); cls.def_prop_ro("typeid", [](PyType &self) { - return nanobind::cast(nanobind::cast(self).attr("typeid")); + return nanobind::cast(nanobind::cast(self).attr("typeid")); }); cls.def("__repr__", [](DerivedTy &self) { PyPrintAccumulator printAccum; @@ -1011,6 +1020,8 @@ class PyAttribute : public BaseContextObject { /// is taken by calling this function. static PyAttribute createFromCapsule(nanobind::object capsule); + nanobind::typed maybeDownCast(); + private: MlirAttribute attr; }; @@ -1088,19 +1099,23 @@ class PyConcreteAttribute : public BaseTy { return DerivedTy::isaFunction(otherAttr); }, nanobind::arg("other")); - cls.def_prop_ro( - "type", [](PyAttribute &attr) { return mlirAttributeGetType(attr); }); + cls.def_prop_ro("type", [](PyAttribute &attr) { + return PyType(attr.getContext(), mlirAttributeGetType(attr)) + .maybeDownCast(); + }); cls.def_prop_ro_static( - "static_typeid", [](nanobind::object & /*class*/) -> MlirTypeID { + "static_typeid", + [](nanobind::object & /*class*/) -> PyTypeID { if (DerivedTy::getTypeIdFunction) - return DerivedTy::getTypeIdFunction(); + return PyTypeID(DerivedTy::getTypeIdFunction()); throw nanobind::attribute_error( (DerivedTy::pyClassName + llvm::Twine(" has no typeid.")) .str() .c_str()); - }); + }, + nanobind::sig("def static_typeid(/) -> TypeID")); cls.def_prop_ro("typeid", [](PyAttribute &self) { - return nanobind::cast(nanobind::cast(self).attr("typeid")); + return nanobind::cast(nanobind::cast(self).attr("typeid")); }); cls.def("__repr__", [](DerivedTy &self) { PyPrintAccumulator printAccum; @@ -1128,6 +1143,17 @@ class PyConcreteAttribute : public BaseTy { static void bindDerived(ClassTy &m) {} }; +class PyStringAttribute : public PyConcreteAttribute { +public: + static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAString; + static constexpr const char *pyClassName = "StringAttr"; + using PyConcreteAttribute::PyConcreteAttribute; + static constexpr GetTypeIDFunctionTy getTypeIdFunction = + mlirStringAttrGetTypeID; + + static void bindDerived(ClassTy &c); +}; + /// Wrapper around the generic MlirValue. /// Values are managed completely by the operation that resulted in their /// definition. For op result value, this is the operation that defines the @@ -1152,7 +1178,7 @@ class PyValue { /// Gets a capsule wrapping the void* within the MlirValue. nanobind::object getCapsule(); - nanobind::object maybeDownCast(); + nanobind::typed maybeDownCast(); /// Creates a PyValue from the MlirValue wrapped by a capsule. Ownership of /// the underlying MlirValue is still tied to the owning operation. @@ -1243,7 +1269,8 @@ class PySymbolTable { /// Returns the symbol (opview) with the given name, throws if there is no /// such symbol in the table. - nanobind::object dunderGetItem(const std::string &name); + nanobind::typed + dunderGetItem(const std::string &name); /// Removes the given operation from the symbol table and erases it. void erase(PyOperationBase &symbol); @@ -1254,14 +1281,14 @@ class PySymbolTable { /// Inserts the given operation into the symbol table. The operation must have /// the symbol trait. - MlirAttribute insert(PyOperationBase &symbol); + PyStringAttribute insert(PyOperationBase &symbol); /// Gets and sets the name of a symbol op. - static MlirAttribute getSymbolName(PyOperationBase &symbol); + static PyStringAttribute getSymbolName(PyOperationBase &symbol); static void setSymbolName(PyOperationBase &symbol, const std::string &name); /// Gets and sets the visibility of a symbol op. - static MlirAttribute getVisibility(PyOperationBase &symbol); + static PyStringAttribute getVisibility(PyOperationBase &symbol); static void setVisibility(PyOperationBase &symbol, const std::string &visibility); diff --git a/mlir/lib/Bindings/Python/IRTypes.cpp b/mlir/lib/Bindings/Python/IRTypes.cpp index 745b9c48c5165..09ef64d4e0baf 100644 --- a/mlir/lib/Bindings/Python/IRTypes.cpp +++ b/mlir/lib/Bindings/Python/IRTypes.cpp @@ -501,7 +501,10 @@ class PyComplexType : public PyConcreteType { "Create a complex type"); c.def_prop_ro( "element_type", - [](PyComplexType &self) { return mlirComplexTypeGetElementType(self); }, + [](PyComplexType &self) { + return PyType(self.getContext(), mlirComplexTypeGetElementType(self)) + .maybeDownCast(); + }, "Returns element type."); } }; @@ -512,7 +515,10 @@ class PyComplexType : public PyConcreteType { void mlir::PyShapedType::bindDerived(ClassTy &c) { c.def_prop_ro( "element_type", - [](PyShapedType &self) { return mlirShapedTypeGetElementType(self); }, + [](PyShapedType &self) { + return PyType(self.getContext(), mlirShapedTypeGetElementType(self)) + .maybeDownCast(); + }, "Returns the element type of the shaped type."); c.def_prop_ro( "has_rank", @@ -718,14 +724,15 @@ class PyRankedTensorType nb::arg("shape"), nb::arg("element_type"), nb::arg("encoding") = nb::none(), nb::arg("loc") = nb::none(), "Create a ranked tensor type"); - c.def_prop_ro("encoding", - [](PyRankedTensorType &self) -> std::optional { - MlirAttribute encoding = - mlirRankedTensorTypeGetEncoding(self.get()); - if (mlirAttributeIsNull(encoding)) - return std::nullopt; - return encoding; - }); + c.def_prop_ro( + "encoding", + [](PyRankedTensorType &self) + -> std::optional> { + MlirAttribute encoding = mlirRankedTensorTypeGetEncoding(self.get()); + if (mlirAttributeIsNull(encoding)) + return std::nullopt; + return PyAttribute(self.getContext(), encoding).maybeDownCast(); + }); } }; @@ -785,8 +792,10 @@ class PyMemRefType : public PyConcreteType { nb::arg("loc") = nb::none(), "Create a memref type") .def_prop_ro( "layout", - [](PyMemRefType &self) -> MlirAttribute { - return mlirMemRefTypeGetLayout(self); + [](PyMemRefType &self) -> nb::typed { + return PyAttribute(self.getContext(), + mlirMemRefTypeGetLayout(self)) + .maybeDownCast(); }, "The layout of the MemRef type.") .def( @@ -810,11 +819,12 @@ class PyMemRefType : public PyConcreteType { "The layout of the MemRef type as an affine map.") .def_prop_ro( "memory_space", - [](PyMemRefType &self) -> std::optional { + [](PyMemRefType &self) + -> std::optional> { MlirAttribute a = mlirMemRefTypeGetMemorySpace(self); if (mlirAttributeIsNull(a)) return std::nullopt; - return a; + return PyAttribute(self.getContext(), a).maybeDownCast(); }, "Returns the memory space of the given MemRef type."); } @@ -850,11 +860,12 @@ class PyUnrankedMemRefType nb::arg("loc") = nb::none(), "Create a unranked memref type") .def_prop_ro( "memory_space", - [](PyUnrankedMemRefType &self) -> std::optional { + [](PyUnrankedMemRefType &self) + -> std::optional> { MlirAttribute a = mlirUnrankedMemrefGetMemorySpace(self); if (mlirAttributeIsNull(a)) return std::nullopt; - return a; + return PyAttribute(self.getContext(), a).maybeDownCast(); }, "Returns the memory space of the given Unranked MemRef type."); } @@ -872,9 +883,14 @@ class PyTupleType : public PyConcreteType { static void bindDerived(ClassTy &c) { c.def_static( "get_tuple", - [](std::vector elements, DefaultingPyMlirContext context) { + [](const std::vector &elements, + DefaultingPyMlirContext context) { + std::vector mlirElements; + mlirElements.reserve(elements.size()); + for (const auto &element : elements) + mlirElements.push_back(element.get()); MlirType t = mlirTupleTypeGet(context->get(), elements.size(), - elements.data()); + mlirElements.data()); return PyTupleType(context->getRef(), t); }, nb::arg("elements"), nb::arg("context") = nb::none(), @@ -882,7 +898,8 @@ class PyTupleType : public PyConcreteType { c.def( "get_type", [](PyTupleType &self, intptr_t pos) { - return mlirTupleTypeGetType(self, pos); + return PyType(self.getContext(), mlirTupleTypeGetType(self, pos)) + .maybeDownCast(); }, nb::arg("pos"), "Returns the pos-th type in the tuple type."); c.def_prop_ro( @@ -906,11 +923,20 @@ class PyFunctionType : public PyConcreteType { static void bindDerived(ClassTy &c) { c.def_static( "get", - [](std::vector inputs, std::vector results, + [](std::vector inputs, std::vector results, DefaultingPyMlirContext context) { - MlirType t = - mlirFunctionTypeGet(context->get(), inputs.size(), inputs.data(), - results.size(), results.data()); + std::vector mlirInputs; + mlirInputs.reserve(inputs.size()); + for (const auto &input : inputs) + mlirInputs.push_back(input.get()); + std::vector mlirResults; + mlirResults.reserve(results.size()); + for (const auto &result : results) + mlirResults.push_back(result.get()); + + MlirType t = mlirFunctionTypeGet(context->get(), inputs.size(), + mlirInputs.data(), results.size(), + mlirResults.data()); return PyFunctionType(context->getRef(), t); }, nb::arg("inputs"), nb::arg("results"), nb::arg("context") = nb::none(), diff --git a/mlir/lib/Bindings/Python/MainModule.cpp b/mlir/lib/Bindings/Python/MainModule.cpp index a14f09f77d2c3..52656138843b9 100644 --- a/mlir/lib/Bindings/Python/MainModule.cpp +++ b/mlir/lib/Bindings/Python/MainModule.cpp @@ -115,6 +115,9 @@ NB_MODULE(_mlir, m) { }); }, "typeid"_a, nb::kw_only(), "replace"_a = false, + // clang-format off + nb::sig("def register_type_caster(typeid: " MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID") ", *, replace: bool = False) -> object"), + // clang-format on "Register a type caster for casting MLIR types to custom user types."); m.def( MLIR_PYTHON_CAPI_VALUE_CASTER_REGISTER_ATTR, @@ -127,6 +130,9 @@ NB_MODULE(_mlir, m) { }); }, "typeid"_a, nb::kw_only(), "replace"_a = false, + // clang-format off + nb::sig("def register_value_caster(typeid: " MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID") ", *, replace: bool = False) -> object"), + // clang-format on "Register a value caster for casting MLIR values to custom user values."); // Define and populate IR submodule. diff --git a/mlir/lib/Bindings/Python/NanobindUtils.h b/mlir/lib/Bindings/Python/NanobindUtils.h index 64ea4329f65f1..40b3215f6f5fe 100644 --- a/mlir/lib/Bindings/Python/NanobindUtils.h +++ b/mlir/lib/Bindings/Python/NanobindUtils.h @@ -276,7 +276,7 @@ class Sliceable { /// Returns the element at the given slice index. Supports negative indices /// by taking elements in inverse order. Returns a nullptr object if out /// of bounds. - nanobind::object getItem(intptr_t index) { + nanobind::typed getItem(intptr_t index) { // Negative indices mean we count from the end. index = wrapIndex(index); if (index < 0) { diff --git a/mlir/lib/Bindings/Python/Pass.cpp b/mlir/lib/Bindings/Python/Pass.cpp index 47ef5d8e9dd3b..e489585fd5f50 100644 --- a/mlir/lib/Bindings/Python/Pass.cpp +++ b/mlir/lib/Bindings/Python/Pass.cpp @@ -77,6 +77,9 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) { new (&self) PyPassManager(passManager); }, "anchor_op"_a = nb::str("any"), "context"_a = nb::none(), + // clang-format off + nb::sig("def __init__(self, anchor_op: str = 'any', context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " | None = None) -> None"), + // clang-format on "Create a new PassManager for the current (or provided) Context.") .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyPassManager::getCapsule) .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyPassManager::createFromCapsule) @@ -149,6 +152,9 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) { return new PyPassManager(passManager); }, "pipeline"_a, "context"_a = nb::none(), + // clang-format off + nb::sig("def parse(pipeline: str, context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " | None = None) -> PassManager"), + // clang-format on "Parse a textual pass-pipeline and return a top-level PassManager " "that can be applied on a Module. Throw a ValueError if the pipeline " "can't be parsed") @@ -217,6 +223,9 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) { errors.take()); }, "operation"_a, + // clang-format off + nb::sig("def run(self, operation: " MAKE_MLIR_PYTHON_QUALNAME("ir._OperationBase") ") -> None"), + // clang-format on "Run the pass manager on the provided operation, raising an " "MLIRError on failure.") .def( diff --git a/mlir/lib/Bindings/Python/Rewrite.cpp b/mlir/lib/Bindings/Python/Rewrite.cpp index 5b7de50f02e6a..96f00cface64f 100644 --- a/mlir/lib/Bindings/Python/Rewrite.cpp +++ b/mlir/lib/Bindings/Python/Rewrite.cpp @@ -10,8 +10,10 @@ #include "IRModule.h" #include "mlir-c/Rewrite.h" +// clang-format off #include "mlir/Bindings/Python/Nanobind.h" #include "mlir-c/Bindings/Python/Interop.h" // This is expected after nanobind. +// clang-format on #include "mlir/Config/mlir-config.h" namespace nb = nanobind; @@ -83,10 +85,13 @@ void mlir::python::populateRewriteSubmodule(nb::module_ &m) { nb::class_(m, "PDLModule") .def( "__init__", - [](PyPDLPatternModule &self, MlirModule module) { - new (&self) - PyPDLPatternModule(mlirPDLPatternModuleFromModule(module)); + [](PyPDLPatternModule &self, PyModule &module) { + new (&self) PyPDLPatternModule( + mlirPDLPatternModuleFromModule(module.get())); }, + // clang-format off + nb::sig("def __init__(self, module: " MAKE_MLIR_PYTHON_QUALNAME("ir.Module") ") -> None"), + // clang-format on "module"_a, "Create a PDL module from the given module.") .def("freeze", [](PyPDLPatternModule &self) { return new PyFrozenRewritePatternSet(mlirFreezeRewritePattern( @@ -100,24 +105,31 @@ void mlir::python::populateRewriteSubmodule(nb::module_ &m) { &PyFrozenRewritePatternSet::createFromCapsule); m.def( "apply_patterns_and_fold_greedily", - [](PyModule &module, MlirFrozenRewritePatternSet set) { - auto status = mlirApplyPatternsAndFoldGreedily(module.get(), set, {}); + [](PyModule &module, PyFrozenRewritePatternSet &set) { + auto status = + mlirApplyPatternsAndFoldGreedily(module.get(), set.get(), {}); if (mlirLogicalResultIsFailure(status)) throw std::runtime_error("pattern application failed to converge"); }, "module"_a, "set"_a, + // clang-format off + nb::sig("def apply_patterns_and_fold_greedily(module: " MAKE_MLIR_PYTHON_QUALNAME("ir.Module") ", set: FrozenRewritePatternSet) -> None"), + // clang-format on "Applys the given patterns to the given module greedily while folding " "results.") .def( "apply_patterns_and_fold_greedily", - [](PyOperationBase &op, MlirFrozenRewritePatternSet set) { + [](PyOperationBase &op, PyFrozenRewritePatternSet &set) { auto status = mlirApplyPatternsAndFoldGreedilyWithOp( - op.getOperation(), set, {}); + op.getOperation(), set.get(), {}); if (mlirLogicalResultIsFailure(status)) throw std::runtime_error( "pattern application failed to converge"); }, "op"_a, "set"_a, + // clang-format off + nb::sig("def apply_patterns_and_fold_greedily(op: " MAKE_MLIR_PYTHON_QUALNAME("ir._OperationBase") ", set: FrozenRewritePatternSet) -> None"), + // clang-format on "Applys the given patterns to the given op greedily while folding " "results."); } diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt index 7b2e1b8c36f25..97f0778071ef9 100644 --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -1,5 +1,9 @@ include(AddMLIRPython) +# Specifies that all MLIR packages are co-located under the `MLIR_PYTHON_PACKAGE_PREFIX.` +# top level package (the API has been embedded in a relocatable way). +add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=${MLIR_PYTHON_PACKAGE_PREFIX}.") + ################################################################################ # Structural groupings. ################################################################################ @@ -19,15 +23,11 @@ declare_mlir_python_sources(MLIRPythonSources.Core.Python ADD_TO_PARENT MLIRPythonSources.Core SOURCES _mlir_libs/__init__.py + _mlir_libs/_mlir/py.typed ir.py passmanager.py rewrite.py dialects/_ods_common.py - - # The main _mlir module has submodules: include stubs from each. - _mlir_libs/_mlir/__init__.pyi - _mlir_libs/_mlir/ir.pyi - _mlir_libs/_mlir/passmanager.pyi ) declare_mlir_python_sources(MLIRPythonSources.Core.Python.Extras @@ -686,7 +686,7 @@ if(MLIR_ENABLE_EXECUTION_ENGINE) MODULE_NAME _mlirExecutionEngine ADD_TO_PARENT MLIRPythonSources.ExecutionEngine ROOT_DIR "${PYTHON_SOURCE_DIR}" - PYTHON_BINDINGS_LIBRARY nanobind + PYTHON_BINDINGS_LIBRARY nanobind SOURCES ExecutionEngineModule.cpp PRIVATE_LINK_LIBS @@ -837,10 +837,11 @@ endif() # once ready. ################################################################################ +set(MLIRPythonModules_ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}") add_mlir_python_common_capi_library(MLIRPythonCAPI INSTALL_COMPONENT MLIRPythonModules INSTALL_DESTINATION "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs" - OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs" + OUTPUT_DIRECTORY "${MLIRPythonModules_ROOT_PREFIX}/_mlir_libs" RELATIVE_INSTALL_ROOT "../../../.." DECLARED_HEADERS MLIRPythonCAPI.HeaderSources @@ -863,18 +864,104 @@ if(NOT LLVM_ENABLE_IDE) ) endif() +# _mlir stubgen +# Note: All this needs to come before add_mlir_python_modules(MLIRPythonModules so that the install targets for the +# generated type stubs get created. + +set(_core_type_stub_sources + _mlir/__init__.pyi + _mlir/ir.pyi + _mlir/passmanager.pyi + _mlir/rewrite.pyi +) + +# Note 1: INTERFACE_SOURCES is a genex ($ $) +# which will be evaluated by file(GENERATE ...) inside mlir_generate_type_stubs. This will evaluate to the correct +# thing in the build dir (i.e., actual source dir paths) and in the install dir +# (where it's a conventional path; see install/lib/cmake/mlir/MLIRTargets.cmake). +# +# Note 2: MLIRPythonExtension.Core is the target that is defined using target_sources(INTERFACE) +# **NOT** MLIRPythonModules.extension._mlir.dso. So be sure to use the correct target! +get_target_property(_core_extension_srcs MLIRPythonExtension.Core INTERFACE_SOURCES) + +# Why is MODULE_NAME _mlir here but mlir._mlir_libs._mlirPythonTestNanobind below??? +# The _mlir extension can be imported independently of any other python code and/or extension modules. +# I.e., you could do `cd $MLIRPythonModules_ROOT_PREFIX/_mlir_libs && python -c "import _mlir"` (try it!). +# _mlir is also (currently) the only extension for which this is possible because dialect extensions modules, +# which generally make use of `mlir_value_subclass/mlir_type_subclass/mlir_attribute_subclass`, perform an +# `import mlir` right when they're loaded (see the mlir_*_subclass ctors in NanobindAdaptors.h). +# Note, this also why IMPORT_PATHS "${MLIRPythonModules_ROOT_PREFIX}/_mlir_libs" here while below +# "${MLIRPythonModules_ROOT_PREFIX}/.." (because MLIR_BINDINGS_PYTHON_INSTALL_PREFIX, by default, ends at mlir). +# +# Further note: this function creates file targets like +# "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs/_mlir/__init__.pyi". These must match the file targets +# that declare_mlir_python_sources expects, which are like "${ROOT_DIR}/${WHATEVER_SOURCE}". +# This is why _mlir_libs is prepended below. +mlir_generate_type_stubs( + MODULE_NAME _mlir + DEPENDS_TARGETS MLIRPythonModules.extension._mlir.dso + OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs" + OUTPUTS "${_core_type_stub_sources}" + DEPENDS_TARGET_SRC_DEPS "${_core_extension_srcs}" + IMPORT_PATHS "${MLIRPythonModules_ROOT_PREFIX}/_mlir_libs" +) +set(_mlir_typestub_gen_target "${NB_STUBGEN_CUSTOM_TARGET}") + +list(TRANSFORM _core_type_stub_sources PREPEND "_mlir_libs/") +# Note, we do not do ADD_TO_PARENT here so that the type stubs are not associated (as mlir_DEPENDS) with +# MLIRPythonSources.Core (or something) when a distro is installed/created. Otherwise they would not be regenerated +# by users of the distro (the stubs are still installed in the distro - they are just not added to mlir_DEPENDS). +declare_mlir_python_sources( + MLIRPythonExtension.Core.type_stub_gen + ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs" + SOURCES "${_core_type_stub_sources}" +) + +# _mlirPythonTestNanobind stubgen + +if(MLIR_INCLUDE_TESTS) + get_target_property(_test_extension_srcs MLIRPythonTestSources.PythonTestExtensionNanobind INTERFACE_SOURCES) + mlir_generate_type_stubs( + # This is the FQN path because dialect modules import _mlir when loaded. See above. + MODULE_NAME mlir._mlir_libs._mlirPythonTestNanobind + DEPENDS_TARGETS + # You need both _mlir and _mlirPythonTestNanobind because dialect modules import _mlir when loaded + # (so _mlir needs to be built before calling stubgen). + MLIRPythonModules.extension._mlir.dso + MLIRPythonModules.extension._mlirPythonTestNanobind.dso + # You need this one so that ir.py "built" because mlir._mlir_libs.__init__.py import mlir.ir in _site_initialize. + MLIRPythonModules.sources.MLIRPythonSources.Core.Python + OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs" + OUTPUTS _mlirPythonTestNanobind.pyi + DEPENDS_TARGET_SRC_DEPS "${_test_extension_srcs}" + IMPORT_PATHS "${MLIRPythonModules_ROOT_PREFIX}/.." + ) + set(_mlirPythonTestNanobind_typestub_gen_target "${NB_STUBGEN_CUSTOM_TARGET}") + declare_mlir_python_sources( + MLIRPythonTestSources.PythonTestExtensionNanobind.type_stub_gen + ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs" + ADD_TO_PARENT MLIRPythonTestSources.Dialects + SOURCES _mlir_libs/_mlirPythonTestNanobind.pyi + ) +endif() + ################################################################################ # The fully assembled package of modules. # This must come last. ################################################################################ add_mlir_python_modules(MLIRPythonModules - ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir" + ROOT_PREFIX ${MLIRPythonModules_ROOT_PREFIX} INSTALL_PREFIX "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}" DECLARED_SOURCES MLIRPythonSources MLIRPythonExtension.RegisterEverything + MLIRPythonExtension.Core.type_stub_gen ${_ADDL_TEST_SOURCES} COMMON_CAPI_LINK_LIBS MLIRPythonCAPI ) +add_dependencies(MLIRPythonModules "${_mlir_typestub_gen_target}") +if(MLIR_INCLUDE_TESTS) + add_dependencies(MLIRPythonModules "${_mlirPythonTestNanobind_typestub_gen_target}") +endif() diff --git a/mlir/python/mlir/_mlir_libs/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py index 083a9075fe4c5..63244212ba42c 100644 --- a/mlir/python/mlir/_mlir_libs/__init__.py +++ b/mlir/python/mlir/_mlir_libs/__init__.py @@ -147,7 +147,9 @@ def process_initializer_module(module_name): if not process_initializer_module(module_name): break - class Context(ir._BaseContext): + ir._Context = ir.Context + + class Context(ir._Context): def __init__( self, load_on_create_dialects=None, thread_pool=None, *args, **kwargs ): diff --git a/mlir/python/mlir/_mlir_libs/_mlir/__init__.pyi b/mlir/python/mlir/_mlir_libs/_mlir/__init__.pyi deleted file mode 100644 index 03449b70b7fa3..0000000000000 --- a/mlir/python/mlir/_mlir_libs/_mlir/__init__.pyi +++ /dev/null @@ -1,12 +0,0 @@ - -globals: "_Globals" - -class _Globals: - dialect_search_modules: list[str] - def _register_dialect_impl(self, dialect_namespace: str, dialect_class: type) -> None: ... - def _register_operation_impl(self, operation_name: str, operation_class: type) -> None: ... - def append_dialect_search_prefix(self, module_name: str) -> None: ... - def _check_dialect_module_loaded(self, dialect_namespace: str) -> bool: ... - -def register_dialect(dialect_class: type) -> type: ... -def register_operation(dialect_class: type, *, replace: bool = ...) -> type: ... diff --git a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi deleted file mode 100644 index dcae3dd742940..0000000000000 --- a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi +++ /dev/null @@ -1,2846 +0,0 @@ -# Originally imported via: -# pybind11-stubgen --print-invalid-expressions-as-is mlir._mlir_libs._mlir.ir -# but with the following diff (in order to remove pipes from types, -# which we won't support until bumping minimum python to 3.10) -# -# --------------------- diff begins ------------------------------------ -# -# diff --git a/pybind11_stubgen/printer.py b/pybind11_stubgen/printer.py -# index 1f755aa..4924927 100644 -# --- a/pybind11_stubgen/printer.py -# +++ b/pybind11_stubgen/printer.py -# @@ -283,14 +283,6 @@ class Printer: -# return split[0] + "..." -# -# def print_type(self, type_: ResolvedType) -> str: -# - if ( -# - str(type_.name) == "typing.Optional" -# - and type_.parameters is not None -# - and len(type_.parameters) == 1 -# - ): -# - return f"{self.print_annotation(type_.parameters[0])} | None" -# - if str(type_.name) == "typing.Union" and type_.parameters is not None: -# - return " | ".join(self.print_annotation(p) for p in type_.parameters) -# if type_.parameters: -# param_str = ( -# "[" -# -# --------------------- diff ends ------------------------------------ -# -# Local modifications: -# * Rewrite references to 'mlir.ir' to local types. -# * Drop `typing.` everywhere (top-level import instead). -# * List -> List, dict -> Dict, Tuple -> Tuple. -# * copy-paste Buffer type from typing_extensions. -# * Shuffle _OperationBase, AffineExpr, Attribute, Type, Value to the top. -# * Patch raw C++ types (like "PyAsmState") with a regex like `Py(.*)`. -# * _BaseContext -> Context, MlirType -> Type, MlirTypeID -> TypeID, MlirAttribute -> Attribute. -# * Local edits to signatures and types that pybind11-stubgen did not auto detect (or detected incorrectly). -# * Add MLIRError, _GlobalDebug, _OperationBase to __all__ by hand. -# * Fill in `Any`s where possible. -# * black formatting. - -from __future__ import annotations - -import abc -import collections -from collections.abc import Callable, Sequence -from pathlib import Path -from typing import Any, BinaryIO, ClassVar, Literal, TypeVar, overload - -__all__ = [ - "AffineAddExpr", - "AffineBinaryExpr", - "AffineCeilDivExpr", - "AffineConstantExpr", - "AffineDimExpr", - "AffineExpr", - "AffineExprList", - "AffineFloorDivExpr", - "AffineMap", - "AffineMapAttr", - "AffineModExpr", - "AffineMulExpr", - "AffineSymbolExpr", - "ArrayAttr", - "ArrayAttributeIterator", - "AsmState", - "AttrBuilder", - "Attribute", - "BF16Type", - "Block", - "BlockArgument", - "BlockArgumentList", - "BlockIterator", - "BlockList", - "BoolAttr", - "ComplexType", - "Context", - "DenseBoolArrayAttr", - "DenseBoolArrayIterator", - "DenseElementsAttr", - "DenseF32ArrayAttr", - "DenseF32ArrayIterator", - "DenseF64ArrayAttr", - "DenseF64ArrayIterator", - "DenseFPElementsAttr", - "DenseI16ArrayAttr", - "DenseI16ArrayIterator", - "DenseI32ArrayAttr", - "DenseI32ArrayIterator", - "DenseI64ArrayAttr", - "DenseI64ArrayIterator", - "DenseI8ArrayAttr", - "DenseI8ArrayIterator", - "DenseIntElementsAttr", - "DenseResourceElementsAttr", - "Diagnostic", - "DiagnosticHandler", - "DiagnosticInfo", - "DiagnosticSeverity", - "Dialect", - "DialectDescriptor", - "DialectRegistry", - "Dialects", - "DictAttr", - "F16Type", - "F32Type", - "F64Type", - "FlatSymbolRefAttr", - "Float4E2M1FNType", - "Float6E2M3FNType", - "Float6E3M2FNType", - "Float8E3M4Type", - "Float8E4M3B11FNUZType", - "Float8E4M3FNType", - "Float8E4M3FNUZType", - "Float8E4M3Type", - "Float8E5M2FNUZType", - "Float8E5M2Type", - "Float8E8M0FNUType", - "FloatAttr", - "FloatTF32Type", - "FloatType", - "FunctionType", - "IndexType", - "InferShapedTypeOpInterface", - "InferTypeOpInterface", - "InsertionPoint", - "IntegerAttr", - "IntegerSet", - "IntegerSetAttr", - "IntegerSetConstraint", - "IntegerSetConstraintList", - "IntegerType", - "Location", - "MemRefType", - "Module", - "NamedAttribute", - "NoneType", - "OpAttributeMap", - "OpOperand", - "OpOperandIterator", - "OpOperandList", - "OpResult", - "OpResultList", - "OpSuccessors", - "OpView", - "OpaqueAttr", - "OpaqueType", - "Operation", - "OperationIterator", - "OperationList", - "RankedTensorType", - "Region", - "RegionIterator", - "RegionSequence", - "ShapedType", - "ShapedTypeComponents", - "StridedLayoutAttr", - "StringAttr", - "SymbolRefAttr", - "SymbolTable", - "TupleType", - "Type", - "TypeAttr", - "TypeID", - "UnitAttr", - "UnrankedMemRefType", - "UnrankedTensorType", - "Value", - "VectorType", - "_GlobalDebug", - "_OperationBase", -] - -if hasattr(collections.abc, "Buffer"): - Buffer = collections.abc.Buffer -else: - class Buffer(abc.ABC): - pass - -class _OperationBase: - @overload - def __eq__(self, arg0: _OperationBase) -> bool: ... - @overload - def __eq__(self, arg0: _OperationBase) -> bool: ... - def __hash__(self) -> int: ... - def __str__(self) -> str: - """ - Returns the assembly form of the operation. - """ - def clone(self, ip: InsertionPoint = None) -> OpView: ... - def detach_from_parent(self) -> OpView: - """ - Detaches the operation from its parent block. - """ - - @property - def attached(self) -> bool: - """ - Reports if the operation is attached to its parent block. - """ - - def erase(self) -> None: ... - - @overload - def get_asm( - binary: Literal[True], - large_elements_limit: int | None = None, - large_resource_limit: int | None = None, - enable_debug_info: bool = False, - pretty_debug_info: bool = False, - print_generic_op_form: bool = False, - use_local_scope: bool = False, - assume_verified: bool = False, - skip_regions: bool = False, - ) -> bytes: ... - @overload - def get_asm( - self, - binary: bool = False, - large_elements_limit: int | None = None, - large_resource_limit: int | None = None, - enable_debug_info: bool = False, - pretty_debug_info: bool = False, - print_generic_op_form: bool = False, - use_local_scope: bool = False, - assume_verified: bool = False, - skip_regions: bool = False, - ) -> str: - """ - Returns the assembly form of the operation. - - See the print() method for common keyword arguments for configuring - the output. - """ - - def move_after(self, other: _OperationBase) -> None: - """ - Puts self immediately after the other operation in its parent block. - """ - def move_before(self, other: _OperationBase) -> None: - """ - Puts self immediately before the other operation in its parent block. - """ - @overload - def print( - self, - state: AsmState, - file: Any | None = None, - binary: bool = False, - ) -> None: - """ - Prints the assembly form of the operation to a file like object. - - Args: - file: The file like object to write to. Defaults to sys.stdout. - binary: Whether to write bytes (True) or str (False). Defaults to False. - state: AsmState capturing the operation numbering and flags. - """ - @overload - def print( - self, - large_elements_limit: int | None = None, - large_resource_limit: int | None = None, - enable_debug_info: bool = False, - pretty_debug_info: bool = False, - print_generic_op_form: bool = False, - use_local_scope: bool = False, - assume_verified: bool = False, - file: Any | None = None, - binary: bool = False, - skip_regions: bool = False, - ) -> None: - """ - Prints the assembly form of the operation to a file like object. - - Args: - file: The file like object to write to. Defaults to sys.stdout. - binary: Whether to write bytes (True) or str (False). Defaults to False. - large_elements_limit: Whether to elide elements attributes above this - number of elements. Defaults to None (no limit). - large_resource_limit: Whether to elide resource strings above this - number of characters. Defaults to None (no limit). If large_elements_limit - is set and this is None, the behavior will be to use large_elements_limit - as large_resource_limit. - enable_debug_info: Whether to print debug/location information. Defaults - to False. - pretty_debug_info: Whether to format debug information for easier reading - by a human (warning: the result is unparseable). - print_generic_op_form: Whether to print the generic assembly forms of all - ops. Defaults to False. - use_local_Scope: Whether to print in a way that is more optimized for - multi-threaded access but may not be consistent with how the overall - module prints. - assume_verified: By default, if not printing generic form, the verifier - will be run and if it fails, generic form will be printed with a comment - about failed verification. While a reasonable default for interactive use, - for systematic use, it is often better for the caller to verify explicitly - and report failures in a more robust fashion. Set this to True if doing this - in order to avoid running a redundant verification. If the IR is actually - invalid, behavior is undefined. - skip_regions: Whether to skip printing regions. Defaults to False. - """ - def verify(self) -> bool: - """ - Verify the operation. Raises MLIRError if verification fails, and returns true otherwise. - """ - def write_bytecode(self, file: BinaryIO | str, desired_version: int | None = None) -> None: - """ - Write the bytecode form of the operation to a file like object. - - Args: - file: The file like object or path to write to. - desired_version: The version of bytecode to emit. - Returns: - The bytecode writer status. - """ - @property - def _CAPIPtr(self) -> object: ... - @property - def attributes(self) -> OpAttributeMap: ... - @property - def context(self) -> Context: - """ - Context that owns the Operation - """ - @property - def location(self) -> Location: - """ - Returns the source location the operation was defined or derived from. - """ - @property - def name(self) -> str: ... - @property - def operands(self) -> OpOperandList: ... - @property - def parent(self) -> _OperationBase | None: ... - @property - def regions(self) -> RegionSequence: ... - @property - def result(self) -> OpResult: - """ - Shortcut to get an op result if it has only one (throws an error otherwise). - """ - @property - def results(self) -> OpResultList: - """ - Returns the List of Operation results. - """ - -_TOperation = TypeVar("_TOperation", bound=_OperationBase) - -class AffineExpr: - @staticmethod - @overload - def get_add(arg0: AffineExpr, arg1: AffineExpr) -> AffineAddExpr: - """ - Gets an affine expression containing a sum of two expressions. - """ - @staticmethod - @overload - def get_add(arg0: int, arg1: AffineExpr) -> AffineAddExpr: - """ - Gets an affine expression containing a sum of a constant and another expression. - """ - @staticmethod - @overload - def get_add(arg0: AffineExpr, arg1: int) -> AffineAddExpr: - """ - Gets an affine expression containing a sum of an expression and a constant. - """ - @staticmethod - @overload - def get_ceil_div(arg0: AffineExpr, arg1: AffineExpr) -> AffineCeilDivExpr: - """ - Gets an affine expression containing the rounded-up result of dividing one expression by another. - """ - @staticmethod - @overload - def get_ceil_div(arg0: int, arg1: AffineExpr) -> AffineCeilDivExpr: - """ - Gets a semi-affine expression containing the rounded-up result of dividing a constant by an expression. - """ - @staticmethod - @overload - def get_ceil_div(arg0: AffineExpr, arg1: int) -> AffineCeilDivExpr: - """ - Gets an affine expression containing the rounded-up result of dividing an expression by a constant. - """ - @staticmethod - def get_constant( - value: int, context: Context | None = None - ) -> AffineConstantExpr: - """ - Gets a constant affine expression with the given value. - """ - @staticmethod - def get_dim(position: int, context: Context | None = None) -> AffineDimExpr: - """ - Gets an affine expression of a dimension at the given position. - """ - @staticmethod - @overload - def get_floor_div(arg0: AffineExpr, arg1: AffineExpr) -> AffineFloorDivExpr: - """ - Gets an affine expression containing the rounded-down result of dividing one expression by another. - """ - @staticmethod - @overload - def get_floor_div(arg0: int, arg1: AffineExpr) -> AffineFloorDivExpr: - """ - Gets a semi-affine expression containing the rounded-down result of dividing a constant by an expression. - """ - @staticmethod - @overload - def get_floor_div(arg0: AffineExpr, arg1: int) -> AffineFloorDivExpr: - """ - Gets an affine expression containing the rounded-down result of dividing an expression by a constant. - """ - @staticmethod - @overload - def get_mod(arg0: AffineExpr, arg1: AffineExpr) -> AffineModExpr: - """ - Gets an affine expression containing the modulo of dividing one expression by another. - """ - @staticmethod - @overload - def get_mod(arg0: int, arg1: AffineExpr) -> AffineModExpr: - """ - Gets a semi-affine expression containing the modulo of dividing a constant by an expression. - """ - @staticmethod - @overload - def get_mod(arg0: AffineExpr, arg1: int) -> AffineModExpr: - """ - Gets an affine expression containing the module of dividingan expression by a constant. - """ - @staticmethod - @overload - def get_mul(arg0: AffineExpr, arg1: AffineExpr) -> AffineMulExpr: - """ - Gets an affine expression containing a product of two expressions. - """ - @staticmethod - @overload - def get_mul(arg0: int, arg1: AffineExpr) -> AffineMulExpr: - """ - Gets an affine expression containing a product of a constant and another expression. - """ - @staticmethod - @overload - def get_mul(arg0: AffineExpr, arg1: int) -> AffineMulExpr: - """ - Gets an affine expression containing a product of an expression and a constant. - """ - @staticmethod - def get_symbol( - position: int, context: Context | None = None - ) -> AffineSymbolExpr: - """ - Gets an affine expression of a symbol at the given position. - """ - def _CAPICreate(self) -> AffineExpr: ... - @overload - def __add__(self, arg0: AffineExpr) -> AffineAddExpr: ... - @overload - def __add__(self, arg0: int) -> AffineAddExpr: ... - @overload - def __eq__(self, arg0: AffineExpr) -> bool: ... - @overload - def __eq__(self, arg0: Any) -> bool: ... - def __hash__(self) -> int: ... - @overload - def __mod__(self, arg0: AffineExpr) -> AffineModExpr: ... - @overload - def __mod__(self, arg0: int) -> AffineModExpr: ... - @overload - def __mul__(self, arg0: AffineExpr) -> AffineMulExpr: ... - @overload - def __mul__(self, arg0: int) -> AffineMulExpr: ... - def __radd__(self, arg0: int) -> AffineAddExpr: ... - def __rmod__(self, arg0: int) -> AffineModExpr: ... - def __rmul__(self, arg0: int) -> AffineMulExpr: ... - def __rsub__(self, arg0: int) -> AffineAddExpr: ... - @overload - def __sub__(self, arg0: AffineExpr) -> AffineAddExpr: ... - @overload - def __sub__(self, arg0: int) -> AffineAddExpr: ... - def compose(self, arg0: AffineMap) -> AffineExpr: ... - def dump(self) -> None: - """ - Dumps a debug representation of the object to stderr. - """ - @property - def _CAPIPtr(self) -> object: ... - @property - def context(self) -> Context: ... - -class Attribute: - @staticmethod - def parse(asm: str | bytes, context: Context | None = None) -> Attribute: - """ - Parses an attribute from an assembly form. Raises an MLIRError on failure. - """ - def _CAPICreate(self) -> Attribute: ... - @overload - def __eq__(self, arg0: Attribute) -> bool: ... - @overload - def __eq__(self, arg0: object) -> bool: ... - def __hash__(self) -> int: ... - def __init__(self, cast_from_type: Attribute) -> None: - """ - Casts the passed attribute to the generic Attribute - """ - def __str__(self) -> str: - """ - Returns the assembly form of the Attribute. - """ - def dump(self) -> None: - """ - Dumps a debug representation of the object to stderr. - """ - def get_named(self, arg0: str) -> NamedAttribute: - """ - Binds a name to the attribute - """ - def maybe_downcast(self) -> Any: ... - @property - def _CAPIPtr(self) -> object: ... - @property - def context(self) -> Context: - """ - Context that owns the Attribute - """ - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class Type: - @staticmethod - def parse(asm: str | bytes, context: Context | None = None) -> Type: - """ - Parses the assembly form of a type. - - Returns a Type object or raises an MLIRError if the type cannot be parsed. - - See also: https://mlir.llvm.org/docs/LangRef/#type-system - """ - def _CAPICreate(self) -> Type: ... - @overload - def __eq__(self, arg0: Type) -> bool: ... - @overload - def __eq__(self, arg0: object) -> bool: ... - def __hash__(self) -> int: ... - def __init__(self, cast_from_type: Type) -> None: - """ - Casts the passed type to the generic Type - """ - def __str__(self) -> str: - """ - Returns the assembly form of the type. - """ - def dump(self) -> None: - """ - Dumps a debug representation of the object to stderr. - """ - def maybe_downcast(self) -> Any: ... - @property - def _CAPIPtr(self) -> object: ... - @property - def context(self) -> Context: - """ - Context that owns the Type - """ - @property - def typeid(self) -> TypeID: ... - -class Value: - def _CAPICreate(self) -> Value: ... - @overload - def __eq__(self, arg0: Value) -> bool: ... - @overload - def __eq__(self, arg0: object) -> bool: ... - def __hash__(self) -> int: ... - def __init__(self, value: Value) -> None: ... - def __str__(self) -> str: - """ - Returns the string form of the value. - - If the value is a block argument, this is the assembly form of its type and the - position in the argument List. If the value is an operation result, this is - equivalent to printing the operation that produced it. - """ - def dump(self) -> None: - """ - Dumps a debug representation of the object to stderr. - """ - @overload - def get_name(self, use_local_scope: bool = False, use_name_loc_as_prefix: bool = True) -> str: ... - @overload - def get_name(self, state: AsmState) -> str: - """ - Returns the string form of value as an operand (i.e., the ValueID). - """ - def maybe_downcast(self) -> Any: ... - def replace_all_uses_with(self, arg0: Value) -> None: - """ - Replace all uses of value with the new value, updating anything in - the IR that uses 'self' to use the other value instead. - """ - def set_type(self, type: Type) -> None: ... - @property - def _CAPIPtr(self) -> object: ... - @property - def context(self) -> Context: - """ - Context in which the value lives. - """ - @property - def owner(self) -> _OperationBase: ... - @property - def type(self) -> Type: ... - @property - def uses(self) -> OpOperandIterator: ... - -class AffineAddExpr(AffineBinaryExpr): - @staticmethod - def get(arg0: AffineExpr, arg1: AffineExpr) -> AffineAddExpr: ... - @staticmethod - def isinstance(other: AffineExpr) -> bool: ... - def __init__(self, expr: AffineExpr) -> None: ... - -class AffineBinaryExpr(AffineExpr): - @staticmethod - def isinstance(other: AffineExpr) -> bool: ... - def __init__(self, expr: AffineExpr) -> None: ... - @property - def lhs(self) -> AffineExpr: ... - @property - def rhs(self) -> AffineExpr: ... - -class AffineCeilDivExpr(AffineBinaryExpr): - @staticmethod - def get(arg0: AffineExpr, arg1: AffineExpr) -> AffineCeilDivExpr: ... - @staticmethod - def isinstance(other: AffineExpr) -> bool: ... - def __init__(self, expr: AffineExpr) -> None: ... - -class AffineConstantExpr(AffineExpr): - @staticmethod - def get(value: int, context: Context | None = None) -> AffineConstantExpr: ... - @staticmethod - def isinstance(other: AffineExpr) -> bool: ... - def __init__(self, expr: AffineExpr) -> None: ... - @property - def value(self) -> int: ... - -class AffineDimExpr(AffineExpr): - @staticmethod - def get(position: int, context: Context | None = None) -> AffineDimExpr: ... - @staticmethod - def isinstance(other: AffineExpr) -> bool: ... - def __init__(self, expr: AffineExpr) -> None: ... - @property - def position(self) -> int: ... - -class AffineExprList: - def __add__(self, arg0: AffineExprList) -> list[AffineExpr]: ... - -class AffineFloorDivExpr(AffineBinaryExpr): - @staticmethod - def get(arg0: AffineExpr, arg1: AffineExpr) -> AffineFloorDivExpr: ... - @staticmethod - def isinstance(other: AffineExpr) -> bool: ... - def __init__(self, expr: AffineExpr) -> None: ... - -class AffineMap: - @staticmethod - def compress_unused_symbols( - arg0: list, arg1: Context | None - ) -> list[AffineMap]: ... - @staticmethod - def get( - dim_count: int, - symbol_count: int, - exprs: list, - context: Context | None = None, - ) -> AffineMap: - """ - Gets a map with the given expressions as results. - """ - @staticmethod - def get_constant(value: int, context: Context | None = None) -> AffineMap: - """ - Gets an affine map with a single constant result - """ - @staticmethod - def get_empty(context: Context | None = None) -> AffineMap: - """ - Gets an empty affine map. - """ - @staticmethod - def get_identity(n_dims: int, context: Context | None = None) -> AffineMap: - """ - Gets an identity map with the given number of dimensions. - """ - @staticmethod - def get_minor_identity( - n_dims: int, n_results: int, context: Context | None = None - ) -> AffineMap: - """ - Gets a minor identity map with the given number of dimensions and results. - """ - @staticmethod - def get_permutation( - permutation: list[int], context: Context | None = None - ) -> AffineMap: - """ - Gets an affine map that permutes its inputs. - """ - def _CAPICreate(self) -> AffineMap: ... - @overload - def __eq__(self, arg0: AffineMap) -> bool: ... - @overload - def __eq__(self, arg0: object) -> bool: ... - def __hash__(self) -> int: ... - def dump(self) -> None: - """ - Dumps a debug representation of the object to stderr. - """ - def get_major_submap(self, n_results: int) -> AffineMap: ... - def get_minor_submap(self, n_results: int) -> AffineMap: ... - def get_submap(self, result_positions: list[int]) -> AffineMap: ... - def replace( - self, - expr: AffineExpr, - replacement: AffineExpr, - n_result_dims: int, - n_result_syms: int, - ) -> AffineMap: ... - @property - def _CAPIPtr(self) -> object: ... - @property - def context(self) -> Context: - """ - Context that owns the Affine Map - """ - @property - def is_permutation(self) -> bool: ... - @property - def is_projected_permutation(self) -> bool: ... - @property - def n_dims(self) -> int: ... - @property - def n_inputs(self) -> int: ... - @property - def n_symbols(self) -> int: ... - @property - def results(self) -> AffineMapExprList: ... - -class AffineMapAttr(Attribute): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(affine_map: AffineMap) -> AffineMapAttr: - """ - Gets an attribute wrapping an AffineMap. - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class AffineModExpr(AffineBinaryExpr): - @staticmethod - def get(arg0: AffineExpr, arg1: AffineExpr) -> AffineModExpr: ... - @staticmethod - def isinstance(other: AffineExpr) -> bool: ... - def __init__(self, expr: AffineExpr) -> None: ... - -class AffineMulExpr(AffineBinaryExpr): - @staticmethod - def get(arg0: AffineExpr, arg1: AffineExpr) -> AffineMulExpr: ... - @staticmethod - def isinstance(other: AffineExpr) -> bool: ... - def __init__(self, expr: AffineExpr) -> None: ... - -class AffineSymbolExpr(AffineExpr): - @staticmethod - def get(position: int, context: Context | None = None) -> AffineSymbolExpr: ... - @staticmethod - def isinstance(other: AffineExpr) -> bool: ... - def __init__(self, expr: AffineExpr) -> None: ... - @property - def position(self) -> int: ... - -class ArrayAttr(Attribute): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(attributes: list, context: Context | None = None) -> ArrayAttr: - """ - Gets a uniqued Array attribute - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __add__(self, arg0: list) -> ArrayAttr: ... - def __getitem__(self, arg0: int) -> Attribute: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - def __iter__( - self, - ) -> ArrayAttributeIterator: ... - def __len__(self) -> int: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class ArrayAttributeIterator: - def __iter__(self) -> ArrayAttributeIterator: ... - def __next__(self) -> Attribute: ... - -class AsmState: - @overload - def __init__(self, value: Value, use_local_scope: bool = False) -> None: ... - @overload - def __init__(self, op: _OperationBase, use_local_scope: bool = False) -> None: ... - -class AttrBuilder: - @staticmethod - def contains(arg0: str) -> bool: ... - @staticmethod - def get(arg0: str) -> Callable: ... - @staticmethod - def insert( - attribute_kind: str, attr_builder: Callable, replace: bool = False - ) -> None: - """ - Register an attribute builder for building MLIR attributes from python values. - """ - -class BF16Type(Type): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> BF16Type: - """ - Create a bf16 type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class Block: - @staticmethod - def create_at_start( - parent: Region, - arg_types: list[Type], - arg_locs: Sequence | None = None, - ) -> Block: - """ - Creates and returns a new Block at the beginning of the given region (with given argument types and locations). - """ - @overload - def __eq__(self, arg0: Block) -> bool: ... - @overload - def __eq__(self, arg0: Any) -> bool: ... - def __hash__(self) -> int: ... - def __iter__(self) -> OperationIterator: - """ - Iterates over operations in the block. - """ - def __str__(self) -> str: - """ - Returns the assembly form of the block. - """ - def append(self, operation: _OperationBase) -> None: - """ - Appends an operation to this block. If the operation is currently in another block, it will be moved. - """ - def append_to(self, arg0: Region) -> None: - """ - Append this block to a region, transferring ownership if necessary - """ - def create_after(self, *args, arg_locs: Sequence | None = None) -> Block: - """ - Creates and returns a new Block after this block (with given argument types and locations). - """ - def create_before(self, *args, arg_locs: Sequence | None = None) -> Block: - """ - Creates and returns a new Block before this block (with given argument types and locations). - """ - @property - def _CAPIPtr(self) -> object: ... - @property - def arguments(self) -> BlockArgumentList: - """ - Returns a List of block arguments. - """ - @property - def operations(self) -> OperationList: - """ - Returns a forward-optimized sequence of operations. - """ - @property - def owner(self) -> OpView: - """ - Returns the owning operation of this block. - """ - @property - def region(self) -> Region: - """ - Returns the owning region of this block. - """ - -class BlockArgument(Value): - @staticmethod - def isinstance(other_value: Value) -> bool: ... - def __init__(self, value: Value) -> None: ... - def maybe_downcast(self) -> Any: ... - def set_type(self, type: Type) -> None: ... - @property - def arg_number(self) -> int: ... - @property - def owner(self) -> Block: ... - -class BlockArgumentList: - @overload - def __getitem__(self, arg0: int) -> BlockArgument: ... - @overload - def __getitem__(self, arg0: slice) -> BlockArgumentList: ... - def __len__(self) -> int: ... - def __add__(self, arg0: BlockArgumentList) -> list[BlockArgument]: ... - @property - def types(self) -> list[Type]: ... - -class BlockIterator: - def __iter__(self) -> BlockIterator: ... - def __next__(self) -> Block: ... - -class BlockList: - def __getitem__(self, arg0: int) -> Block: ... - def __iter__(self) -> BlockIterator: ... - def __len__(self) -> int: ... - def append(self, *args, arg_locs: Sequence | None = None) -> Block: - """ - Appends a new block, with argument types as positional args. - - Returns: - The created block. - """ - -class BoolAttr(Attribute): - @staticmethod - def get(value: bool, context: Context | None = None) -> BoolAttr: - """ - Gets an uniqued bool attribute - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __bool__(self: Attribute) -> bool: - """ - Converts the value of the bool attribute to a Python bool - """ - def __init__(self, cast_from_attr: Attribute) -> None: ... - @property - def static_typeid(self) -> TypeID: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - @property - def value(self) -> bool: - """ - Returns the value of the bool attribute - """ - -class ComplexType(Type): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(arg0: Type) -> ComplexType: - """ - Create a complex type - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def element_type(self) -> Type: - """ - Returns element type. - """ - @property - def typeid(self) -> TypeID: ... - -class Context: - current: ClassVar[Context] = ... # read-only - allow_unregistered_dialects: bool - @staticmethod - def _get_live_count() -> int: ... - def _CAPICreate(self) -> object: ... - def __enter__(self) -> Context: ... - def __exit__(self, arg0: Any, arg1: Any, arg2: Any) -> None: ... - def __init__(self) -> None: ... - def _clear_live_operations(self) -> int: ... - def _get_context_again(self) -> Context: ... - def _get_live_module_count(self) -> int: ... - def _get_live_operation_count(self) -> int: ... - def _get_live_operation_objects(self) -> list[Operation]: ... - def append_dialect_registry(self, registry: DialectRegistry) -> None: ... - def attach_diagnostic_handler( - self, callback: Callable[[Diagnostic], bool] - ) -> DiagnosticHandler: - """ - Attaches a diagnostic handler that will receive callbacks - """ - def enable_multithreading(self, enable: bool) -> None: ... - def get_dialect_descriptor(self, dialect_name: str) -> DialectDescriptor: - """ - Gets or loads a dialect by name, returning its descriptor object - """ - def is_registered_operation(self, operation_name: str) -> bool: ... - def load_all_available_dialects(self) -> None: ... - @property - def _CAPIPtr(self) -> object: ... - @property - def d(self) -> Dialects: - """ - Alias for 'dialect' - """ - @property - def dialects(self) -> Dialects: - """ - Gets a container for accessing dialects by name - """ - -class DenseBoolArrayAttr(Attribute): - @staticmethod - def get( - values: Sequence[bool], context: Context | None = None - ) -> DenseBoolArrayAttr: - """ - Gets a uniqued dense array attribute - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __add__(self, arg0: list) -> DenseBoolArrayAttr: ... - def __getitem__(self, arg0: int) -> bool: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - def __iter__( - self, - ) -> DenseBoolArrayIterator: ... - def __len__(self) -> int: ... - @property - def static_typeid(self) -> TypeID: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class DenseBoolArrayIterator: - def __iter__(self) -> DenseBoolArrayIterator: ... - def __next__(self) -> bool: ... - -class DenseElementsAttr(Attribute): - @staticmethod - def get( - array: Buffer, - signless: bool = True, - type: Type | None = None, - shape: list[int] | None = None, - context: Context | None = None, - ) -> DenseElementsAttr: - """ - Gets a DenseElementsAttr from a Python buffer or array. - - When `type` is not provided, then some limited type inferencing is done based - on the buffer format. Support presently exists for 8/16/32/64 signed and - unsigned integers and float16/float32/float64. DenseElementsAttrs of these - types can also be converted back to a corresponding buffer. - - For conversions outside of these types, a `type=` must be explicitly provided - and the buffer contents must be bit-castable to the MLIR internal - representation: - - * Integer types (except for i1): the buffer must be byte aligned to the - next byte boundary. - * Floating point types: Must be bit-castable to the given floating point - size. - * i1 (bool): Bit packed into 8bit words where the bit pattern matches a - row major ordering. An arbitrary Numpy `bool_` array can be bit packed to - this specification with: `np.packbits(ary, axis=None, bitorder='little')`. - - If a single element buffer is passed (or for i1, a single byte with value 0 - or 255), then a splat will be created. - - Args: - array: The array or buffer to convert. - signless: If inferring an appropriate MLIR type, use signless types for - integers (defaults True). - type: Skips inference of the MLIR element type and uses this instead. The - storage size must be consistent with the actual contents of the buffer. - shape: Overrides the shape of the buffer when constructing the MLIR - shaped type. This is needed when the physical and logical shape differ (as - for i1). - context: Explicit context, if not from context manager. - - Returns: - DenseElementsAttr on success. - - Raises: - ValueError: If the type of the buffer or array cannot be matched to an MLIR - type or if the buffer does not meet expectations. - """ - @staticmethod - def get_splat(shaped_type: Type, element_attr: Attribute) -> DenseElementsAttr: - """ - Gets a DenseElementsAttr where all values are the same - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - def __len__(self) -> int: ... - def get_splat_value(self) -> Attribute: ... - @property - def is_splat(self) -> bool: ... - @property - def static_typeid(self) -> TypeID: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class DenseF32ArrayAttr(Attribute): - @staticmethod - def get( - values: Sequence[float], context: Context | None = None - ) -> DenseF32ArrayAttr: - """ - Gets a uniqued dense array attribute - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __add__(self, arg0: list) -> DenseF32ArrayAttr: ... - def __getitem__(self, arg0: int) -> float: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - def __iter__( - self, - ) -> DenseF32ArrayIterator: ... - def __len__(self) -> int: ... - @property - def static_typeid(self) -> TypeID: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class DenseF32ArrayIterator: - def __iter__(self) -> DenseF32ArrayIterator: ... - def __next__(self) -> float: ... - -class DenseF64ArrayAttr(Attribute): - @staticmethod - def get( - values: Sequence[float], context: Context | None = None - ) -> DenseF64ArrayAttr: - """ - Gets a uniqued dense array attribute - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __add__(self, arg0: list) -> DenseF64ArrayAttr: ... - def __getitem__(self, arg0: int) -> float: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - def __iter__( - self, - ) -> DenseF64ArrayIterator: ... - def __len__(self) -> int: ... - @property - def static_typeid(self) -> TypeID: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class DenseF64ArrayIterator: - def __iter__(self) -> DenseF64ArrayIterator: ... - def __next__(self) -> float: ... - -class DenseFPElementsAttr(DenseElementsAttr): - @staticmethod - def get( - array: Buffer, - signless: bool = True, - type: Type | None = None, - shape: list[int] | None = None, - context: Context | None = None, - ) -> DenseFPElementsAttr: ... - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __getitem__(self, arg0: int) -> float: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - @property - def static_typeid(self) -> TypeID: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class DenseI16ArrayAttr(Attribute): - @staticmethod - def get(values: Sequence[int], context: Context | None = None) -> DenseI16ArrayAttr: - """ - Gets a uniqued dense array attribute - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __add__(self, arg0: list) -> DenseI16ArrayAttr: ... - def __getitem__(self, arg0: int) -> int: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - def __iter__( - self, - ) -> DenseI16ArrayIterator: ... - def __len__(self) -> int: ... - @property - def static_typeid(self) -> TypeID: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class DenseI16ArrayIterator: - def __iter__(self) -> DenseI16ArrayIterator: ... - def __next__(self) -> int: ... - -class DenseI32ArrayAttr(Attribute): - @staticmethod - def get(values: Sequence[int], context: Context | None = None) -> DenseI32ArrayAttr: - """ - Gets a uniqued dense array attribute - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __add__(self, arg0: list) -> DenseI32ArrayAttr: ... - def __getitem__(self, arg0: int) -> int: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - def __iter__( - self, - ) -> DenseI32ArrayIterator: ... - def __len__(self) -> int: ... - @property - def static_typeid(self) -> TypeID: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class DenseI32ArrayIterator: - def __iter__(self) -> DenseI32ArrayIterator: ... - def __next__(self) -> int: ... - -class DenseI64ArrayAttr(Attribute): - @staticmethod - def get(values: Sequence[int], context: Context | None = None) -> DenseI64ArrayAttr: - """ - Gets a uniqued dense array attribute - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __add__(self, arg0: list) -> DenseI64ArrayAttr: ... - def __getitem__(self, arg0: int) -> int: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - def __iter__( - self, - ) -> DenseI16ArrayIterator: ... - def __len__(self) -> int: ... - @property - def static_typeid(self) -> TypeID: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class DenseI64ArrayIterator: - def __iter__(self) -> DenseI64ArrayIterator: ... - def __next__(self) -> int: ... - -class DenseI8ArrayAttr(Attribute): - @staticmethod - def get(values: Sequence[int], context: Context | None = None) -> DenseI8ArrayAttr: - """ - Gets a uniqued dense array attribute - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __add__(self, arg0: list) -> DenseI8ArrayAttr: ... - def __getitem__(self, arg0: int) -> int: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - def __iter__( - self, - ) -> DenseI8ArrayIterator: ... - def __len__(self) -> int: ... - @property - def static_typeid(self) -> TypeID: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class DenseI8ArrayIterator: - def __iter__(self) -> DenseI8ArrayIterator: ... - def __next__(self) -> int: ... - -class DenseIntElementsAttr(DenseElementsAttr): - @staticmethod - def get( - array: Buffer, - signless: bool = True, - type: Type | None = None, - shape: list[int] | None = None, - context: Context | None = None, - ) -> DenseIntElementsAttr: ... - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __getitem__(self, arg0: int) -> int: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - @property - def static_typeid(self) -> TypeID: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class DenseResourceElementsAttr(Attribute): - @staticmethod - def get_from_buffer( - array: Buffer, - name: str, - type: Type, - alignment: int | None = None, - is_mutable: bool = False, - context: Context | None = None, - ) -> DenseResourceElementsAttr: - """ - Gets a DenseResourceElementsAttr from a Python buffer or array. - - This function does minimal validation or massaging of the data, and it is - up to the caller to ensure that the buffer meets the characteristics - implied by the shape. - - The backing buffer and any user objects will be retained for the lifetime - of the resource blob. This is typically bounded to the context but the - resource can have a shorter lifespan depending on how it is used in - subsequent processing. - - Args: - buffer: The array or buffer to convert. - name: Name to provide to the resource (may be changed upon collision). - type: The explicit ShapedType to construct the attribute with. - context: Explicit context, if not from context manager. - - Returns: - DenseResourceElementsAttr on success. - - Raises: - ValueError: If the type of the buffer or array cannot be matched to an MLIR - type or if the buffer does not meet expectations. - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - @property - def static_typeid(self) -> TypeID: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class Diagnostic: - @property - def location(self) -> Location: ... - @property - def message(self) -> str: ... - @property - def notes(self) -> tuple[Diagnostic]: ... - @property - def severity(self) -> DiagnosticSeverity: ... - -class DiagnosticHandler: - def __enter__(self) -> DiagnosticHandler: ... - def __exit__(self, arg0: object, arg1: object, arg2: object) -> None: ... - def detach(self) -> None: ... - @property - def attached(self) -> bool: ... - @property - def had_error(self) -> bool: ... - -class DiagnosticInfo: - def __init__(self, arg0: Diagnostic) -> None: ... - @property - def location(self) -> Location: ... - @property - def message(self) -> str: ... - @property - def notes(self) -> list[DiagnosticInfo]: ... - @property - def severity(self) -> DiagnosticSeverity: ... - -class DiagnosticSeverity: - """ - Members: - - ERROR - - WARNING - - NOTE - - REMARK - """ - - ERROR: ClassVar[DiagnosticSeverity] # value = - NOTE: ClassVar[DiagnosticSeverity] # value = - REMARK: ClassVar[DiagnosticSeverity] # value = - WARNING: ClassVar[DiagnosticSeverity] # value = - __members__: ClassVar[ - dict[str, DiagnosticSeverity] - ] # value = {'ERROR': , 'WARNING': , 'NOTE': , 'REMARK': } - def __eq__(self, other: Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... - def __init__(self, value: int) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: Any) -> bool: ... - def __setstate__(self, state: int) -> None: ... - @property - def name(self) -> str: ... - @property - def value(self) -> int: ... - -class Dialect: - def __init__(self, descriptor: DialectDescriptor) -> None: ... - @property - def descriptor(self) -> DialectDescriptor: ... - -class DialectDescriptor: - @property - def namespace(self) -> str: ... - -class DialectRegistry: - def _CAPICreate(self) -> DialectRegistry: ... - def __init__(self) -> None: ... - @property - def _CAPIPtr(self) -> object: ... - -class Dialects: - def __getattr__(self, arg0: str) -> Dialect: ... - def __getitem__(self, arg0: str) -> Dialect: ... - -class DictAttr(Attribute): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(value: dict = {}, context: Context | None = None) -> DictAttr: - """ - Gets an uniqued Dict attribute - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __contains__(self, arg0: str) -> bool: ... - @overload - def __getitem__(self, arg0: str) -> Attribute: ... - @overload - def __getitem__(self, arg0: int) -> NamedAttribute: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - def __len__(self) -> int: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class FloatType(Type): - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def width(self) -> int: - """ - Returns the width of the floating-point type. - """ - -class F16Type(FloatType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> F16Type: - """ - Create a f16 type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class F32Type(FloatType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> F32Type: - """ - Create a f32 type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class F64Type(FloatType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> F64Type: - """ - Create a f64 type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class FlatSymbolRefAttr(Attribute): - @staticmethod - def get(value: str, context: Context | None = None) -> FlatSymbolRefAttr: - """ - Gets a uniqued FlatSymbolRef attribute - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - @property - def static_typeid(self) -> TypeID: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - @property - def value(self) -> str: - """ - Returns the value of the FlatSymbolRef attribute as a string - """ - -class Float4E2M1FNType(FloatType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> Float4E2M1FNType: - """ - Create a float4_e2m1fn type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class Float6E2M3FNType(FloatType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> Float6E2M3FNType: - """ - Create a float6_e2m3fn type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class Float6E3M2FNType(FloatType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> Float6E3M2FNType: - """ - Create a float6_e3m2fn type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class Float8E3M4Type(FloatType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> Float8E3M4Type: - """ - Create a float8_e3m4 type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class Float8E4M3B11FNUZType(FloatType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> Float8E4M3B11FNUZType: - """ - Create a float8_e4m3b11fnuz type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class Float8E4M3FNType(FloatType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> Float8E4M3FNType: - """ - Create a float8_e4m3fn type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class Float8E4M3FNUZType(FloatType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> Float8E4M3FNUZType: - """ - Create a float8_e4m3fnuz type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class Float8E4M3Type(FloatType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> Float8E4M3Type: - """ - Create a float8_e4m3 type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class Float8E5M2FNUZType(FloatType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> Float8E5M2FNUZType: - """ - Create a float8_e5m2fnuz type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class Float8E5M2Type(FloatType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> Float8E5M2Type: - """ - Create a float8_e5m2 type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class Float8E8M0FNUType(FloatType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> Float8E8M0FNUType: - """ - Create a float8_e8m0fnu type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class FloatAttr(Attribute): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(type: Type, value: float, loc: Location | None = None) -> FloatAttr: - """ - Gets an uniqued float point attribute associated to a type - """ - @staticmethod - def get_f32(value: float, context: Context | None = None) -> FloatAttr: - """ - Gets an uniqued float point attribute associated to a f32 type - """ - @staticmethod - def get_f64(value: float, context: Context | None = None) -> FloatAttr: - """ - Gets an uniqued float point attribute associated to a f64 type - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __float__(self: Attribute) -> float: - """ - Converts the value of the float attribute to a Python float - """ - def __init__(self, cast_from_attr: Attribute) -> None: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - @property - def value(self) -> float: - """ - Returns the value of the float attribute - """ - -class FloatTF32Type(FloatType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> FloatTF32Type: - """ - Create a tf32 type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class FunctionType(Type): - static_typeid: ClassVar[TypeID] - @staticmethod - def get( - inputs: list[Type], results: list[Type], context: Context | None = None - ) -> FunctionType: - """ - Gets a FunctionType from a List of input and result types - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def inputs(self) -> list: - """ - Returns the List of input types in the FunctionType. - """ - @property - def results(self) -> list: - """ - Returns the List of result types in the FunctionType. - """ - @property - def typeid(self) -> TypeID: ... - -class IndexType(Type): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> IndexType: - """ - Create a index type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class InferShapedTypeOpInterface: - def __init__(self, object: object, context: Context | None = None) -> None: - """ - Creates an interface from a given operation/opview object or from a - subclass of OpView. Raises ValueError if the operation does not implement the - interface. - """ - def inferReturnTypeComponents( - self, - operands: list | None = None, - attributes: Attribute | None = None, - properties=None, - regions: list[Region] | None = None, - context: Context | None = None, - loc: Location | None = None, - ) -> list[ShapedTypeComponents]: - """ - Given the arguments required to build an operation, attempts to infer - its return shaped type components. Raises ValueError on failure. - """ - @property - def operation(self) -> Operation: - """ - Returns an Operation for which the interface was constructed. - """ - @property - def opview(self) -> OpView: - """ - Returns an OpView subclass _instance_ for which the interface was - constructed - """ - -class InferTypeOpInterface: - def __init__(self, object: object, context: Context | None = None) -> None: - """ - Creates an interface from a given operation/opview object or from a - subclass of OpView. Raises ValueError if the operation does not implement the - interface. - """ - def inferReturnTypes( - self, - operands: list | None = None, - attributes: Attribute | None = None, - properties=None, - regions: list[Region] | None = None, - context: Context | None = None, - loc: Location | None = None, - ) -> list[Type]: - """ - Given the arguments required to build an operation, attempts to infer - its return types. Raises ValueError on failure. - """ - @property - def operation(self) -> Operation: - """ - Returns an Operation for which the interface was constructed. - """ - @property - def opview(self) -> OpView: - """ - Returns an OpView subclass _instance_ for which the interface was - constructed - """ - -class InsertionPoint: - current: ClassVar[InsertionPoint] = ... # read-only - @staticmethod - def at_block_begin(block: Block) -> InsertionPoint: - """ - Inserts at the beginning of the block. - """ - @staticmethod - def at_block_terminator(block: Block) -> InsertionPoint: - """ - Inserts before the block terminator. - """ - def __enter__(self) -> InsertionPoint: ... - def __exit__(self, arg0: Any, arg1: Any, arg2: Any) -> None: ... - @overload - def __init__(self, block: Block) -> None: - """ - Inserts after the last operation but still inside the block. - """ - @overload - def __init__(self, beforeOperation: _OperationBase) -> None: - """ - Inserts before a referenced operation. - """ - def insert(self, operation: _OperationBase) -> None: - """ - Inserts an operation. - """ - @property - def block(self) -> Block: - """ - Returns the block that this InsertionPoint points to. - """ - @property - def ref_operation(self) -> _OperationBase | None: - """ - The reference operation before which new operations are inserted, or None if the insertion point is at the end of the block - """ - -class IntegerAttr(Attribute): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(type: Type, value: int) -> IntegerAttr: - """ - Gets an uniqued integer attribute associated to a type - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - def __int__(self) -> int: - """ - Converts the value of the integer attribute to a Python int - """ - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - @property - def value(self) -> int: - """ - Returns the value of the integer attribute - """ - -class IntegerSet: - @staticmethod - def get( - num_dims: int, - num_symbols: int, - exprs: list, - eq_flags: list[bool], - context: Context | None = None, - ) -> IntegerSet: ... - @staticmethod - def get_empty( - num_dims: int, num_symbols: int, context: Context | None = None - ) -> IntegerSet: ... - def _CAPICreate(self) -> IntegerSet: ... - @overload - def __eq__(self, arg0: IntegerSet) -> bool: ... - @overload - def __eq__(self, arg0: object) -> bool: ... - def __hash__(self) -> int: ... - def dump(self) -> None: - """ - Dumps a debug representation of the object to stderr. - """ - def get_replaced( - self, - dim_exprs: list, - symbol_exprs: list, - num_result_dims: int, - num_result_symbols: int, - ) -> IntegerSet: ... - @property - def _CAPIPtr(self) -> object: ... - @property - def constraints(self) -> IntegerSetConstraintList: ... - @property - def context(self) -> Context: ... - @property - def is_canonical_empty(self) -> bool: ... - @property - def n_dims(self) -> int: ... - @property - def n_equalities(self) -> int: ... - @property - def n_inequalities(self) -> int: ... - @property - def n_inputs(self) -> int: ... - @property - def n_symbols(self) -> int: ... - -class IntegerSetAttr(Attribute): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(integer_set) -> IntegerSetAttr: - """ - Gets an attribute wrapping an IntegerSet. - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class IntegerSetConstraint: - def __init__(self, *args, **kwargs) -> None: ... - @property - def expr(self) -> AffineExpr: ... - @property - def is_eq(self) -> bool: ... - -class IntegerSetConstraintList: - def __init__(self, *args, **kwargs) -> None: ... - def __add__(self, arg0: IntegerSetConstraintList) -> list[IntegerSetConstraint]: ... - @overload - def __getitem__(self, arg0: int) -> IntegerSetConstraint: ... - @overload - def __getitem__(self, arg0: slice) -> IntegerSetConstraintList: ... - def __len__(self) -> int: ... - -class IntegerType(Type): - static_typeid: ClassVar[TypeID] - @staticmethod - def get_signed(width: int, context: Context | None = None) -> IntegerType: - """ - Create a signed integer type - """ - @staticmethod - def get_signless(width: int, context: Context | None = None) -> IntegerType: - """ - Create a signless integer type - """ - @staticmethod - def get_unsigned(width: int, context: Context | None = None) -> IntegerType: - """ - Create an unsigned integer type - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def is_signed(self) -> bool: - """ - Returns whether this is a signed integer - """ - @property - def is_signless(self) -> bool: - """ - Returns whether this is a signless integer - """ - @property - def is_unsigned(self) -> bool: - """ - Returns whether this is an unsigned integer - """ - @property - def typeid(self) -> TypeID: ... - @property - def width(self) -> int: - """ - Returns the width of the integer type - """ - -class Location: - current: ClassVar[Location] = ... # read-only - __hash__: ClassVar[None] = None - @staticmethod - def callsite( - callee: Location, frames: Sequence[Location], context: Context | None = None - ) -> Location: - """ - Gets a Location representing a caller and callsite - """ - @staticmethod - def file( - filename: str, line: int, col: int, context: Context | None = None - ) -> Location: - """ - Gets a Location representing a file, line and column - """ - @staticmethod - def from_attr(attribute: Attribute, context: Context | None = None) -> Location: - """ - Gets a Location from a LocationAttr - """ - @staticmethod - def fused( - locations: Sequence[Location], - metadata: Attribute | None = None, - context: Context | None = None, - ) -> Location: - """ - Gets a Location representing a fused location with optional metadata - """ - @staticmethod - def name( - name: str, - childLoc: Location | None = None, - context: Context | None = None, - ) -> Location: - """ - Gets a Location representing a named location with optional child location - """ - @staticmethod - def unknown(context: Context | None = None) -> Location: - """ - Gets a Location representing an unknown location - """ - def _CAPICreate(self) -> Location: ... - def __enter__(self) -> Location: ... - @overload - def __eq__(self, arg0: Location) -> bool: ... - @overload - def __eq__(self, arg0: Location) -> bool: ... - def __exit__(self, arg0: object, arg1: object, arg2: object) -> None: ... - def emit_error(self, message: str) -> None: - """ - Emits an error at this location - """ - @property - def _CAPIPtr(self) -> object: ... - @property - def attr(self) -> Attribute: - """ - Get the underlying LocationAttr - """ - @property - def context(self) -> Context: - """ - Context that owns the Location - """ - -class MemRefType(ShapedType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get( - shape: list[int], - element_type: Type, - layout: Attribute = None, - memory_space: Attribute = None, - loc: Location | None = None, - ) -> MemRefType: - """ - Create a memref type - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def affine_map(self) -> AffineMap: - """ - The layout of the MemRef type as an affine map. - """ - @property - def layout(self) -> Attribute: - """ - The layout of the MemRef type. - """ - @property - def memory_space(self) -> Attribute | None: - """ - Returns the memory space of the given MemRef type. - """ - @property - def typeid(self) -> TypeID: ... - def get_strides_and_offset(self) -> tuple[list[int], int]: - """ - The strides and offset of the MemRef type. - """ - -class Module: - @staticmethod - def create(loc: Location | None = None) -> Module: - """ - Creates an empty module - """ - @staticmethod - def parse(asm: str | bytes, context: Context | None = None) -> Module: - """ - Parses a module's assembly format from a string. - - Returns a new MlirModule or raises an MLIRError if the parsing fails. - - See also: https://mlir.llvm.org/docs/LangRef/ - """ - @staticmethod - def parseFile(path: str, context: Context | None = None) -> Module: - """ - Parses a module's assembly format from file. - - Returns a new MlirModule or raises an MLIRError if the parsing fails. - - See also: https://mlir.llvm.org/docs/LangRef/ - """ - def _CAPICreate(self) -> Any: ... - def __str__(self) -> str: - """ - Gets the assembly form of the operation with default options. - - If more advanced control over the assembly formatting or I/O options is needed, - use the dedicated print or get_asm method, which supports keyword arguments to - customize behavior. - """ - def dump(self) -> None: - """ - Dumps a debug representation of the object to stderr. - """ - @property - def _CAPIPtr(self) -> object: ... - @property - def body(self) -> Block: - """ - Return the block for this module - """ - @property - def context(self) -> Context: - """ - Context that created the Module - """ - @property - def operation(self) -> Operation: - """ - Accesses the module as an operation - """ - -class MLIRError(Exception): - def __init__( - self, message: str, error_diagnostics: list[DiagnosticInfo] - ) -> None: ... - -class NamedAttribute: - @property - def attr(self) -> Attribute: - """ - The underlying generic attribute of the NamedAttribute binding - """ - @property - def name(self) -> str: - """ - The name of the NamedAttribute binding - """ - -class NoneType(Type): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> NoneType: - """ - Create a none type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class OpAttributeMap: - def __contains__(self, arg0: str) -> bool: ... - def __delitem__(self, arg0: str) -> None: ... - @overload - def __getitem__(self, arg0: str) -> Attribute: ... - @overload - def __getitem__(self, arg0: int) -> NamedAttribute: ... - def __len__(self) -> int: ... - def __setitem__(self, arg0: str, arg1: Attribute) -> None: ... - -class OpOperand: - @property - def operand_number(self) -> int: ... - @property - def owner(self) -> _OperationBase: ... - -class OpOperandIterator: - def __iter__(self) -> OpOperandIterator: ... - def __next__(self) -> OpOperand: ... - -class OpOperandList: - def __add__(self, arg0: OpOperandList) -> list[Value]: ... - @overload - def __getitem__(self, arg0: int) -> Value: ... - @overload - def __getitem__(self, arg0: slice) -> OpOperandList: ... - def __len__(self) -> int: ... - def __setitem__(self, arg0: int, arg1: Value) -> None: ... - -class OpResult(Value): - @staticmethod - def isinstance(other_value: Value) -> bool: ... - def __init__(self, value: Value) -> None: ... - @staticmethod - def isinstance(arg: Any) -> bool: ... - @property - def owner(self) -> _OperationBase: ... - @property - def result_number(self) -> int: ... - -class OpResultList: - def __add__(self, arg0: OpResultList) -> list[OpResult]: ... - @overload - def __getitem__(self, arg0: int) -> OpResult: ... - @overload - def __getitem__(self, arg0: slice) -> OpResultList: ... - def __len__(self) -> int: ... - @property - def owner(self) -> _OperationBase: ... - @property - def types(self) -> list[Type]: ... - -class OpSuccessors: - def __add__(self, arg0: OpSuccessors) -> list[Block]: ... - @overload - def __getitem__(self, arg0: int) -> Block: ... - @overload - def __getitem__(self, arg0: slice) -> OpSuccessors: ... - def __setitem__(self, arg0: int, arg1: Block) -> None: ... - def __len__(self) -> int: ... - -class OpView(_OperationBase): - _ODS_OPERAND_SEGMENTS: ClassVar[None] = ... - _ODS_REGIONS: ClassVar[tuple] = ... - _ODS_RESULT_SEGMENTS: ClassVar[None] = ... - def __init__(self, operation: _OperationBase) -> None: ... - @classmethod - def build_generic( - cls: type[_TOperation], - results: Sequence[Type] | None = None, - operands: Sequence[Value] | None = None, - attributes: dict[str, Attribute] | None = None, - successors: Sequence[Block] | None = None, - regions: int | None = None, - loc: Location | None = None, - ip: InsertionPoint | None = None, - ) -> _TOperation: - """ - Builds a specific, generated OpView based on class level attributes. - """ - @classmethod - def parse( - cls: type[_TOperation], - source: str | bytes, - *, - source_name: str = "", - context: Context | None = None, - ) -> _TOperation: - """ - Parses a specific, generated OpView based on class level attributes - """ - def __init__(self, operation: _OperationBase) -> None: ... - @property - def operation(self) -> _OperationBase: ... - @property - def opview(self) -> OpView: ... - @property - def successors(self) -> OpSuccessors: - """ - Returns the List of Operation successors. - """ - -class OpaqueAttr(Attribute): - static_typeid: ClassVar[TypeID] - @staticmethod - def get( - dialect_namespace: str, - buffer: Buffer, - type: Type, - context: Context | None = None, - ) -> OpaqueAttr: - """ - Gets an Opaque attribute. - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - @property - def data(self) -> bytes: - """ - Returns the data for the Opaqued attributes as `bytes` - """ - @property - def dialect_namespace(self) -> str: - """ - Returns the dialect namespace for the Opaque attribute as a string - """ - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class OpaqueType(Type): - static_typeid: ClassVar[TypeID] - @staticmethod - def get( - dialect_namespace: str, buffer: str, context: Context | None = None - ) -> OpaqueType: - """ - Create an unregistered (opaque) dialect type. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def data(self) -> str: - """ - Returns the data for the Opaque type as a string. - """ - @property - def dialect_namespace(self) -> str: - """ - Returns the dialect namespace for the Opaque type as a string. - """ - @property - def typeid(self) -> TypeID: ... - -class Operation(_OperationBase): - def _CAPICreate(self) -> object: ... - @staticmethod - def create( - name: str, - results: Sequence[Type] | None = None, - operands: Sequence[Value] | None = None, - attributes: dict[str, Attribute] | None = None, - successors: Sequence[Block] | None = None, - regions: int = 0, - loc: Location | None = None, - ip: InsertionPoint | None = None, - infer_type: bool = False, - ) -> Operation: - """ - Creates a new operation. - - Args: - name: Operation name (e.g. "dialect.operation"). - results: Sequence of Type representing op result types. - attributes: Dict of str:Attribute. - successors: List of Block for the operation's successors. - regions: Number of regions to create. - loc: A Location object (defaults to resolve from context manager). - ip: An InsertionPoint (defaults to resolve from context manager or set to - False to disable insertion, even with an insertion point set in the - context manager). - infer_type: Whether to infer result types. - Returns: - A new "detached" Operation object. Detached operations can be added - to blocks, which causes them to become "attached." - """ - @staticmethod - def parse( - source: str | bytes, *, source_name: str = "", context: Context | None = None - ) -> Operation: - """ - Parses an operation. Supports both text assembly format and binary bytecode format. - """ - def _CAPICreate(self) -> object: ... - @property - def _CAPIPtr(self) -> object: ... - @property - def operation(self) -> Operation: ... - @property - def opview(self) -> OpView: ... - @property - def successors(self) -> OpSuccessors: - """ - Returns the List of Operation successors. - """ - -class OperationIterator: - def __iter__(self) -> OperationIterator: ... - def __next__(self) -> OpView: ... - -class OperationList: - def __getitem__(self, arg0: int) -> OpView: ... - def __iter__(self) -> OperationIterator: ... - def __len__(self) -> int: ... - -class RankedTensorType(ShapedType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get( - shape: list[int], - element_type: Type, - encoding: Attribute | None = None, - loc: Location | None = None, - ) -> RankedTensorType: - """ - Create a ranked tensor type - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def encoding(self) -> Attribute | None: ... - @property - def typeid(self) -> TypeID: ... - -class Region: - __hash__: ClassVar[None] = None - @overload - def __eq__(self, arg0: Region) -> bool: ... - @overload - def __eq__(self, arg0: object) -> bool: ... - def __iter__(self) -> BlockIterator: - """ - Iterates over blocks in the region. - """ - @property - def blocks(self) -> BlockList: - """ - Returns a forward-optimized sequence of blocks. - """ - @property - def owner(self) -> OpView: - """ - Returns the operation owning this region. - """ - -class RegionIterator: - def __iter__(self) -> RegionIterator: ... - def __next__(self) -> Region: ... - -class RegionSequence: - @overload - def __getitem__(self, arg0: int) -> Region: ... - @overload - def __getitem__(self, arg0: slice) -> Sequence[Region]: ... - def __iter__(self) -> RegionIterator: ... - def __len__(self) -> int: ... - -class ShapedType(Type): - @staticmethod - def get_dynamic_size() -> int: - """ - Returns the value used to indicate dynamic dimensions in shaped types. - """ - @staticmethod - def get_dynamic_stride_or_offset() -> int: - """ - Returns the value used to indicate dynamic strides or offsets in shaped types. - """ - @staticmethod - def is_dynamic_size(dim_size: int) -> bool: - """ - Returns whether the given dimension size indicates a dynamic dimension. - """ - @staticmethod - def is_static_size(dim_size: int) -> bool: - """ - Returns whether the given dimension size indicates a static dimension. - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - def get_dim_size(self, dim: int) -> int: - """ - Returns the dim-th dimension of the given ranked shaped type. - """ - def is_dynamic_dim(self, dim: int) -> bool: - """ - Returns whether the dim-th dimension of the given shaped type is dynamic. - """ - def is_static_dim(self, dim: int) -> bool: - """ - Returns whether the dim-th dimension of the given shaped type is static. - """ - def is_dynamic_stride_or_offset(self, dim_size: int) -> bool: - """ - Returns whether the given value is used as a placeholder for dynamic strides and offsets in shaped types. - """ - def is_static_stride_or_offset(self, dim_size: int) -> bool: - """ - Returns whether the given shaped type stride or offset value is statically-sized. - """ - @property - def element_type(self) -> Type: - """ - Returns the element type of the shaped type. - """ - @property - def has_rank(self) -> bool: - """ - Returns whether the given shaped type is ranked. - """ - @property - def has_static_shape(self) -> bool: - """ - Returns whether the given shaped type has a static shape. - """ - @property - def rank(self) -> int: - """ - Returns the rank of the given ranked shaped type. - """ - @property - def shape(self) -> list[int]: - """ - Returns the shape of the ranked shaped type as a List of integers. - """ - @property - def static_typeid(self) -> TypeID: ... - @property - def typeid(self) -> TypeID: ... - -class ShapedTypeComponents: - @staticmethod - @overload - def get(element_type: Type) -> ShapedTypeComponents: - """ - Create an shaped type components object with only the element type. - """ - @staticmethod - @overload - def get(shape: list, element_type: Type) -> ShapedTypeComponents: - """ - Create a ranked shaped type components object. - """ - @staticmethod - @overload - def get( - shape: list, element_type: Type, attribute: Attribute - ) -> ShapedTypeComponents: - """ - Create a ranked shaped type components object with attribute. - """ - @property - def element_type(self) -> Type: - """ - Returns the element type of the shaped type components. - """ - @property - def has_rank(self) -> bool: - """ - Returns whether the given shaped type component is ranked. - """ - @property - def rank(self) -> int: - """ - Returns the rank of the given ranked shaped type components. If the shaped type components does not have a rank, None is returned. - """ - @property - def shape(self) -> list[int]: - """ - Returns the shape of the ranked shaped type components as a List of integers. Returns none if the shaped type component does not have a rank. - """ - -class StridedLayoutAttr(Attribute): - static_typeid: ClassVar[TypeID] - @staticmethod - def get( - offset: int, strides: list[int], context: Context | None = None - ) -> StridedLayoutAttr: - """ - Gets a strided layout attribute. - """ - @staticmethod - def get_fully_dynamic( - rank: int, context: Context | None = None - ) -> StridedLayoutAttr: - """ - Gets a strided layout attribute with dynamic offset and strides of a given rank. - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - @property - def offset(self) -> int: - """ - Returns the value of the float point attribute - """ - @property - def strides(self) -> list[int]: - """ - Returns the value of the float point attribute - """ - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class StringAttr(Attribute): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(value: str | bytes, context: Context | None = None) -> StringAttr: - """ - Gets a uniqued string attribute - """ - @staticmethod - def get_typed(type: Type, value: str) -> StringAttr: - """ - Gets a uniqued string attribute associated to a type - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - @property - def value(self) -> str: - """ - Returns the value of the string attribute - """ - @property - def value_bytes(self) -> bytes: - """ - Returns the value of the string attribute as `bytes` - """ - -class SymbolRefAttr(Attribute): - @staticmethod - def get(symbols: list[str], context: Context | None = None) -> Attribute: - """ - Gets a uniqued SymbolRef attribute from a List of symbol names - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - @property - def static_typeid(self) -> TypeID: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - @property - def value(self) -> list[str]: - """ - Returns the value of the SymbolRef attribute as a List[str] - """ - -class SymbolTable: - @staticmethod - def get_symbol_name(symbol: _OperationBase) -> Attribute: ... - @staticmethod - def get_visibility(symbol: _OperationBase) -> Attribute: ... - @staticmethod - def replace_all_symbol_uses( - old_symbol: str, new_symbol: str, from_op: _OperationBase - ) -> None: ... - @staticmethod - def set_symbol_name(symbol: _OperationBase, name: str) -> None: ... - @staticmethod - def set_visibility(symbol: _OperationBase, visibility: str) -> None: ... - @staticmethod - def walk_symbol_tables( - from_op: _OperationBase, - all_sym_uses_visible: bool, - callback: Callable[[_OperationBase, bool], None], - ) -> None: ... - def __contains__(self, arg0: str) -> bool: ... - def __delitem__(self, arg0: str) -> None: ... - def __getitem__(self, arg0: str) -> OpView: ... - def __init__(self, arg0: _OperationBase) -> None: ... - def erase(self, operation: _OperationBase) -> None: ... - def insert(self, operation: _OperationBase) -> Attribute: ... - -class TupleType(Type): - static_typeid: ClassVar[TypeID] - @staticmethod - def get_tuple(elements: list[Type], context: Context | None = None) -> TupleType: - """ - Create a Tuple type - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - def get_type(self, pos: int) -> Type: - """ - Returns the pos-th type in the Tuple type. - """ - @property - def num_types(self) -> int: - """ - Returns the number of types contained in a Tuple. - """ - @property - def typeid(self) -> TypeID: ... - -class TypeAttr(Attribute): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(value: Type, context: Context | None = None) -> TypeAttr: - """ - Gets a uniqued Type attribute - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - @property - def value(self) -> Type: ... - -class TypeID: - def _CAPICreate(self) -> TypeID: ... - @overload - def __eq__(self, arg0: TypeID) -> bool: ... - @overload - def __eq__(self, arg0: Any) -> bool: ... - def __hash__(self) -> int: ... - @property - def _CAPIPtr(self) -> object: ... - -class UnitAttr(Attribute): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(context: Context | None = None) -> UnitAttr: - """ - Create a Unit attribute. - """ - @staticmethod - def isinstance(other: Attribute) -> bool: ... - def __init__(self, cast_from_attr: Attribute) -> None: ... - @property - def type(self) -> Type: ... - @property - def typeid(self) -> TypeID: ... - -class UnrankedMemRefType(ShapedType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get( - element_type: Type, memory_space: Attribute, loc: Location | None = None - ) -> UnrankedMemRefType: - """ - Create a unranked memref type - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def memory_space(self) -> Attribute | None: - """ - Returns the memory space of the given Unranked MemRef type. - """ - @property - def typeid(self) -> TypeID: ... - -class UnrankedTensorType(ShapedType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get(element_type: Type, loc: Location | None = None) -> UnrankedTensorType: - """ - Create a unranked tensor type - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def typeid(self) -> TypeID: ... - -class VectorType(ShapedType): - static_typeid: ClassVar[TypeID] - @staticmethod - def get( - shape: list[int], - element_type: Type, - *, - scalable: list | None = None, - scalable_dims: list[int] | None = None, - loc: Location | None = None, - ) -> VectorType: - """ - Create a vector type - """ - @staticmethod - def isinstance(other: Type) -> bool: ... - def __init__(self, cast_from_type: Type) -> None: ... - @property - def scalable(self) -> bool: ... - @property - def scalable_dims(self) -> list[bool]: ... - @property - def typeid(self) -> TypeID: ... - -class _GlobalDebug: - flag: ClassVar[bool] = False diff --git a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi b/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi deleted file mode 100644 index 1010daddae2aa..0000000000000 --- a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi +++ /dev/null @@ -1,36 +0,0 @@ -# Originally imported via: -# stubgen {...} -m mlir._mlir_libs._mlir.passmanager -# Local modifications: -# * Relative imports for cross-module references. -# * Add __all__ - - -from . import ir as _ir - -__all__ = [ - "PassManager", -] - -class PassManager: - def __init__(self, context: _ir.Context | None = None) -> None: ... - def _CAPICreate(self) -> object: ... - def _testing_release(self) -> None: ... - def enable_ir_printing( - self, - print_before_all: bool = False, - print_after_all: bool = True, - print_module_scope: bool = False, - print_after_change: bool = False, - print_after_failure: bool = False, - large_elements_limit: int | None = None, - large_resource_limit: int | None = None, - enable_debug_info: bool = False, - print_generic_op_form: bool = False, - tree_printing_dir_path: str | None = None, - ) -> None: ... - def enable_verifier(self, enable: bool) -> None: ... - @staticmethod - def parse(pipeline: str, context: _ir.Context | None = None) -> PassManager: ... - def run(self, module: _ir._OperationBase) -> None: ... - @property - def _CAPIPtr(self) -> object: ... diff --git a/mlir/python/mlir/_mlir_libs/_mlir/py.typed b/mlir/python/mlir/_mlir_libs/_mlir/py.typed new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/mlir/python/requirements.txt b/mlir/python/requirements.txt index 1a0075e829aef..abe09259bb1e8 100644 --- a/mlir/python/requirements.txt +++ b/mlir/python/requirements.txt @@ -1,6 +1,7 @@ -nanobind>=2.4, <3.0 +nanobind>=2.9, <3.0 numpy>=1.19.5, <=2.1.2 pybind11>=2.10.0, <=2.13.6 PyYAML>=5.4.0, <=6.0.1 ml_dtypes>=0.1.0, <=0.6.0; python_version<"3.13" # provides several NumPy dtype extensions, including the bf16 -ml_dtypes>=0.5.0, <=0.6.0; python_version>="3.13" \ No newline at end of file +ml_dtypes>=0.5.0, <=0.6.0; python_version>="3.13" +typing_extensions>=4.12.2 diff --git a/mlir/test/Examples/standalone/test.toy b/mlir/test/Examples/standalone/test.toy index 8b6d9dd62b76d..a88c115ebf197 100644 --- a/mlir/test/Examples/standalone/test.toy +++ b/mlir/test/Examples/standalone/test.toy @@ -3,6 +3,7 @@ # RUN: -DCMAKE_CXX_COMPILER=%host_cxx -DCMAKE_C_COMPILER=%host_cc \ # RUN: -DLLVM_ENABLE_LIBCXX=%enable_libcxx -DMLIR_DIR=%mlir_cmake_dir \ # RUN: -DLLVM_USE_LINKER=%llvm_use_linker \ +# RUN: -DMLIR_PYTHON_PACKAGE_PREFIX=mlir_standalone \ # RUN: -DPython3_EXECUTABLE=%python \ # RUN: -DPython_EXECUTABLE=%python # RUN: "%cmake_exe" --build . --target check-standalone | tee %t diff --git a/mlir/test/python/lib/PythonTestModuleNanobind.cpp b/mlir/test/python/lib/PythonTestModuleNanobind.cpp index 108df8da86a7e..a497fcccf13d7 100644 --- a/mlir/test/python/lib/PythonTestModuleNanobind.cpp +++ b/mlir/test/python/lib/PythonTestModuleNanobind.cpp @@ -37,7 +37,10 @@ NB_MODULE(_mlirPythonTestNanobind, m) { mlirDialectHandleLoadDialect(pythonTestDialect, context); } }, - nb::arg("context"), nb::arg("load") = true); + nb::arg("context"), nb::arg("load") = true, + // clang-format off + nb::sig("def register_python_test_dialect(context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") ", load: bool = True) -> None")); + // clang-format on m.def( "register_dialect", @@ -46,14 +49,21 @@ NB_MODULE(_mlirPythonTestNanobind, m) { mlirGetDialectHandle__python_test__(); mlirDialectHandleInsertDialect(pythonTestDialect, registry); }, - nb::arg("registry")); + nb::arg("registry"), + // clang-format off + nb::sig("def register_dialect(registry: " MAKE_MLIR_PYTHON_QUALNAME("ir.DialectRegistry") ") -> None")); + // clang-format on - m.def("test_diagnostics_with_errors_and_notes", [](MlirContext ctx) { - mlir::python::CollectDiagnosticsToStringScope handler(ctx); - - mlirPythonTestEmitDiagnosticWithNote(ctx); - throw nb::value_error(handler.takeMessage().c_str()); - }); + m.def( + "test_diagnostics_with_errors_and_notes", + [](MlirContext ctx) { + mlir::python::CollectDiagnosticsToStringScope handler(ctx); + mlirPythonTestEmitDiagnosticWithNote(ctx); + throw nb::value_error(handler.takeMessage().c_str()); + }, + // clang-format off + nb::sig("def test_diagnostics_with_errors_and_notes(arg: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") ", /) -> None")); + // clang-format on mlir_attribute_subclass(m, "TestAttr", mlirAttributeIsAPythonTestTestAttribute, @@ -63,6 +73,9 @@ NB_MODULE(_mlirPythonTestNanobind, m) { [](const nb::object &cls, MlirContext ctx) { return cls(mlirPythonTestTestAttributeGet(ctx)); }, + // clang-format off + nb::sig("def get(cls: object, context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " | None = None) -> object"), + // clang-format on nb::arg("cls"), nb::arg("context").none() = nb::none()); mlir_type_subclass(m, "TestType", mlirTypeIsAPythonTestTestType, @@ -72,6 +85,9 @@ NB_MODULE(_mlirPythonTestNanobind, m) { [](const nb::object &cls, MlirContext ctx) { return cls(mlirPythonTestTestTypeGet(ctx)); }, + // clang-format off + nb::sig("def get(cls: object, context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " | None = None) -> object"), + // clang-format on nb::arg("cls"), nb::arg("context").none() = nb::none()); auto typeCls = @@ -88,6 +104,9 @@ NB_MODULE(_mlirPythonTestNanobind, m) { shape.size(), shape.data(), mlirIntegerTypeGet(ctx, width), encoding)); }, + // clang-format off + nb::sig("def get(cls: object, shape: collections.abc.Sequence[int], width: int, context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " | None = None) -> object"), + // clang-format on nb::arg("cls"), nb::arg("shape"), nb::arg("width"), nb::arg("context").none() = nb::none());