# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from typing import Any, Sequence
import os
_this_dir = os.path.dirname(__file__)
def get_lib_dirs() -> Sequence[str]:
"""Gets the lib directory for linking to shared libraries.
On some platforms, the package may need to be built specially to export
development libraries.
"""
return [_this_dir]
def get_include_dirs() -> Sequence[str]:
"""Gets the include directory for compiling against exported C libraries.
Depending on how the package was build, development C libraries may or may
not be present.
"""
return [os.path.join(_this_dir, "include")]
# Perform Python level site initialization. This involves:
# 1. Attempting to load initializer modules, specific to the distribution.
# 2. Defining the concrete mlir.ir.Context that does site specific
# initialization.
#
# Aside from just being far more convenient to do this at the Python level,
# it is actually quite hard/impossible to have such __init__ hooks, given
# the pybind memory model (i.e. there is not a Python reference to the object
# in the scope of the base class __init__).
#
# For #1, we:
# a. Probe for modules named '_mlirRegisterEverything' and
# '_site_initialize_{i}', where 'i' is a number starting at zero and
# proceeding so long as a module with the name is found.
# b. If the module has a 'register_dialects' attribute, it will be called
# immediately with a DialectRegistry to populate.
# c. If the module has a 'context_init_hook', it will be added to a list
# of callbacks that are invoked as the last step of Context
# initialization (and passed the Context under construction).
# d. If the module has a 'disable_multithreading' attribute, it will be
# taken as a boolean. If it is True for any initializer, then the
# default behavior of enabling multithreading on the context
# will be suppressed. This complies with the original behavior of all
# contexts being created with multithreading enabled while allowing
# this behavior to be changed if needed (i.e. if a context_init_hook
# explicitly sets up multithreading).
#
# This facility allows downstreams to customize Context creation to their
# needs.
_dialect_registry = None
def get_dialect_registry():
global _dialect_registry
if _dialect_registry is None:
from ._mlir import ir
_dialect_registry = ir.DialectRegistry()
return _dialect_registry
def _site_initialize():
import importlib
import itertools
import logging
from ._mlir import ir
logger = logging.getLogger(__name__)
post_init_hooks = []
disable_multithreading = False
def process_initializer_module(module_name):
nonlocal disable_multithreading
try:
m = importlib.import_module(f".{module_name}", __name__)
except ModuleNotFoundError:
return False
except ImportError:
message = (
f"Error importing mlir initializer {module_name}. This may "
"happen in unclean incremental builds but is likely a real bug if "
"encountered otherwise and the MLIR Python API may not function."
)
logger.warning(message, exc_info=True)
return False
logger.debug("Initializing MLIR with module: %s", module_name)
if hasattr(m, "register_dialects"):
logger.debug("Registering dialects from initializer %r", m)
m.register_dialects(get_dialect_registry())
if hasattr(m, "context_init_hook"):
logger.debug("Adding context init hook from %r", m)
post_init_hooks.append(m.context_init_hook)
if hasattr(m, "disable_multithreading"):
if bool(m.disable_multithreading):
logger.debug("Disabling multi-threading for context")
disable_multithreading = True
return True
# If _mlirRegisterEverything is built, then include it as an initializer
# module.
init_module = None
if process_initializer_module("_mlirRegisterEverything"):
init_module = importlib.import_module(f"._mlirRegisterEverything", __name__)
# Load all _site_initialize_{i} modules, where 'i' is a number starting
# at 0.
for i in itertools.count():
module_name = f"_site_initialize_{i}"
if not process_initializer_module(module_name):
break
class Context(ir._BaseContext):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.append_dialect_registry(get_dialect_registry())
for hook in post_init_hooks:
hook(self)
if not disable_multithreading:
self.enable_multithreading(True)
# TODO: There is some debate about whether we should eagerly load
# all dialects. It is being done here in order to preserve existing
# behavior. See: https://github.com/llvm/llvm-project/issues/56037
self.load_all_available_dialects()
if init_module:
logger.debug(
"Registering translations from initializer %r", init_module
)
init_module.register_llvm_translations(self)
ir.Context = Context
class MLIRError(Exception):
"""
An exception with diagnostic information. Has the following fields:
message: str
error_diagnostics: List[ir.DiagnosticInfo]
"""
def __init__(self, message, error_diagnostics):
self.message = message
self.error_diagnostics = error_diagnostics
super().__init__(message, error_diagnostics)
def __str__(self):
s = self.message
if self.error_diagnostics:
s += ":"
for diag in self.error_diagnostics:
s += (
"\nerror: "
+ str(diag.location)[4:-1]
+ ": "
+ diag.message.replace("\n", "\n ")
)
for note in diag.notes:
s += (
"\n note: "
+ str(note.location)[4:-1]
+ ": "
+ note.message.replace("\n", "\n ")
)
return s
ir.MLIRError = MLIRError
_site_initialize()
Codebase Browser
llvm