Merge upstream MLCommons chakra changes

README.md

@@ -4,7 +4,7 @@ Chakra is an open and interoperable graph-based representation of AI/ML workload
 
 This is a repository of Chakra schema and a complementary set of tools and capabilities to enable the collection, analysis, generation, and adoption of Chakra execution traces by a broad range of simulators, emulators, and replay tools.
 
-Chakra is under active development as a [MLCommons](https://mlcommons.org/en) research project. Please see [MLCommons Chakra Working Group](https://mlcommons.org/en/groups/research-chakratracebench/) for more details for participating in this effort.
+Chakra is under active development as a [MLCommons](https://mlcommons.org/en)® research project. Please see [MLCommons Chakra Working Group](https://mlcommons.org/en/groups/research-chakratracebench/) for more details for participating in this effort.
 
 A detailed description of the original motivation and guiding principles can be found [here](https://arxiv.org/abs/2305.14516). The paper was published prior to Chakra becoming a MLCommons project. Please cite this repository to refer to the latest Chakra schema and tools.
 

pyproject.toml

@@ -1,5 +1,8 @@
 [build-system]
-requires = ["setuptools", "setuptools-grpc"]
+requires = [ "setuptools>=61",
+  "wheel",
+  "setuptools-grpc",
+  "grpcio-tools"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -12,10 +15,16 @@ authors = [
     {name = "MLCommons", email = "chakra@mlcommons.org"},
 ]
 dependencies = [
-    "protobuf==5.*",
+    "protobuf"
+]
+
+[project.optional-dependencies]
+visualizer = [
     "graphviz",
     "networkx",
-    "pydot",
+    "pydot"
+]
+trace_link = [
     "HolisticTraceAnalysis @ git+https://github.com/facebookresearch/HolisticTraceAnalysis.git@d731cc2e2249976c97129d409a83bd53d93051f6"
 ]
 

setup.py

@@ -1,7 +1,5 @@
-from distutils.command.build import build
-
 from setuptools import setup
-
+from setuptools.command.build import build
 
 class build_grpc(build):
     """

src/feeder/et_feeder.h

@@ -10,10 +10,7 @@
 #include "protoio.hh"
 
 namespace Chakra {
-struct CompareNodes : public std::binary_function<
-                          std::shared_ptr<ETFeederNode>,
-                          std::shared_ptr<ETFeederNode>,
-                          bool> {
+struct CompareNodes {
   bool operator()(
       const std::shared_ptr<ETFeederNode> lhs,
       const std::shared_ptr<ETFeederNode> rhs) const {

src/feeder/et_feeder_node.cpp

@@ -9,6 +9,7 @@ ETFeederNode::ETFeederNode(std::shared_ptr<ChakraProtoMsg::Node> node) {
   this->name_ = node->name();
   this->runtime_ = node->duration_micros();
   this->is_cpu_op_ = 0;
+  this->num_ops_ = 0;
 
   if (node->has_inputs()) {
     this->inputs_values_ = static_cast<string>(node->inputs().values());

src/feeder_v3/cache.h

@@ -0,0 +1,101 @@
+#ifndef CHAKRA_FEEDER_V3_CACHE_H
+#define CHAKRA_FEEDER_V3_CACHE_H
+
+#include <list>
+#include <memory>
+#include <mutex>
+#include <shared_mutex>
+#include <unordered_map>
+
+namespace Chakra {
+namespace FeederV3 {
+
+template <typename K, typename V>
+class Cache {
+ public:
+  Cache(size_t capacity) : capacity(capacity) {}
+  void put(const K& key, const V& value) {
+    std::unique_lock lock(cache_mutex);
+    if (this->cache.find(key) != this->cache.end()) {
+      // hit and update
+      this->lru.erase(this->cache[key].second);
+      this->lru.push_back(key);
+      this->cache[key].second = --this->lru.end();
+      this->cache[key].first = std::make_shared<V>(value);
+    } else {
+      // miss
+      while (this->cache.size() >= this->capacity) {
+        // evict
+        auto victim = this->lru.front();
+        this->cache.erase(victim);
+        this->lru.pop_front();
+      }
+      // and put new
+      this->lru.push_back(key);
+      this->cache[key] =
+          std::make_pair(std::make_shared<V>(value), --this->lru.end());
+    }
+  }
+  bool has(const K& key) {
+    std::shared_lock lock(cache_mutex);
+    return this->cache.find(key) != this->cache.end();
+  }
+  std::weak_ptr<const V> get(const K& key) {
+    std::shared_lock lock(cache_mutex);
+    if (this->cache.find(key) == this->cache.end()) {
+      throw std::runtime_error("Key not found in cache");
+    }
+    std::weak_ptr<const V> value(this->cache.at(key).first);
+    return value;
+  }
+  std::shared_ptr<const V> get_locked(const K& key) {
+    std::shared_lock lock(cache_mutex);
+    if (this->cache.find(key) == this->cache.end()) {
+      throw std::runtime_error("Key not found in cache");
+    }
+    return this->cache.at(key).first;
+  }
+  std::weak_ptr<const V> get_or_null(const K& key) {
+    std::shared_lock lock(cache_mutex);
+    if (this->cache.find(key) == this->cache.end()) {
+      return std::weak_ptr<V>();
+    }
+    std::weak_ptr<const V> value(this->cache.at(key).first);
+    return value;
+  }
+  std::shared_ptr<const V> get_or_null_locked(const K& key) {
+    std::shared_lock lock(cache_mutex);
+    if (this->cache.find(key) == this->cache.end()) {
+      return std::shared_ptr<const V>();
+    }
+    return this->cache.at(key).first;
+  }
+
+  void remove(const K& key) {
+    std::unique_lock lock(cache_mutex);
+    if (this->cache.find(key) == this->cache.end()) {
+      throw std::runtime_error("Key not found in cache");
+    }
+    this->lru.erase(this->cache[key].second);
+    this->cache.erase(key);
+  }
+
+  ~Cache() {
+    this->cache.clear();
+    this->lru.clear();
+  }
+
+ private:
+  size_t capacity;
+  std::unordered_map<
+      K,
+      std::pair<std::shared_ptr<V>, typename std::list<K>::iterator>>
+      cache;
+  std::list<K> lru;
+  mutable std::shared_mutex cache_mutex;
+};
+
+} // namespace FeederV3
+} // namespace Chakra
+
+#endif

src/feeder_v3/common.h

@@ -0,0 +1,32 @@
+#ifndef CHAKRA_FEEDER_V3_COMMON_H
+#define CHAKRA_FEEDER_V3_COMMON_H
+#include <cstdint>
+#include "et_def.pb.h"
+
+namespace Chakra {
+namespace FeederV3 {
+using NodeId = uint64_t;
+using ETFeederId = uint64_t;
+using ChakraNode = ChakraProtoMsg::Node;
+using ChakraGlobalMetadata = ChakraProtoMsg::GlobalMetadata;
+using ChakraAttr = ChakraProtoMsg::AttributeProto;
+
+constexpr static bool ALLOW_IMPLICIT_INTEGER_CONVERSION = true;
+constexpr static bool ALLOW_IMPLICIT_FLOAT_CONVERSION = true;
+constexpr static bool ALLOW_IMPLICIT_INTEGER_TO_FLOAT_CONVERSION = true;
+constexpr static bool ALLOW_IMPLICIT_FLOAT_TO_INTEGER_CONVERSION = false;
+constexpr static bool NO_IMPLICIT_CONVERSION = false;
+constexpr static bool DEFAULT_STRICT_TYPING = false;
+
+constexpr static size_t DEFAULT_ETFEEDER_CACHE_SIZE = 16384;
+constexpr static bool RESOLVE_DATA_DEPS = true;
+constexpr static bool RESOLVE_CTRL_DEPS = false;
+
+constexpr static size_t DEFAULT_PROTOBUF_BUFFER_SIZE = 16384;
+
+constexpr static bool SOFT_SANITY_CHECK = true;
+
+} // namespace FeederV3
+} // namespace Chakra
+
+#endif