[nativert] Move GraphExecutorBase to PyTorch core (#156196)

Summary:
Moves GraphExecutorBase class to PyTorch core.
GraphExecutorBase is a lightweight abstraction to execute a graph with  execution frames without actually owning the graph nor the weights. This is introduced to decouple the state management of the top level runtime from the kernel executions so that sub graphs from higher order ops can be supported.

Torch Native Runtime RFC: pytorch/rfcs#72

Test Plan:
CI

Rollback Plan:

Differential Revision: D76830436

Pull Request resolved: #156196
Approved by: https://github.com/zhxchen17

Author: yiming0416

build_variables.bzl

@@ -600,6 +600,7 @@ libtorch_nativert_sources = [
     "torch/nativert/executor/Placement.cpp",
     "torch/nativert/executor/ExecutionPlanner.cpp",
     "torch/nativert/executor/ExecutionFrame.cpp",
+    "torch/nativert/executor/GraphExecutorBase.cpp",
     "torch/nativert/executor/OpKernel.cpp",
     "torch/nativert/executor/PlacementUtils.cpp",
     "torch/nativert/executor/Weights.cpp",

torch/nativert/executor/GraphExecutorBase.cpp

@@ -0,0 +1,120 @@
+#include <ATen/record_function.h>
+#include <torch/nativert/executor/GraphExecutorBase.h>
+
+#include <c10/util/Logging.h>
+#include <caffe2/core/timer.h>
+
+namespace torch::nativert {
+
+GraphExecutorBase::GraphExecutorBase(
+    const Graph& graph,
+    std::vector<std::unique_ptr<OpKernel>> nodeKernels,
+    const ExecutorConfig& executorConfig)
+    : graph_(graph),
+      nodeKernels_(std::move(nodeKernels)),
+      executorConfig_(executorConfig),
+      execPlan_(ExecutionPlanner{graph_}.createPlan()) {};
+
+void GraphExecutorBase::fillUserInputs(
+    ExecutionFrame& frame,
+    std::vector<c10::IValue> inputs) {
+  RECORD_USER_SCOPE("Executor::fillUserInputs");
+  const auto& inputValues = graph_.userInputs();
+  TORCH_CHECK_EQ(inputValues.size(), inputs.size());
+
+  // load user input tensor into execution frame
+  for (size_t i = 0; i < inputValues.size(); i++) {
+    if (inputValues[i]) {
+      frame.setIValue(inputValues[i]->id(), std::move(inputs[i]));
+    }
+  }
+}
+
+ProfileMetrics GraphExecutorBase::benchmarkIndividualNodes(
+    ExecutionFrame& executionFrame,
+    std::vector<std::vector<c10::IValue>> inputsList,
+    const uint32_t warmupRuns,
+    const uint32_t mainRuns) {
+  // TODO: add support for memory profiling
+  TORCH_CHECK(warmupRuns >= 1 && mainRuns >= 1);
+
+  ProfileMetrics results;
+  const auto numNodes = static_cast<uint32_t>(nodeKernels_.size());
+  results.timePerNode.resize(numNodes, 0);
+  if (inputsList.empty()) {
+    auto i = 0;
+    for (const auto& nodeKernel : nodeKernels_) {
+      std::string target(nodeKernel->node()->target());
+      results.timePerNode[i] = 0;
+      results.timePerNodeType[target] = 0;
+      results.instancesPerNodeType[target]++;
+      if (nodeKernel->hasPrimKernel()) {
+        results.primNodesCount++;
+        results.primNodes.insert(target);
+      } else if (nodeKernel->hasStaticDispatch()) {
+        results.staticDispatchNodesCount++;
+        results.staticDispatchNodes.insert(target);
+      }
+      i++;
+    }
+    results.totalNodesCount = numNodes;
+    for (const auto& p : results.timePerNodeType) {
+      const std::string& kind = p.first;
+      results.percentPerNodeType[kind] = 0;
+    }
+    return results;
+  }
+
+  // Warmup
+  for (uint32_t i = 0; i < warmupRuns; i++) {
+    for (const auto& inputs : inputsList) {
+      execute(executionFrame, inputs);
+    }
+  }
+
+  // Execute kernels
+  caffe2::Timer timer;
+  for (uint32_t i = 0; i < mainRuns; i++) {
+    for (auto inputs : inputsList) {
+      const auto& inputValues = graph_.userInputs();
+
+      TORCH_CHECK_EQ(inputValues.size(), inputs.size());
+      for (size_t j = 0; j < inputValues.size(); j++) {
+        executionFrame.setIValue(inputValues[j]->id(), std::move(inputs[j]));
+      }
+      for (NodeIndex nodeIdx = 0; nodeIdx < nodeKernels_.size(); ++nodeIdx) {
+        timer.Start();
+        nodeKernels_[nodeIdx]->compute(executionFrame);
+        float millis = timer.MilliSeconds();
+        results.timePerNode[nodeIdx] += millis;
+      }
+    }
+  }
+
+  // Summarize results
+  const float numTotalIters =
+      (static_cast<float>(mainRuns) * static_cast<float>(inputsList.size()));
+  for (const auto i : c10::irange(numNodes)) {
+    const Node* node = nodeKernels_[i]->node();
+    std::string target(node->target());
+    results.timePerNode[i] /= numTotalIters;
+    results.timePerNodeType[target] += results.timePerNode[i];
+    results.instancesPerNodeType[target]++;
+    if (nodeKernels_[i]->hasPrimKernel()) {
+      results.primNodes.insert(target);
+      results.primNodesCount++;
+    } else if (nodeKernels_[i]->hasStaticDispatch()) {
+      results.staticDispatchNodes.insert(target);
+      results.staticDispatchNodesCount++;
+    }
+    results.totalTime += results.timePerNode[i];
+  }
+  results.totalNodesCount = numNodes;
+  for (const auto& r : results.timePerNodeType) {
+    const std::string& target = r.first;
+    results.percentPerNodeType[target] = r.second * 100.0 / results.totalTime;
+  }
+  return results;
+}
+
+} // namespace torch::nativert

torch/nativert/executor/GraphExecutorBase.h

@@ -0,0 +1,79 @@
+#pragma once
+
+#include <torch/nativert/executor/ExecutionFrame.h>
+#include <torch/nativert/executor/ExecutionPlanner.h>
+#include <torch/nativert/executor/ExecutorConfig.h>
+#include <torch/nativert/executor/OpKernel.h>
+#include <torch/nativert/graph/Graph.h>
+#include <torch/nativert/graph/GraphSignature.h>
+
+namespace torch::nativert {
+
+struct ProfileMetrics {
+  size_t primNodesCount{0};
+  size_t staticDispatchNodesCount{0};
+  size_t totalNodesCount{0};
+  std::vector<float> timePerNode;
+  std::unordered_map<std::string, float> timePerNodeType;
+  std::unordered_map<std::string, float> percentPerNodeType;
+  std::unordered_map<std::string, int> instancesPerNodeType;
+  std::unordered_set<std::string> staticDispatchNodes;
+  std::unordered_set<std::string> primNodes;
+  float totalTime{0};
+};
+
+/**
+ * GraphExecutor is a lightweight abstraction to execute a graph with
+ * execution frames without actually owning the graph nor the weights. This is
+ * introduced to decouple the state management of the top level runtime from the
+ * kernel executions so that sub graphs from higher order ops can be supported.
+ */
+class GraphExecutorBase {
+ public:
+  GraphExecutorBase(
+      const Graph& graph,
+      std::vector<std::unique_ptr<OpKernel>> nodeKernels,
+      const ExecutorConfig& executorConfig);
+  virtual ~GraphExecutorBase() = default;
+
+  const Graph& graph() const {
+    return graph_;
+  }
+
+  // This API only returns the flattened UserOutputs,
+  // intended to be used for Inference path
+  virtual std::vector<c10::IValue> execute(
+      ExecutionFrame& frame,
+      std::vector<c10::IValue> inputs) = 0;
+
+  virtual std::vector<c10::IValue> executeWithPrefilledFrame(
+      ExecutionFrame& frame) = 0;
+
+  ProfileMetrics benchmarkIndividualNodes(
+      ExecutionFrame& executionFrame,
+      std::vector<std::vector<c10::IValue>> inputs,
+      const uint32_t warmup_runs,
+      const uint32_t main_runs);
+
+  std::vector<std::unique_ptr<OpKernel>> stealKernels() {
+    return std::move(nodeKernels_);
+  }
+
+  void setKernels(std::vector<std::unique_ptr<OpKernel>>&& kernels) {
+    nodeKernels_ = std::move(kernels);
+  }
+
+ protected:
+  void fillUserInputs(ExecutionFrame& frame, std::vector<c10::IValue> inputs);
+
+  const Graph& graph_;
+
+  // cache of the constructed kernels to avoid reconstruction per execution
+  std::vector<std::unique_ptr<OpKernel>> nodeKernels_;
+
+  const ExecutorConfig& executorConfig_;
+
+  std::unique_ptr<ExecutionPlan> execPlan_;
+};
+
+} // namespace torch::nativert