1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
|
//===- 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 <string>
#include <vector>
#include "PybindUtils.h"
#include "llvm/ADT/Optional.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/StringSet.h"
namespace mlir {
namespace python {
/// Globals that are always accessible once the extension has been initialized.
class PyGlobals {
public:
PyGlobals();
~PyGlobals();
/// Most code should get the globals via this static accessor.
static PyGlobals &get() {
assert(instance && "PyGlobals is null");
return *instance;
}
/// Get and set the list of parent modules to search for dialect
/// implementation classes.
std::vector<std::string> &getDialectSearchPrefixes() {
return dialectSearchPrefixes;
}
void setDialectSearchPrefixes(std::vector<std::string> newValues) {
dialectSearchPrefixes.swap(newValues);
}
/// Clears positive and negative caches regarding what implementations are
/// available. Future lookups will do more expensive existence checks.
void clearImportCache();
/// 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.
void loadDialectModule(llvm::StringRef dialectNamespace);
/// Decorator for registering a custom Dialect class. The class object must
/// have a DIALECT_NAMESPACE attribute.
pybind11::object registerDialectDecorator(pybind11::object pyClass);
/// Adds a concrete implementation dialect class.
/// Raises an exception if the mapping already exists.
/// This is intended to be called by implementation code.
void registerDialectImpl(const std::string &dialectNamespace,
pybind11::object pyClass);
/// Adds a concrete implementation operation class.
/// Raises an exception if the mapping already exists.
/// This is intended to be called by implementation code.
void registerOperationImpl(const std::string &operationName,
pybind11::object pyClass,
pybind11::object rawOpViewClass);
/// Looks up a registered dialect class by namespace. Note that this may
/// trigger loading of the defining module and can arbitrarily re-enter.
llvm::Optional<pybind11::object>
lookupDialectClass(const std::string &dialectNamespace);
/// Looks up a registered raw OpView class by operation name. Note that this
/// may trigger a load of the dialect, which can arbitrarily re-enter.
llvm::Optional<pybind11::object>
lookupRawOpViewClass(llvm::StringRef operationName);
private:
static PyGlobals *instance;
/// Module name prefixes to search under for dialect implementation modules.
std::vector<std::string> dialectSearchPrefixes;
/// Map of dialect namespace to external dialect class object.
llvm::StringMap<pybind11::object> dialectClassMap;
/// Map of full operation name to external operation class object.
llvm::StringMap<pybind11::object> operationClassMap;
/// Map of operation name to custom subclass that directly initializes
/// the OpView base class (bypassing the user class constructor).
llvm::StringMap<pybind11::object> rawOpViewClassMap;
/// Set of dialect namespaces that we have attempted to import implementation
/// modules for.
llvm::StringSet<> loadedDialectModulesCache;
/// Cache of operation name to custom OpView subclass that directly
/// initializes the OpView base class (or an undefined object for negative
/// lookup). This is maintained on loopup as a shadow of rawOpViewClassMap
/// in order for repeat lookups of the OpView classes to only incur the cost
/// of one hashtable lookup.
llvm::StringMap<pybind11::object> rawOpViewClassMapCache;
};
} // namespace python
} // namespace mlir
#endif // MLIR_BINDINGS_PYTHON_GLOBALS_H
|
