kparzysz-quic commented on code in PR #12066:
URL: https://github.com/apache/tvm/pull/12066#discussion_r931252691


##########
src/ir/module.cc:
##########
@@ -40,6 +40,8 @@
 #include <sstream>
 #include <unordered_set>
 
+#include "tvm/ir/global_var_supply.h"

Review Comment:
   This should be <>, please move it to where the other includes from tvm/ir 
are.



##########
src/ir/global_var_supply.cc:
##########
@@ -0,0 +1,115 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file global_var_supply.cc
+ * \brief GlobalVarSupply that can be used to generate unique GlobalVars.
+ */
+#include "tvm/ir/global_var_supply.h"
+
+#include <tvm/runtime/registry.h>
+
+#include <utility>
+
+#include "tvm/ir/expr.h"
+
+namespace tvm {
+GlobalVarSupply::GlobalVarSupply(const NameSupply& name_supply,
+                                 std::unordered_map<std::string, GlobalVar> 
name_to_var_map) {
+  auto n = make_object<GlobalVarSupplyNode>(name_supply);
+  n->name_to_var_map_ = std::move(name_to_var_map);
+  data_ = std::move(n);
+}
+
+std::string GetModuleName(const IRModule& module) {
+  return 
module->GetAttr<String>(tvm::attr::kModuleName).value_or("tvmgen_default");
+}
+
+GlobalVarSupply::GlobalVarSupply(const Array<IRModule>& modules) : 
GlobalVarSupply(NameSupply("")) {
+  if (!modules.empty()) {
+    IRModule first_mod = modules.front();
+    this->operator->()->name_supply_->prefix_ = GetModuleName(first_mod);
+  }
+  for (auto& mod : modules) {
+    for (auto kv : mod->functions) {
+      this->operator->()->ReserveGlobalVar(kv.first);
+    }
+  }
+}
+
+GlobalVarSupply::GlobalVarSupply(const IRModule module)
+    : GlobalVarSupply(Array<IRModule>{module}) {}
+
+void GlobalVarSupplyNode::ReserveGlobalVar(const GlobalVar& var, bool 
allow_conflict) {
+  name_supply_->ReserveName(var->name_hint, false);
+  if (!allow_conflict) {
+    ICHECK(name_to_var_map_.count(var->name_hint) == 0)
+        << "GlobalVar " << var << " conflicts by name in this supply.";
+  }
+  name_to_var_map_[var->name_hint] = var;
+}
+
+GlobalVarSupplyNode::GlobalVarSupplyNode(NameSupply name_supply)
+    : name_supply_(std::move(name_supply)) {}
+
+GlobalVar GlobalVarSupplyNode::UniqueGlobalFor(const String& name, bool 
add_prefix) {
+  String final_name = name_supply_->ReserveName(name, add_prefix);
+
+  auto it = name_to_var_map_.find(final_name);
+  if (it != name_to_var_map_.end()) {
+    return it->second;
+  } else {
+    GlobalVar var = GlobalVar(final_name);
+    name_to_var_map_[final_name] = var;

Review Comment:
   Do `emplace` instead.



##########
src/ir/global_var_supply.cc:
##########
@@ -0,0 +1,115 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file global_var_supply.cc
+ * \brief GlobalVarSupply that can be used to generate unique GlobalVars.
+ */
+#include "tvm/ir/global_var_supply.h"
+
+#include <tvm/runtime/registry.h>
+
+#include <utility>
+
+#include "tvm/ir/expr.h"
+
+namespace tvm {
+GlobalVarSupply::GlobalVarSupply(const NameSupply& name_supply,
+                                 std::unordered_map<std::string, GlobalVar> 
name_to_var_map) {
+  auto n = make_object<GlobalVarSupplyNode>(name_supply);
+  n->name_to_var_map_ = std::move(name_to_var_map);
+  data_ = std::move(n);
+}
+
+std::string GetModuleName(const IRModule& module) {
+  return 
module->GetAttr<String>(tvm::attr::kModuleName).value_or("tvmgen_default");
+}
+
+GlobalVarSupply::GlobalVarSupply(const Array<IRModule>& modules) : 
GlobalVarSupply(NameSupply("")) {
+  if (!modules.empty()) {
+    IRModule first_mod = modules.front();
+    this->operator->()->name_supply_->prefix_ = GetModuleName(first_mod);
+  }
+  for (auto& mod : modules) {
+    for (auto kv : mod->functions) {
+      this->operator->()->ReserveGlobalVar(kv.first);
+    }
+  }
+}
+
+GlobalVarSupply::GlobalVarSupply(const IRModule module)
+    : GlobalVarSupply(Array<IRModule>{module}) {}
+
+void GlobalVarSupplyNode::ReserveGlobalVar(const GlobalVar& var, bool 
allow_conflict) {
+  name_supply_->ReserveName(var->name_hint, false);
+  if (!allow_conflict) {
+    ICHECK(name_to_var_map_.count(var->name_hint) == 0)
+        << "GlobalVar " << var << " conflicts by name in this supply.";
+  }
+  name_to_var_map_[var->name_hint] = var;
+}
+
+GlobalVarSupplyNode::GlobalVarSupplyNode(NameSupply name_supply)
+    : name_supply_(std::move(name_supply)) {}
+
+GlobalVar GlobalVarSupplyNode::UniqueGlobalFor(const String& name, bool 
add_prefix) {
+  String final_name = name_supply_->ReserveName(name, add_prefix);
+
+  auto it = name_to_var_map_.find(final_name);
+  if (it != name_to_var_map_.end()) {
+    return it->second;
+  } else {
+    GlobalVar var = GlobalVar(final_name);
+    name_to_var_map_[final_name] = var;
+    return var;
+  }
+}
+
+GlobalVar GlobalVarSupplyNode::FreshGlobal(String name, bool add_prefix) {
+  String final_name = name_supply_->FreshName(name, add_prefix);
+  ICHECK(name_to_var_map_.find(final_name) == name_to_var_map_.end())
+      << "GlobalVar already exists for name " << final_name;
+  GlobalVar var = GlobalVar(final_name);
+  name_to_var_map_[final_name] = var;

Review Comment:
   Same here.



##########
src/ir/name_supply.cc:
##########
@@ -0,0 +1,112 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file name_supply.cc
+ * \brief NameSupply that can be used to generate unique variable names.
+ */
+#include "tvm/ir/name_supply.h"
+
+#include <tvm/runtime/registry.h>
+
+#include <utility>
+
+namespace tvm {
+
+NameSupply::NameSupply(const String& prefix, std::unordered_map<std::string, 
int> name_map) {
+  auto n = make_object<NameSupplyNode>();
+  n->prefix_ = prefix;
+  n->name_map = std::move(name_map);
+  data_ = std::move(n);
+}
+
+String NameSupplyNode::ReserveName(const String& name, bool add_prefix) {
+  String final_name = name;
+  if (add_prefix) {
+    final_name = add_prefix_to_name(name);
+  }
+  name_map[final_name] = 0;
+  return final_name;
+}
+
+String NameSupplyNode::FreshName(const String& name, bool add_prefix) {
+  String unique_name = name;
+  if (add_prefix) {
+    unique_name = add_prefix_to_name(name);
+  }
+  unique_name = GetUniqueName(unique_name);
+  return unique_name;
+}
+
+bool NameSupplyNode::ContainsName(const String& name, bool add_prefix) {
+  String unique_name = name;
+  if (add_prefix) {
+    unique_name = add_prefix_to_name(name);
+  }
+
+  return name_map.count(unique_name);
+}
+
+String NameSupplyNode::add_prefix_to_name(const String& name) {
+  if (prefix_.empty()) {
+    return name;
+  }
+
+  std::ostringstream ss;
+  ICHECK(name.defined());
+  ss << prefix_ << "_" << name;
+  return ss.str();
+}
+
+std::string NameSupplyNode::GetUniqueName(std::string prefix) {

Review Comment:
   "Prefix" already has a different meaning in the context of name supply, so 
please use a different name here.



##########
src/ir/name_supply.cc:
##########
@@ -0,0 +1,112 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file name_supply.cc
+ * \brief NameSupply that can be used to generate unique variable names.
+ */
+#include "tvm/ir/name_supply.h"
+
+#include <tvm/runtime/registry.h>
+
+#include <utility>
+
+namespace tvm {
+
+NameSupply::NameSupply(const String& prefix, std::unordered_map<std::string, 
int> name_map) {
+  auto n = make_object<NameSupplyNode>();
+  n->prefix_ = prefix;
+  n->name_map = std::move(name_map);
+  data_ = std::move(n);
+}
+
+String NameSupplyNode::ReserveName(const String& name, bool add_prefix) {
+  String final_name = name;
+  if (add_prefix) {
+    final_name = add_prefix_to_name(name);
+  }
+  name_map[final_name] = 0;
+  return final_name;
+}
+
+String NameSupplyNode::FreshName(const String& name, bool add_prefix) {
+  String unique_name = name;
+  if (add_prefix) {
+    unique_name = add_prefix_to_name(name);
+  }
+  unique_name = GetUniqueName(unique_name);
+  return unique_name;
+}
+
+bool NameSupplyNode::ContainsName(const String& name, bool add_prefix) {
+  String unique_name = name;
+  if (add_prefix) {
+    unique_name = add_prefix_to_name(name);
+  }
+
+  return name_map.count(unique_name);
+}
+
+String NameSupplyNode::add_prefix_to_name(const String& name) {
+  if (prefix_.empty()) {
+    return name;
+  }
+
+  std::ostringstream ss;
+  ICHECK(name.defined());
+  ss << prefix_ << "_" << name;
+  return ss.str();
+}
+
+std::string NameSupplyNode::GetUniqueName(std::string prefix) {
+  for (size_t i = 0; i < prefix.size(); ++i) {
+    if (prefix[i] == '.') prefix[i] = '_';
+  }

Review Comment:
   Why are you doing this replacement?



##########
src/ir/name_supply.cc:
##########
@@ -0,0 +1,112 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file name_supply.cc
+ * \brief NameSupply that can be used to generate unique variable names.
+ */
+#include "tvm/ir/name_supply.h"
+
+#include <tvm/runtime/registry.h>
+
+#include <utility>
+
+namespace tvm {
+
+NameSupply::NameSupply(const String& prefix, std::unordered_map<std::string, 
int> name_map) {
+  auto n = make_object<NameSupplyNode>();
+  n->prefix_ = prefix;
+  n->name_map = std::move(name_map);
+  data_ = std::move(n);
+}
+
+String NameSupplyNode::ReserveName(const String& name, bool add_prefix) {
+  String final_name = name;
+  if (add_prefix) {
+    final_name = add_prefix_to_name(name);
+  }
+  name_map[final_name] = 0;
+  return final_name;
+}
+
+String NameSupplyNode::FreshName(const String& name, bool add_prefix) {
+  String unique_name = name;
+  if (add_prefix) {
+    unique_name = add_prefix_to_name(name);
+  }
+  unique_name = GetUniqueName(unique_name);
+  return unique_name;
+}
+
+bool NameSupplyNode::ContainsName(const String& name, bool add_prefix) {
+  String unique_name = name;
+  if (add_prefix) {
+    unique_name = add_prefix_to_name(name);
+  }
+
+  return name_map.count(unique_name);
+}
+
+String NameSupplyNode::add_prefix_to_name(const String& name) {
+  if (prefix_.empty()) {
+    return name;
+  }
+
+  std::ostringstream ss;
+  ICHECK(name.defined());
+  ss << prefix_ << "_" << name;
+  return ss.str();
+}
+
+std::string NameSupplyNode::GetUniqueName(std::string prefix) {
+  for (size_t i = 0; i < prefix.size(); ++i) {
+    if (prefix[i] == '.') prefix[i] = '_';
+  }
+  auto it = name_map.find(prefix);
+  if (it != name_map.end()) {
+    while (true) {

Review Comment:
   You could just do
   ```
   new_name = prefix;
   while (!name_map.insert({new_name, 0}).second) {
      new_name = prefix + ...
   }
   return new_name;
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscr...@tvm.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org

Reply via email to