Add tests for supply

Author: Florin-Gabriel Blanaru

python/tvm/ir/supply.py

@@ -62,6 +62,8 @@ def __init__(self, value=None):
         if value is None:
             name_supply = NameSupply("")
             self.__init_handle_by_constructor__(_ffi_api.GlobalVarSupply_NameSupply, name_supply)
+        elif isinstance(value, NameSupply):
+            self.__init_handle_by_constructor__(_ffi_api.GlobalVarSupply_NameSupply, value)
         elif isinstance(value, (list, tvm.container.Array)):
             self.__init_handle_by_constructor__(_ffi_api.GlobalVarSupply_IRModules, value)
         elif isinstance(value, IRModule):

src/ir/name_supply.cc

@@ -40,7 +40,7 @@ String NameSupplyNode::ReserveName(const String& name, bool add_prefix) {
   if (add_prefix) {
     final_name = prefix_module_name(name);
   }
-  name_map[final_name] = 1;
+  name_map[final_name] = 0;
   return final_name;
 }
 

tests/cpp/name_supply_test.cc

@@ -0,0 +1,110 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <tvm/ir/global_var_supply.h>
+#include <tvm/ir/module.h>
+#include <tvm/ir/name_supply.h>
+#include <tvm/relay/expr.h>
+#include <tvm/relay/function.h>
+
+using namespace tvm;
+
+NameSupply preambleNameSupply() {
+  NameSupply name_supply = NameSupply("prefix");
+  name_supply->FreshName("test");
+  return name_supply;
+}
+
+TEST(NameSupply, FreshName) {
+  NameSupply name_supply = preambleNameSupply();
+  String fresh = name_supply->FreshName("test");
+
+  EXPECT_EQ(fresh.compare("prefix_test1"), 0);
+}
+
+TEST(NameSupply, ContainsName) {
+  NameSupply name_supply = preambleNameSupply();
+
+  EXPECT_TRUE(name_supply->ContainsName("test"));
+  EXPECT_FALSE(name_supply->ContainsName("test1"));
+}
+
+TEST(NameSupply, ReserveName) {
+  NameSupply name_supply = preambleNameSupply();
+  name_supply->ReserveName("otherTest", false);
+
+  EXPECT_TRUE(name_supply->ContainsName("otherTest", false));
+  EXPECT_FALSE(name_supply->ContainsName("otherTest"));
+}
+
+GlobalVarSupply preambleVarSupply() {
+  GlobalVarSupply global_var_supply = GlobalVarSupply();
+  global_var_supply->FreshGlobal("test");
+  return global_var_supply;
+}
+
+TEST(GlobalVarSupply, FreshGlobal) {
+  GlobalVarSupply global_var_supply = preambleVarSupply();
+  GlobalVar first_var = global_var_supply->FreshGlobal("test");
+  GlobalVar second_var = global_var_supply->FreshGlobal("test");
+
+  EXPECT_FALSE(tvm::StructuralEqual()(first_var, second_var));
+  EXPECT_EQ(first_var->name_hint.compare("test1"), 0);
+  EXPECT_EQ(second_var->name_hint.compare("test2"), 0);
+}
+
+TEST(GlobalVarSupply, UniqueGlobalFor) {
+  GlobalVarSupply global_var_supply = preambleVarSupply();
+  GlobalVar first_var = global_var_supply->UniqueGlobalFor("someName");
+  GlobalVar second_var = global_var_supply->UniqueGlobalFor("someName");
+
+  EXPECT_TRUE(tvm::StructuralEqual()(first_var, second_var));
+  EXPECT_EQ(first_var->name_hint.compare("someName"), 0);
+  EXPECT_EQ(second_var->name_hint.compare("someName"), 0);
+}
+
+TEST(GlobalVarSupply, ReserveGlobal) {
+  GlobalVarSupply global_var_supply = preambleVarSupply();
+  GlobalVar var = GlobalVar("someName");
+  global_var_supply->ReserveGlobalVar(var);
+  GlobalVar second_var = global_var_supply->UniqueGlobalFor("someName");
+  GlobalVar third_var = global_var_supply->FreshGlobal("someName");
+
+  EXPECT_TRUE(tvm::StructuralEqual()(var, second_var));
+  EXPECT_FALSE(tvm::StructuralEqual()(var, third_var));
+  EXPECT_EQ(second_var->name_hint.compare("someName"), 0);
+  EXPECT_EQ(third_var->name_hint.compare("someName1"), 0);
+}
+
+TEST(GlobalVarSupply, BuildIRModule) {
+  auto x = relay::Var("x", relay::Type());
+  auto f = relay::Function(tvm::Array<relay::Var>{x}, x, relay::Type(), {});
+  GlobalVar var = GlobalVar("test");
+  IRModule module = IRModule({{var, f}});
+
+  GlobalVarSupply global_var_supply = GlobalVarSupply(module);
+  GlobalVar second_var = global_var_supply->UniqueGlobalFor("test", false);
+  GlobalVar third_var = global_var_supply->FreshGlobal("test", false);
+
+  EXPECT_TRUE(tvm::StructuralEqual()(var, second_var));
+  EXPECT_FALSE(tvm::StructuralEqual()(var, third_var));
+  EXPECT_EQ(second_var->name_hint.compare("test"), 0);
+  EXPECT_EQ(third_var->name_hint.compare("test1"), 0);
+}

tests/python/relay/test_name_supply.py

@@ -0,0 +1,72 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import tvm
+import tvm.testing
+
+from tvm import relay
+from tvm.ir import GlobalVar, structural_equal
+from tvm.ir.supply import NameSupply
+from tvm.ir.supply import GlobalVarSupply
+
+
+def test_name_supply():
+    name_supply = NameSupply("prefix")
+    name_supply.reserve_name("test")
+
+    assert name_supply.contains_name("test")
+    assert name_supply.fresh_name("test") == "prefix_test1"
+    assert name_supply.contains_name("test1")
+    assert not name_supply.contains_name("test1", False)
+    assert not name_supply.contains_name("test2")
+
+
+def test_global_var_supply_from_none():
+    var_supply = GlobalVarSupply()
+    global_var = GlobalVar("test")
+    var_supply.reserve_global(global_var)
+
+    assert structural_equal(var_supply.unique_global_for("test"), global_var)
+    assert not structural_equal(var_supply.fresh_global("test"), global_var)
+
+
+def test_global_var_supply_from_name_supply():
+    name_supply = NameSupply("prefix")
+    var_supply = GlobalVarSupply(name_supply)
+    global_var = GlobalVar("test")
+    var_supply.reserve_global(global_var)
+
+    assert structural_equal(var_supply.unique_global_for("test", False), global_var)
+    assert not structural_equal(var_supply.unique_global_for("test"), global_var)
+
+
+def test_global_var_supply_from_ir_mod():
+    x = relay.var("x")
+    y = relay.var("y")
+    mod = tvm.IRModule()
+    global_var = GlobalVar("test")
+    mod[global_var] = relay.Function([x, y], relay.add(x, y))
+    var_supply = GlobalVarSupply(mod)
+
+    second_global_var = var_supply.fresh_global("test", False)
+
+    assert structural_equal(var_supply.unique_global_for("test", False), global_var)
+    assert not structural_equal(var_supply.unique_global_for("test"), global_var)
+    assert not structural_equal(second_global_var, global_var)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()