Expose list of PassContext configurations to the Python APIs (#8212)
* Expose C++ PassContext::ListAllConfigs via its Python counterpart tvm.ir.transform.PassContext.list_configs() * Add unit tests for the C++ and Python layers
This commit is contained in:
@@ -183,6 +183,12 @@ class PassContext : public ObjectRef {
|
||||
*/
|
||||
TVM_DLL static PassContext Current();
|
||||
|
||||
/*!
|
||||
* \brief Get all supported configuration names, registered within the PassContext.
|
||||
* \return List of all configuration names.
|
||||
*/
|
||||
TVM_DLL static Array<String> ListConfigNames();
|
||||
|
||||
/*!
|
||||
* \brief Call instrument implementations' callbacks when entering PassContext.
|
||||
* The callbacks are called in order, and if one raises an exception, the rest will not be
|
||||
|
||||
@@ -120,6 +120,11 @@ class PassContext(tvm.runtime.Object):
|
||||
"""Return the current pass context."""
|
||||
return _ffi_transform_api.GetCurrentPassContext()
|
||||
|
||||
@staticmethod
|
||||
def list_config_names():
|
||||
"""List all registered `PassContext` configuration names"""
|
||||
return list(_ffi_transform_api.ListConfigNames())
|
||||
|
||||
|
||||
@tvm._ffi.register_object("transform.Pass")
|
||||
class Pass(tvm.runtime.Object):
|
||||
|
||||
@@ -145,6 +145,14 @@ class PassConfigManager {
|
||||
}
|
||||
}
|
||||
|
||||
Array<String> ListConfigNames() {
|
||||
Array<String> config_keys;
|
||||
for (const auto& kv : key2vtype_) {
|
||||
config_keys.push_back(kv.first);
|
||||
}
|
||||
return config_keys;
|
||||
}
|
||||
|
||||
static PassConfigManager* Global() {
|
||||
static auto* inst = new PassConfigManager();
|
||||
return inst;
|
||||
@@ -163,6 +171,10 @@ void PassContext::RegisterConfigOption(const char* key, uint32_t value_type_inde
|
||||
PassConfigManager::Global()->Register(key, value_type_index);
|
||||
}
|
||||
|
||||
Array<String> PassContext::ListConfigNames() {
|
||||
return PassConfigManager::Global()->ListConfigNames();
|
||||
}
|
||||
|
||||
PassContext PassContext::Create() { return PassContext(make_object<PassContextNode>()); }
|
||||
|
||||
void PassContext::InstrumentEnterPassContext() {
|
||||
@@ -607,5 +619,7 @@ Pass PrintIR(String header, bool show_meta_data) {
|
||||
|
||||
TVM_REGISTER_GLOBAL("transform.PrintIR").set_body_typed(PrintIR);
|
||||
|
||||
TVM_REGISTER_GLOBAL("transform.ListConfigNames").set_body_typed(PassContext::ListConfigNames);
|
||||
|
||||
} // namespace transform
|
||||
} // namespace tvm
|
||||
|
||||
@@ -121,6 +121,13 @@ TEST(Relay, Sequential) {
|
||||
ICHECK(tvm::StructuralEqual()(f, expected));
|
||||
}
|
||||
|
||||
TEST(PassContextListConfigNames, Basic) {
|
||||
Array<String> configs = relay::transform::PassContext::ListConfigNames();
|
||||
ICHECK_EQ(configs.empty(), false);
|
||||
ICHECK_EQ(std::count(std::begin(configs), std::end(configs), "relay.backend.use_auto_scheduler"),
|
||||
1);
|
||||
}
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
testing::InitGoogleTest(&argc, argv);
|
||||
testing::FLAGS_gtest_death_test_style = "threadsafe";
|
||||
|
||||
@@ -182,6 +182,13 @@ def test_instrument_pass_counts():
|
||||
assert passes_counter.run_after_count == 0
|
||||
|
||||
|
||||
def test_list_pass_configs():
|
||||
config_names = tvm.transform.PassContext.list_config_names()
|
||||
|
||||
assert len(config_names) > 0
|
||||
assert "relay.backend.use_auto_scheduler" in config_names
|
||||
|
||||
|
||||
def test_enter_pass_ctx_exception():
|
||||
events = []
|
||||
|
||||
|
||||
Reference in New Issue
Block a user