[Partitioner] Partitioner refactoring 3

include/glow/Partitioner/Partitioner.h

@@ -16,7 +16,7 @@
 #ifndef GLOW_PARTITIONER_PARTITIONER_H
 #define GLOW_PARTITIONER_PARTITIONER_H
 
-#include "glow/Partitioner/PartitionerTypes.h"
+#include "glow/Partitioner/PartitionerBase.h"
 #include "glow/Support/Error.h"
 
 namespace glow {
@@ -25,7 +25,7 @@ using namespace runtime;
 
 /// Given a module, partitions each of the its functions into multiple ones
 /// based on memory constraints and minimizes the communication cost.
-class Partitioner {
+class Partitioner final : public PartitionerBase {
   /// The module that needs to be decomposed.
   Module *module_;
 
@@ -36,7 +36,10 @@ class Partitioner {
   /// The cost model related to device.
   std::vector<DeviceInfo> deviceInfo_;
 
-  /// The backend pointers.
+  /// The backends created in Partitioner. Used for function optimization.
+  std::vector<std::unique_ptr<Backend>> backendHolder;
+
+  /// The raw backend pointers.
   std::vector<Backend *> backends_;
 
   /// The map between backend name and BackendInfo.
@@ -50,9 +53,6 @@ class Partitioner {
   /// needed after partitions.
   DeviceIDTy logicalDeviceID_;
 
-  /// The result of module partitioning.
-  DAGListTy partitions_;
-
   /// Total memory (bytes) requested by one module.
   uint64_t memSize_;
 
@@ -82,11 +82,19 @@ class Partitioner {
   NodeToFunctionMap selectPartitions(Function *F, uint64_t availableMemory,
                                      llvm::StringRef backendName);
 
-  /// Duplicates all networks in the module order to saturate the Host.
-  void saturateHost(unsigned logicalDeviceCount);
-
-  FunctionToBackendNameMap
-  backendBasedPartition(Function *F, std::vector<Backend *> &backends,
+  /// Duplicates \p partitions in the module order to saturate the Host. \p
+  /// logicalDeviceCount is the number of logical devices used by the current
+  /// partitions. For example: If a network is partitioned into two parts (\p
+  /// logicalDeviceCount) and there are six devices this would duplicate the
+  /// network three times.
+  void saturateHost(unsigned logicalDeviceCount, const DAGListTy &partitions);
+
+  /// Partition a function \p F based on backends \p backends. \returns the
+  /// final partition result(or an err) and a map between partitions and backend
+  /// names. \p cctx is used for functions optimization.
+  llvm::Expected<DAGListTy>
+  backendBasedPartition(FunctionToBackendNameMap &funcToBackend, Function *F,
+                        std::vector<Backend *> &backends,
                         CompilationContext &cctx);
 
   /// Performs a load balancing optimization pass to optimize for load
@@ -96,26 +104,20 @@ class Partitioner {
                                        llvm::StringRef backendName,
                                        NodeToFunctionMap &mapping);
 
-  /// Given the node-function mapping, do the actual partitioning. If \p saveDAG
-  /// is true, the DAG will be saved into partitions_, which is the final
-  /// partition result.
-  void doPartitioning(llvm::StringRef funcName, std::vector<Function *>,
-                      NodeToFunctionMap &mapping, bool saveDAG);
-
   /// If there is no need to do any partition, just generate the DAGNode based
   /// on current functions in this module for backend \p backendName found in \p
-  /// backendMap. \p cctx is used during optimization of the Function. \returns
-  /// whether there was an error encountered.
-  llvm::Error
+  /// backendMap. \p cctx is used for function optimization. \returns the
+  /// partition result or an error.
+  llvm::Expected<DAGListTy>
   createDAGWithoutPartition(llvm::StringRef backendName,
                             std::map<std::string, BackendInfo> &backendMap,
                             CompilationContext &cctx);
 
-  /// Get the map between the backend name and the concrete backend info (e.g.
-  /// backend pointer, mem, number) used in this partiton. If there are backends
-  /// need to be created, we use \p backendsHolder to hold them for memory
-  /// purpose.
-  void getBackendMap(std::map<std::string, BackendInfo> &backendMap,
+  /// Create the map between the backend name and the concrete backend info
+  /// (e.g. backend pointer, mem, number) used in this partiton. If there are
+  /// backends need to be created, we use \p backendsHolder to hold them for
+  /// memory purpose.
+  void genBackendMap(std::map<std::string, BackendInfo> &backendMap,
                      std::vector<std::unique_ptr<Backend>> &backendsHolder,
                      std::vector<Backend *> &backends);
 
@@ -141,30 +143,30 @@ class Partitioner {
               const std::vector<Backend *> &backends, bool saturateHost = false,
               bool optimized = false);
 
-  /// Based on partitionConfig_ passed into Partitioner, do the user-defined
+  /// Based on \p partitionConfig passed into Partitioner, do user-defined
   /// partition.
-  llvm::Error PartitionFromConfig();
-
-  /// Decompose each function in a module. Now we support partitioning a module
-  /// among different type of devices. \p cctx is used during optimization of
-  /// the Function. \returns whether there was an error encountered.
-  llvm::Error Partition(CompilationContext &cctx);
+  llvm::Expected<DAGListTy>
+  partitionFromConfig(const PartitionConfig &partitionConfig);
 
   /// This partition approach is used in Glow Quantization Profiling flow. The
   /// backendBasedPartition is applied first in case there are heterogeneous
   /// backends. Then each sub-function will be compiled and run in CPU backend
-  /// for profiling.
-  llvm::Error QuantizationProfilingPartition(CompilationContext &cctx,
-                                             Function *F,
-                                             std::vector<Backend *> backends);
-
-  /// Get the final partitions.
-  DAGListTy &getPartitionResult() { return partitions_; }
-
-  /// Dump the partition result to a dot file. Since now all functions belong to
-  /// a function family and they have the same partition, we only dump the one
-  /// function's partition.
-  void dumpDAG(llvm::StringRef dotFilename) const;
+  /// for profiling. \p cctx is used for function optimization. \returns the
+  /// partition result or an error.
+  llvm::Expected<DAGListTy>
+  quantizationProfilingPartition(CompilationContext &cctx);
+
+  /// This partition approch first do the partition based on backend types, and
+  /// then based on cost models(memory usage and performance). \p cctx is used
+  /// for function optimization. \returns the partition result or an error.
+  llvm::Expected<DAGListTy> heterogeneousPartition(CompilationContext &cctx);
+
+  /// Decompose each function in a module. Given the parameters, this function
+  /// will choose different partition approches supported in this class:
+  /// heterogeneous partition, user-defined partition or quantization profiling.
+  /// \p cctx is used for function optimization. \returns the partition result
+  /// or an error.
+  llvm::Expected<DAGListTy> partition(CompilationContext &cctx) override;
 };
 } // namespace glow
 #endif // GLOW_PARTITIONER_PARTITIONER_H

include/glow/Partitioner/PartitionerBase.h

@@ -0,0 +1,49 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc.
+ *
+ * Licensed 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.
+ */
+#ifndef GLOW_PARTITIONER_PARTITIONERBASE_H
+#define GLOW_PARTITIONER_PARTITIONERBASE_H
+
+#include "glow/Partitioner/PartitionerTypes.h"
+#include "glow/Support/Error.h"
+
+namespace glow {
+
+using namespace runtime;
+/// Given a module, partitions each of the its functions into multiple ones
+/// based on memory constraints and minimizes the communication cost.
+class PartitionerBase {
+public:
+  virtual ~PartitionerBase() = default;
+
+  /// Decompose each function in a module. \p cctx is used in function
+  /// optimization. \returns the partition result.
+  virtual llvm::Expected<DAGListTy> partition(CompilationContext &cctx) = 0;
+
+  /// Dump the partition result \p partitions to a dot file with name \p
+  /// dotFilename. Since now all functions belong to a function family and they
+  /// have the same partition, we only dump the one function's partition.
+  void dumpDAG(llvm::StringRef dotFilename, const DAGListTy &partitions) const;
+
+protected:
+  /// Given the node-function mapping \p mapping, do the actual partitioning. If
+  /// \p saveDAG is true, the DAG will be generated. \returns the final
+  /// partitions or an empty partition (If \p saveDAG is false).
+  DAGListTy doPartitioning(llvm::StringRef funcName, std::vector<Function *>,
+                           Module *module, NodeToFunctionMap &mapping,
+                           bool saveDAG);
+};
+} // namespace glow
+#endif // GLOW_PARTITIONER_PARTITIONER_H

lib/Partitioner/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_library(Partitioner
+              PartitionerBase.cpp
               PartitionerUtils.cpp
               PartitionerOptimizer.cpp
               PartitionerValidation.cpp