General Layout Support

include/nnvm/compiler/op_attr_types.h

@@ -16,6 +16,7 @@
 #include <nnvm/graph.h>
 #include <vector>
 #include <string>
+#include "packed_func_ext.h"
 
 namespace nnvm {
 namespace compiler {
@@ -73,19 +74,17 @@ using FTVMSchedule = std::function<
            const Array<Tensor>& outs,
            const std::string& target)>;
 
-/*! \brief Layout Information about an entry */
-using TLayoutInfo = std::string;
-
 /*!
- * \brief The producer consumer function of node layout
- * \param attrs The attribute of the node.
- * \param ilayouts The input layouts that the node request.
- * \param olayouts The output layouts that the node produce.
- * \return bool The success flag.
+ * \brief Modify the op node to alter its input layout.
+ *  it is invoked in AlterOpLayout pass.
+ * \param attrs The attribute of the original node.
+ * \param inputs The input symbols of the original node.
+ * \param tinfos The inferred shape and dtype of the inputs.
  */
-using FTVMLayoutRequest = std::function<bool (const NodeAttrs& attrs,
-                                              std::vector<TLayoutInfo> *ilayouts,
-                                              std::vector<TLayoutInfo> *olayouts)>;
+using FTVMAlterOpLayout = std::function<
+  Symbol(const NodeAttrs& attrs,
+         const SymbolArray& inputs,
+         const Array<Tensor>& tinfos)>;
 
 /*!
  * \brief Transform from normal operator to vectorized operator

include/nnvm/compiler/packed_func_ext.h

@@ -11,6 +11,7 @@
 #include <nnvm/graph.h>
 #include <nnvm/symbolic.h>
 #include <string>
+#include <vector>
 #include <unordered_map>
 
 namespace nnvm {
@@ -19,6 +20,7 @@ namespace compiler {
 using tvm::runtime::PackedFunc;
 
 using AttrDict = std::unordered_map<std::string, std::string>;
+using SymbolArray = std::vector<const Symbol*>;
 
 /*!
  * \brief Get PackedFunction from global registry and
@@ -52,6 +54,12 @@ template<>
 struct extension_class_info<nnvm::compiler::AttrDict> {
   static const int code = 18;
 };
+
+template<>
+struct extension_class_info<nnvm::compiler::SymbolArray> {
+  static const int code = 19;
+};
+
 }  // namespace runtime
 }  // namespace tvm
 #endif  // NNVM_COMPILER_PACKED_FUNC_EXT_H_

include/nnvm/graph_attr_types.h

@@ -9,6 +9,7 @@
 #include <vector>
 #include <string>
 #include "./tuple.h"
+#include "./layout.h"
 
 namespace nnvm {
 
@@ -46,14 +47,29 @@ using ShapeVector = std::vector<TShape>;
  * \code
  *  Graph g = ApplyPass(src_graph, "InferType");
  *  const DTypeVector& types = g.GetAttr<DTypeVector>("dtype");
- *  // get shape by entry id
+ *  // get type by entry id
  *  int entry_type = dtypes[g.indexed_graph().entry_id(my_entry)];
  * \endcode
  *
  * \sa FInferType
  */
 using DTypeVector = std::vector<int>;
 
+/*!
+ * \brief The result holder of layout of each NodeEntry in the graph.
+ * \note Stored under graph.attrs["layout"], provided by Pass "InferType"
+ *
+ * \code
+ *  Graph g = ApplyPass(src_graph, "LayoutTransform");
+ *  const LayoutVector& layouts = g.GetAttr<LayoutVector>("layout");
+ *  // get layout by entry id
+ *  int entry_layout = layouts[g.indexed_graph().entry_id(my_entry)];
+ * \endcode
+ *
+ * \sa FInferLayout
+ */
+using LayoutVector = std::vector<Layout>;
+
 /*!
  * \brief The result holder of device of each operator in the graph.
  * \note Stored under graph.attrs["device"], provided by Pass "PlaceDevice"