[Partitioner] Heterogeneous partition testing

include/glow/Backends/DeviceManager.h

@@ -89,6 +89,15 @@ class DeviceManager {
   /// \returns the name of backend that powers this Device.
   llvm::StringRef getBackendName() { return config_.backendName; }
 
+  /// \returns a string with \p name in parameters.
+  llvm::StringRef getParamByName(llvm::StringRef name) const {
+    auto it = config_.parameters.find(name);
+    if (it != config_.parameters.end()) {
+      return it->second;
+    }
+    return "";
+  }
+
   /// \returns the maximum memory (in bytes) available on the device.
   virtual uint64_t getMaximumMemory() const = 0;
 

include/glow/Graph/Node.h

@@ -394,6 +394,16 @@ llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const Node &node);
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const Node *node);
 
+/// Helper to get the Kind of a Node (e.g. Kinded::Kind::AddNodeKind) given its
+/// \p nodeName (e.g. Add).
+inline Kinded::Kind getKindFromNodeName(llvm::StringRef nodeName) {
+#define DEF_NODE(CLASS, NAME)                                                  \
+  if (nodeName == #NAME) {                                                     \
+    return Kinded::Kind::CLASS##Kind;                                          \
+  }
+#include "glow/AutoGenNodes.def"
+  LOG(FATAL) << "Unknown node name: " << nodeName.str();
+}
 } // namespace glow
 
 namespace llvm {

include/glow/Partitioner/Partitioner.h

@@ -34,6 +34,10 @@ struct BackendInfo {
   uint64_t memSize;
   /// Backend pointer.
   Backend *backend = nullptr;
+  /// The non-supported nodes kind.
+  std::set<Kinded::Kind> nonSupportedNodesKinds;
+  /// The supported nodes kind.
+  std::set<Kinded::Kind> supportedNodesKinds;
 };
 
 /// A mapping of newly-created functions along with a set of nodes sets. The

include/glow/Partitioner/PartitionerUtils.h

@@ -82,5 +82,8 @@ GraphMemInfo updateGraphMemInfoByAddingNode(const NodesSet &currNodes,
 /// Return the memory usage of a given nodes set.
 GraphMemInfo getGraphMemInfo(const NodesSet &nodes);
 
+/// Parse a node name string (e.g. "Div,Add") \p names, \returns a set of
+/// NodeKinds corresponding to the names in the string.
+std::set<Kinded::Kind> generateNodeKindsSet(llvm::StringRef names);
 } // namespace glow
 #endif // GLOW_PARTITIONER_PARTITIONUTILS_H

include/glow/Runtime/RuntimeTypes.h

@@ -51,6 +51,14 @@ struct DeviceInfo {
   uint64_t availableMemory;
   /// Backend Type.
   std::string backendName;
+  /// A string contains the node names(e.g. Add, Div) which are separeted by
+  /// ",". E.g. "Div,Add". In Partitioner, those nodes won't be supported in
+  /// this backend.
+  std::string nonSupportedNodes;
+  /// A string contains the node names(e.g. Add, Div) which are separeted by
+  /// ",". E.g. "Div,Add". In Partitioner, the complementary set of those nodes
+  /// won't be supported in this backend.
+  std::string supportedNodes;
   /// Available SRAM capacity in bytes.
   uint64_t sramCapacity;
   /// Peak compute on device in ops/second. Assumes all ops are in int8.

lib/Partitioner/Partitioner.cpp

@@ -912,7 +912,29 @@ Partitioner::backendBasedPartition(Function *F,
   for (auto &N : F->getNodes()) {
     for (auto &backend : backends) {
       // Find the first backend that supports this node. The order of backends
-      // is important.
+      // is important. The check flow is :
+
+      // Step 1: If a node is in pre-defined non-supported nodes set, it can not
+      // be assigned to this backend. Continue.
+      const auto &nonSupportedNodesKinds =
+          backendMap_[backend->getBackendName()].nonSupportedNodesKinds;
+      if (nonSupportedNodesKinds.count(N.getKind())) {
+        // This op is on the pre-defined non-supported op list:
+        continue;
+      }
+      // Step 2: If the pre-defined supported nodes set is empty, it means all
+      // nodes could be assigned to this backend. If the pre-defined supported
+      // nodes set is not empty, we check that if the node from Step 1 is in
+      // this set or not. If not, continue.
+      const auto &supportedNodesKinds =
+          backendMap_[backend->getBackendName()].supportedNodesKinds;
+      if (!supportedNodesKinds.empty() &&
+          !supportedNodesKinds.count(N.getKind())) {
+        // This op is not on the pre-definded supported op list:
+        continue;
+      }
+      // Step 3: Check if the node is actually supported in this backend, if so,
+      // assign it to this backend and break. Otherwise continue.
       // TODO: the logic here need to be improved.
       if (backend->shouldLower(&N) || backend->isOpSupported(N)) {
         // Put this node into a partition for this backend.
@@ -981,6 +1003,10 @@ void Partitioner::getBackendMap(
       // is the same.
       // TODO : will improve the algorithm for different memory size.
       backendInfo.memSize = deviceInfo_[i].availableMemory;
+      backendInfo.nonSupportedNodesKinds =
+          generateNodeKindsSet(deviceInfo_[i].nonSupportedNodes);
+      backendInfo.supportedNodesKinds =
+          generateNodeKindsSet(deviceInfo_[i].supportedNodes);
       if (hasBackends) {
         backendInfo.backend = backends_[i];
       } else {

lib/Partitioner/PartitionerUtils.cpp

@@ -286,4 +286,17 @@ GraphMemInfo getGraphMemInfo(const NodesSet &nodes) {
   return ret;
 }
 
+std::set<Kinded::Kind> generateNodeKindsSet(llvm::StringRef names) {
+  std::set<Kinded::Kind> nodeKindsSet;
+  llvm::StringRef::size_type pos = names.find(',');
+  while (pos != llvm::StringRef::npos) {
+    nodeKindsSet.insert(getKindFromNodeName(names.substr(0, pos)));
+    names = names.substr(pos + 1);
+    pos = names.find(',');
+  }
+  if (!names.empty()) {
+    nodeKindsSet.insert(getKindFromNodeName(names));
+  }
+  return nodeKindsSet;
+}
 } // namespace glow

lib/Runtime/HostManager/HostManager.cpp

@@ -90,7 +90,8 @@ llvm::Error HostManager::addNetwork(std::unique_ptr<Module> module,
     DeviceInfo info = device.second->getDeviceInfo();
     info.availableMemory = device.second->getAvailableMemory();
     info.backendName = device.second->getBackendName();
-    VLOG(1) << "AVAILABLE DEVICE MEMORY " << info.availableMemory << "\n";
+    info.nonSupportedNodes = device.second->getParamByName("nonSupportedNodes");
+    info.supportedNodes = device.second->getParamByName("supportedNodes");
     deviceInfo.push_back(info);
   }
   // Perform a round of target-independent graph optimizations. This helps the

tests/images/run.sh

@@ -145,6 +145,10 @@ done
 ./bin/image-classifier tests/images/imagenet/*.png -expected-labels=${imagenetIdxValues} -image-mode=0to1 -m=quant_resnet50 -model-input-name=gpu_0/data_0 -use-imagenet-normalization "$@"
 num_errors=$(($num_errors + $?))
 
+# Heterogeneous partition Resnet50 Caffe2 model test
+./bin/image-classifier tests/images/imagenet/*.png -image-mode=0to1 -m=resnet50 -model-input-name=gpu_0/data -cpu-memory=100000 -load-device-configs="tests/runtime_test/heterogeneousConfigs.yaml" "$@"
+num_errors=$(($num_errors + $?))
+
 # Emotion_ferplus onnx model test
 i=0
 for png_filename in tests/images/EmotionSampleImages/*.png; do

tests/runtime_test/heterogeneousConfigs.yaml

@@ -0,0 +1,15 @@
+---
+- name:     Device1
+  backendName: CPU
+  parameters: |
+    "nonSupportedNodes":"AvgPool"
+    "deviceID" : "0"
+- name:     Device2
+  backendName: CPU
+  parameters: |
+    "nonSupportedNodes":"AvgPool"
+- name:     Device3
+  backendName: Interpreter
+  parameters: |
+    "supportedNodes":""
+...