Allow all enum int to be exposed

include/tvm/array.h

@@ -31,7 +31,7 @@ class ArrayNode : public Node {
 };
 
 /*!
- * \brief Immutable array container of NodeRef in DSL graph.
+ * \brief Array container of NodeRef in DSL graph.
  *  Array implements copy on write semantics, which means array is mutable
  *  but copy will happen when array is referenced in more than two places.
  * \tparam T The content NodeRef type.

include/tvm/base.h

@@ -23,7 +23,7 @@ class UnaryOp;
 class BinaryOp;
 
 /*! \brief list of all supported data types */
-enum DataType {
+enum DataType : int {
   kUnknown = 0,
   kInt32 = 1,
   kFloat32 = 2
@@ -56,10 +56,17 @@ class AttrVisitor {
 //! \cond Doxygen_Suppress
   virtual void Visit(const char* key, double* value) = 0;
   virtual void Visit(const char* key, int64_t* value) = 0;
-  virtual void Visit(const char* key, DataType* value) = 0;
+  virtual void Visit(const char* key, int* value) = 0;
   virtual void Visit(const char* key, std::string* value) = 0;
   virtual void Visit(const char* key, const UnaryOp** value) = 0;
   virtual void Visit(const char* key, const BinaryOp** value) = 0;
+  template<typename ENum,
+           typename = typename std::enable_if<std::is_enum<ENum>::value>::type>
+  void Visit(const char* key, ENum* ptr) {
+    static_assert(std::is_same<int, typename std::underlying_type<ENum>::type>::value,
+                  "declare enum to be enum int to use visitor");
+    this->Visit(key, reinterpret_cast<int*>(ptr));
+  }
 //! \endcond
 };
 

src/c_api/c_api.cc

@@ -45,7 +45,7 @@ struct APIAttrGetter : public AttrVisitor {
   void Visit(const char* key, int64_t* value) override {
     if (skey == key) *ret = value[0];
   }
-  void Visit(const char* key, DataType* value) override {
+  void Visit(const char* key, int* value) override {
     if (skey == key) *ret = static_cast<int64_t>(value[0]);
   }
   void Visit(const char* key, std::string* value) override {
@@ -68,7 +68,7 @@ struct APIAttrDir : public AttrVisitor {
   void Visit(const char* key, int64_t* value) override {
     names->push_back(key);
   }
-  void Visit(const char* key, DataType* value) override {
+  void Visit(const char* key, int* value) override {
     names->push_back(key);
   }
   void Visit(const char* key, std::string* value) override {