Expose list of PassContext configurations to the Python APIs (apache#…

…8212)

* Expose C++ PassContext::ListAllConfigs via its Python counterpart
   tvm.ir.transform.PassContext.list_configs()

 * Add unit tests for the C++ and Python layers

include/tvm/ir/transform.h

@@ -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

python/tvm/ir/transform.py

@@ -120,6 +120,11 @@ def current():
         """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):

src/ir/transform.cc

@@ -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

tests/cpp/relay_transform_sequential_test.cc

@@ -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";

tests/python/relay/test_pass_instrument.py

@@ -182,6 +182,13 @@ def run_after_pass(self, mod, info):
     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 = []