llvm-project/mlir/include/mlir/Bindings/Python/Globals.h

//===- Globals.h - MLIR Python extension globals --------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_BINDINGS_PYTHON_GLOBALS_H
#define MLIR_BINDINGS_PYTHON_GLOBALS_H

#include <optional>
#include <regex>
#include <string>
#include <string_view>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include "mlir-c/IR.h"
#include "mlir-c/Support.h"
#include "mlir/Bindings/Python/NanobindUtils.h"

namespace mlir {
namespace python {
namespace MLIR_BINDINGS_PYTHON_DOMAIN {
/// Globals that are always accessible once the extension has been initialized.
/// Methods of this class are thread-safe.
class MLIR_PYTHON_API_EXPORTED PyGlobals {
public:
  PyGlobals();
  ~PyGlobals();

  /// Most code should get the globals via this static accessor.
  static PyGlobals &get();

  /// Get and set the list of parent modules to search for dialect
  /// implementation classes.
  std::vector<std::string> getDialectSearchPrefixes() {
    nanobind::ft_lock_guard lock(mutex);
    return dialectSearchPrefixes;
  }
  void setDialectSearchPrefixes(std::vector<std::string> newValues) {
    nanobind::ft_lock_guard lock(mutex);
    dialectSearchPrefixes.swap(newValues);
  }
  void addDialectSearchPrefix(std::string value) {
    nanobind::ft_lock_guard lock(mutex);
    dialectSearchPrefixes.push_back(std::move(value));
  }

  /// Loads a python module corresponding to the given dialect namespace.
  /// No-ops if the module has already been loaded or is not found. Raises
  /// an error on any evaluation issues.
  /// Note that this returns void because it is expected that the module
  /// contains calls to decorators and helpers that register the salient
  /// entities. Returns true if dialect is successfully loaded.
  bool loadDialectModule(std::string_view dialectNamespace);

  /// Adds a user-friendly Attribute builder.
  /// Raises an exception if the mapping already exists and replace == false
  /// and allow_existing == false.
  /// Silently skips registration if allow_existing == true and the mapping
  /// already exists (first registration wins).
  /// This is intended to be called by implementation code.
  void registerAttributeBuilder(const std::string &attributeKind,
                                nanobind::callable pyFunc, bool replace = false,
                                bool allow_existing = false);

  /// Adds a user-friendly type caster. Raises an exception if the mapping
  /// already exists and replace == false. This is intended to be called by
  /// implementation code.
  void registerTypeCaster(MlirTypeID mlirTypeID, nanobind::callable typeCaster,
                          bool replace = false);

  /// Adds a user-friendly value caster. Raises an exception if the mapping
  /// already exists and replace == false. This is intended to be called by
  /// implementation code.
  void registerValueCaster(MlirTypeID mlirTypeID,
                           nanobind::callable valueCaster,
                           bool replace = false);

  /// Adds a concrete implementation dialect class.
  /// Raises an exception if the mapping already exists and replace == false.
  /// This is intended to be called by implementation code.
  void registerDialectImpl(const std::string &dialectNamespace,
                           nanobind::object pyClass, bool replace = false);

  /// Adds a concrete implementation operation class.
  /// Raises an exception if the mapping already exists and replace == false.
  /// This is intended to be called by implementation code.
  void registerOperationImpl(const std::string &operationName,
                             nanobind::object pyClass, bool replace = false);

  /// Adds an operation adaptor class.
  /// Raises an exception if the mapping already exists and replace == false.
  /// This is intended to be called by implementation code.
  void registerOpAdaptorImpl(const std::string &operationName,
                             nanobind::object pyClass, bool replace = false);

  /// Returns the custom Attribute builder for Attribute kind.
  std::optional<nanobind::callable>
  lookupAttributeBuilder(const std::string &attributeKind);

  /// Returns the custom type caster for MlirTypeID mlirTypeID.
  std::optional<nanobind::callable> lookupTypeCaster(MlirTypeID mlirTypeID,
                                                     MlirDialect dialect);

  /// Returns the custom value caster for MlirTypeID mlirTypeID.
  std::optional<nanobind::callable> lookupValueCaster(MlirTypeID mlirTypeID,
                                                      MlirDialect dialect);

  /// Looks up a registered dialect class by namespace. Note that this may
  /// trigger loading of the defining module and can arbitrarily re-enter.
  std::optional<nanobind::object>
  lookupDialectClass(const std::string &dialectNamespace);

  /// Looks up a registered operation class (deriving from OpView) by operation
  /// name. Note that this may trigger a load of the dialect, which can
  /// arbitrarily re-enter.
  std::optional<nanobind::object>
  lookupOperationClass(std::string_view operationName);

  /// Looks up a registered operation adaptor class by operation
  /// name. Note that this may trigger a load of the dialect, which can
  /// arbitrarily re-enter.
  std::optional<nanobind::object>
  lookupOpAdaptorClass(std::string_view operationName);

  class MLIR_PYTHON_API_EXPORTED TracebackLoc {
  public:
    bool locTracebacksEnabled();

    void setLocTracebacksEnabled(bool value);

    size_t locTracebackFramesLimit();

    void setLocTracebackFramesLimit(size_t value);

    void registerTracebackFileInclusion(const std::string &file);

    void registerTracebackFileExclusion(const std::string &file);

    bool isUserTracebackFilename(std::string_view file);

    static constexpr size_t kMaxFrames = 512;

  private:
    nanobind::ft_mutex mutex;
    bool locTracebackEnabled_ = false;
    size_t locTracebackFramesLimit_ = 10;
    std::unordered_set<std::string> userTracebackIncludeFiles;
    std::unordered_set<std::string> userTracebackExcludeFiles;
    std::regex userTracebackIncludeRegex;
    bool rebuildUserTracebackIncludeRegex = false;
    std::regex userTracebackExcludeRegex;
    bool rebuildUserTracebackExcludeRegex = false;
    std::unordered_map<std::string, bool> isUserTracebackFilenameCache;
  };

  TracebackLoc &getTracebackLoc() { return tracebackLoc; }

  class TypeIDAllocator {
  public:
    TypeIDAllocator() : allocator(mlirTypeIDAllocatorCreate()) {}
    ~TypeIDAllocator() {
      if (allocator.ptr)
        mlirTypeIDAllocatorDestroy(allocator);
    }
    TypeIDAllocator(const TypeIDAllocator &) = delete;
    TypeIDAllocator(TypeIDAllocator &&other) : allocator(other.allocator) {
      other.allocator.ptr = nullptr;
    }

    MlirTypeIDAllocator get() { return allocator; }
    MlirTypeID allocate() {
      return mlirTypeIDAllocatorAllocateTypeID(allocator);
    }

  private:
    MlirTypeIDAllocator allocator;
  };

  MlirTypeID allocateTypeID() { return typeIDAllocator.allocate(); }

private:
  static PyGlobals *instance;

  nanobind::ft_mutex mutex;

  /// Module name prefixes to search under for dialect implementation modules.
  std::vector<std::string> dialectSearchPrefixes;
  /// Map of dialect namespace to external dialect class object.
  std::unordered_map<std::string, nanobind::object> dialectClassMap;
  /// Map of full operation name to external operation class object.
  std::unordered_map<std::string, nanobind::object> operationClassMap;
  /// Map of full operation name to external operation adaptor class object.
  std::unordered_map<std::string, nanobind::object> opAdaptorClassMap;
  /// Map of attribute ODS name to custom builder.
  std::unordered_map<std::string, nanobind::callable> attributeBuilderMap;
  /// Map of MlirTypeID to custom type caster.
  std::unordered_map<MlirTypeID, nanobind::callable, MlirTypeIDHash,
                     MlirTypeIDEqual>
      typeCasterMap;
  /// Map of MlirTypeID to custom value caster.
  std::unordered_map<MlirTypeID, nanobind::callable, MlirTypeIDHash,
                     MlirTypeIDEqual>
      valueCasterMap;
  /// Set of dialect namespaces that we have attempted to import implementation
  /// modules for.
  std::unordered_set<std::string> loadedDialectModules;

  TracebackLoc tracebackLoc;
  TypeIDAllocator typeIDAllocator;
};
} // namespace MLIR_BINDINGS_PYTHON_DOMAIN
} // namespace python
} // namespace mlir

#endif // MLIR_BINDINGS_PYTHON_GLOBALS_H