[Relay tests] AlterOpLayout - Temporary attr update

apache · Nov 17, 2019 · 6d727b6 · 6d727b6
1 parent 5d66e7a
commit 6d727b6
Show file tree

Hide file tree

Showing 4 changed files with 102 additions and 32 deletions.
diff --git a/include/tvm/relay/op.h b/include/tvm/relay/op.h
@@ -258,6 +258,12 @@ class OpRegistry {
   inline OpRegistry& set_attr(const std::string& attr_name,  // NOLINT(*)
                               const ValueType& value, int plevel = 10);
 
+  /*!
+   * \brief Resets an attr of the registry.
+   * \param attr_name The name of the attribute.
+   */
+  inline void reset_attr(const std::string& attr_name);
+
   // set the name of the op to be the same as registry
   inline OpRegistry& set_name() {  // NOLINT(*)
     if (get()->name.length() == 0) {

diff --git a/python/tvm/relay/op/op.py b/python/tvm/relay/op/op.py
@@ -64,6 +64,22 @@ def set_attr(self, attr_name, value, plevel=10):
         """
         _OpSetAttr(self, attr_name, value, plevel)
 
+    def reset_attr(self, attr_name):
+        """Reset attribute about the operator.
+
+        Parameters
+        ----------
+        attr_name : str
+            The attribute name
+
+        value : object
+            The attribute value
+
+        plevel : int
+            The priority level
+        """
+        _OpResetAttr(self, attr_name)
+
 
 def get(op_name):
     """Get the Op for a given name

diff --git a/src/relay/ir/op.cc b/src/relay/ir/op.cc
@@ -95,6 +95,17 @@ const bool Op::HasGenericAttr(const std::string& key) {
   return true;
 }
 
+// Resets attr of the OpMap.
+void OpRegistry::reset_attr(const std::string& key) {
+  OpManager* mgr = OpManager::Global();
+  std::lock_guard<std::mutex> lock(mgr->mutex);
+  std::unique_ptr<GenericOpMap>& op_map = mgr->attr[key];
+  if (op_map == nullptr) {
+    return;
+  }
+  op_map->data_.clear();
+}
+
 void OpRegistry::UpdateAttr(const std::string& key,
                             TVMRetValue value,
                             int plevel) {
@@ -152,6 +163,15 @@ TVM_REGISTER_API("relay.op._OpSetAttr")
     reg.set_attr(attr_name, value, plevel);
   });
 
+TVM_REGISTER_API("relay.op._OpResetAttr")
+.set_body([](TVMArgs args, TVMRetValue* rv) {
+    Op op = args[0];
+    std::string attr_name = args[1];
+    auto& reg =
+        OpRegistry::Registry()->__REGISTER_OR_GET__(op->name).set_name();
+    reg.reset_attr(attr_name);
+  });
+
 TVM_REGISTER_API("relay.op._Register")
 .set_body([](TVMArgs args, TVMRetValue* rv) {
     std::string op_name = args[0];