1# Licensed under the Apache License v2.0 with LLVM Exceptions.
2# See https://llvm.org/LICENSE.txt for license information.
3# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
4
5from typing import Any, Sequence
6
7import os
8
9_this_dir = os.path.dirname(__file__)
10
11
12def get_lib_dirs() -> Sequence[str]:
13  """Gets the lib directory for linking to shared libraries.
14
15  On some platforms, the package may need to be built specially to export
16  development libraries.
17  """
18  return [_this_dir]
19
20
21def get_include_dirs() -> Sequence[str]:
22  """Gets the include directory for compiling against exported C libraries.
23
24  Depending on how the package was build, development C libraries may or may
25  not be present.
26  """
27  return [os.path.join(_this_dir, "include")]
28
29
30# Perform Python level site initialization. This involves:
31#   1. Attempting to load initializer modules, specific to the distribution.
32#   2. Defining the concrete mlir.ir.Context that does site specific
33#      initialization.
34#
35# Aside from just being far more convenient to do this at the Python level,
36# it is actually quite hard/impossible to have such __init__ hooks, given
37# the pybind memory model (i.e. there is not a Python reference to the object
38# in the scope of the base class __init__).
39#
40# For #1, we:
41#   a. Probe for modules named '_mlirRegisterEverything' and
42#     '_site_initialize_{i}', where 'i' is a number starting at zero and
43#     proceeding so long as a module with the name is found.
44#   b. If the module has a 'register_dialects' attribute, it will be called
45#     immediately with a DialectRegistry to populate.
46#   c. If the module has a 'context_init_hook', it will be added to a list
47#     of callbacks that are invoked as the last step of Context
48#     initialization (and passed the Context under construction).
49#
50# This facility allows downstreams to customize Context creation to their
51# needs.
52def _site_initialize():
53  import importlib
54  import itertools
55  import logging
56  from ._mlir import ir
57  registry = ir.DialectRegistry()
58  post_init_hooks = []
59
60  def process_initializer_module(module_name):
61    try:
62      m = importlib.import_module(f".{module_name}", __name__)
63    except ModuleNotFoundError:
64      return False
65
66    logging.debug("Initializing MLIR with module: %s", module_name)
67    if hasattr(m, "register_dialects"):
68      logging.debug("Registering dialects from initializer %r", m)
69      m.register_dialects(registry)
70    if hasattr(m, "context_init_hook"):
71      logging.debug("Adding context init hook from %r", m)
72      post_init_hooks.append(m.context_init_hook)
73    return True
74
75
76  # If _mlirRegisterEverything is built, then include it as an initializer
77  # module.
78  process_initializer_module("_mlirRegisterEverything")
79
80  # Load all _site_initialize_{i} modules, where 'i' is a number starting
81  # at 0.
82  for i in itertools.count():
83    module_name = f"_site_initialize_{i}"
84    if not process_initializer_module(module_name):
85      break
86
87  class Context(ir._BaseContext):
88    def __init__(self, *args, **kwargs):
89      super().__init__(*args, **kwargs)
90      self.append_dialect_registry(registry)
91      for hook in post_init_hooks:
92        hook(self)
93      # TODO: There is some debate about whether we should eagerly load
94      # all dialects. It is being done here in order to preserve existing
95      # behavior. See: https://github.com/llvm/llvm-project/issues/56037
96      self.load_all_available_dialects()
97
98  ir.Context = Context
99
100
101_site_initialize()
102