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:
Leandro Nunes
2021-06-09 20:07:45 +01:00
committed by GitHub
parent 53e4c603e6
commit 1f2ca068c7
5 changed files with 39 additions and 0 deletions
+6
View File
@@ -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
+5
View File
@@ -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):
+14
View File
@@ -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 = []