[Profiler] Include ActivityType from Kineto (#80750)

We don't want to compile with Kineto on all platforms, but if we're going to have significant integration between profiler and Kineto profiler will need to be able to rely on simple API constructs like the Kineto enums.

Differential Revision: [D37228314](https://our.internmc.facebook.com/intern/diff/D37228314/)

**NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D37228314/)!
Pull Request resolved: https://github.com/pytorch/pytorch/pull/80750
Approved by: https://github.com/aaronenyeshi

BUILD.bazel

@@ -1605,6 +1605,7 @@ cc_library(
         "torch/csrc/distributed",
         "torch/lib",
         "torch/lib/libshm",
+        "third_party/kineto/libkineto/include",
     ],
     visibility = ["//visibility:public"],
     deps = [
@@ -1656,6 +1657,7 @@ cc_library(
     deps = [
         ":caffe2",
         ":torch_headers",
+        "@kineto",
     ] + if_cuda([
         ":torch_distributed_cuda",
         "@cuda//:nvToolsExt",

WORKSPACE

@@ -126,6 +126,12 @@ new_local_repository(
     path = "third_party/fmt",
 )
 
+new_local_repository(
+    name = "kineto",
+    build_file = "//third_party:kineto.BUILD",
+    path = "third_party/kineto",
+)
+
 new_patched_local_repository(
     name = "tbb",
     patches = [

buckbuild.bzl

@@ -132,6 +132,7 @@ THIRD_PARTY_LIBS = {
     "gmock": ["//xplat/third-party/gmock:gtest", "//third_party:gmock"],
     "gtest": ["//xplat/third-party/gmock:gmock", "//third_party:gtest"],
     "kineto": ["//xplat/kineto/libkineto:libkineto", "//third_party:libkineto"],
+    "libkineto_headers": ["//xplat/kineto/libkineto:libkineto_headers", "//third_party:libkineto_headers"],
     "omp": ["//xplat/third-party/linker_lib:omp", "//third_party:no-op"],
     "psimd": ["//xplat/third-party/psimd:psimd", "//third_party:psimd"],
     "pthreadpool": ["//xplat/third-party/pthreadpool:pthreadpool", "//third_party:pthreadpool"],
@@ -1132,6 +1133,7 @@ def define_buck_targets(
             ":generated-autograd-headers",
             ":torch_headers",
             C10,
+            third_party("libkineto_headers"),
         ],
     )
 

caffe2/CMakeLists.txt

@@ -1137,9 +1137,11 @@ endif()
   target_include_directories(torch_cpu PRIVATE
     ${TORCH_ROOT}/third_party/miniz-2.1.0)
 
+  target_include_directories(torch_cpu PRIVATE
+    ${TORCH_ROOT}/third_party/kineto/libkineto/include)
+
   if(USE_KINETO)
     target_include_directories(torch_cpu PRIVATE
-      ${TORCH_ROOT}/third_party/kineto/libkineto/include
       ${TORCH_ROOT}/third_party/kineto/libkineto/src)
   endif()
 

third_party/kineto.BUILD

@@ -0,0 +1,10 @@
+load("@rules_cc//cc:defs.bzl", "cc_library")
+
+cc_library(
+    name = "kineto",
+    hdrs = glob(["libkineto/include/*.h",]),
+    includes = [
+        "libkineto/include/",
+    ],
+    visibility = ["//visibility:public"],
+)

third_party/kineto.buck.bzl

@@ -168,3 +168,14 @@ def define_kineto():
             ":fmt",
         ],
     )
+
+    cxx_library(
+        name = "libkineto_headers",
+        exported_headers = native.glob([
+            "kineto/include/*.h",
+        ]),
+        public_include_directories = [
+            "kineto/include",
+        ],
+        visibility = ["PUBLIC"],
+    )

torch/CMakeLists.txt

@@ -74,6 +74,7 @@ set(TORCH_PYTHON_INCLUDE_DIRECTORIES
     ${TORCH_ROOT}/third_party/gloo
     ${TORCH_ROOT}/third_party/onnx
     ${TORCH_ROOT}/third_party/flatbuffers/include
+    ${TORCH_ROOT}/third_party/kineto/libkineto/include
 
     ${TORCH_SRC_DIR}/csrc
     ${TORCH_SRC_DIR}/csrc/api/include

torch/csrc/autograd/profiler_kineto.cpp

@@ -62,7 +62,6 @@ using torch::profiler::impl::Result;
 using torch::profiler::impl::shapesToStr;
 using torch::profiler::impl::stacksToStr;
 using torch::profiler::impl::kineto::annotation_t;
-using torch::profiler::impl::kineto::KinetoActivityType;
 
 struct EventFieldsVisitor {
   EventFieldsVisitor(
@@ -341,7 +340,7 @@ struct KinetoThreadLocalState : public ProfilerThreadLocalStateBase {
         int64_t start_us = e->start_time_ns_ / 1000;
         int64_t end_us = e->endTimeNS() / 1000;
         kineto_events_.emplace_back(
-            e->kinetoType() == KinetoActivityType::PYTHON_FUNCTION);
+            e->kinetoType() == libkineto::ActivityType::PYTHON_FUNCTION);
         kineto_events_.back()
             .name(e->name())
             .startUs(start_us)

torch/csrc/profiler/collection.cpp

@@ -210,12 +210,11 @@ std::string toString(const ExtraFields<EventType::PyCall>& e) {
       e.callsite_.funcname_.str());
 }
 
-using torch::profiler::impl::kineto::KinetoActivityType;
 namespace {
-KinetoActivityType scopeToType(at::RecordScope scope) {
+auto scopeToType(at::RecordScope scope) {
   return scope == at::RecordScope::USER_SCOPE
-      ? KinetoActivityType::USER_ANNOTATION
-      : KinetoActivityType::CPU_OP;
+      ? libkineto::ActivityType::USER_ANNOTATION
+      : libkineto::ActivityType::CPU_OP;
 }
 } // namespace
 
@@ -230,9 +229,9 @@ DEFINE_VISITOR(
     kinetoType,
     scopeToType(e.scope_),
     scopeToType(e.scope_),
-    KinetoActivityType::CPU_INSTANT_EVENT,
-    KinetoActivityType::PYTHON_FUNCTION,
-    KinetoActivityType::PYTHON_FUNCTION);
+    libkineto::ActivityType::CPU_INSTANT_EVENT,
+    libkineto::ActivityType::PYTHON_FUNCTION,
+    libkineto::ActivityType::PYTHON_FUNCTION);
 DEFINE_VISITOR(correlationID, e.correlation_id_, 0, 0, 0, 0);
 DEFINE_VISITOR(
     endTimeNS,

torch/csrc/profiler/collection.h

@@ -188,7 +188,7 @@ struct TORCH_API Result : public std::enable_shared_from_this<Result> {
   }
 
   std::string name() const;
-  torch::profiler::impl::kineto::KinetoActivityType kinetoType() const;
+  libkineto::ActivityType kinetoType() const;
   uint64_t correlationID() const;
   int64_t endTimeNS() const;
   uint64_t endTID() const;

torch/csrc/profiler/kineto_shim.cpp

@@ -60,30 +60,9 @@ TraceWrapper::TraceWrapper(const int64_t start_time, const std::string& name)
 }
 #endif // USE_KINETO
 
-#ifdef USE_KINETO
-namespace {
-libkineto::ActivityType toActivityType(const KinetoActivityType type) {
-  switch (type) {
-    case KinetoActivityType::CPU_OP:
-      return libkineto::ActivityType::CPU_OP;
-    case KinetoActivityType::CPU_INSTANT_EVENT:
-      return libkineto::ActivityType::CPU_INSTANT_EVENT;
-    case KinetoActivityType::PYTHON_FUNCTION:
-      return libkineto::ActivityType::PYTHON_FUNCTION;
-    default:
-      TORCH_INTERNAL_ASSERT(
-          type == KinetoActivityType::USER_ANNOTATION,
-          "Invalid KinetoActivityType: ",
-          (int)type);
-      return libkineto::ActivityType::USER_ANNOTATION;
-  }
-}
-} // namespace
-#endif // USE_KINETO
-
 void TraceWrapper::addCPUActivity(
     const std::string& name,
-    const KinetoActivityType kineto_type,
+    const libkineto::ActivityType type,
     const DeviceAndResource device_and_resource,
     const uint64_t correlation_id,
     const int64_t start_time,
@@ -91,7 +70,6 @@ void TraceWrapper::addCPUActivity(
     const annotation_t& annotations) {
 #ifdef USE_KINETO
   TORCH_CHECK((bool)(*this), "Cannot add event to non-existent trace.");
-  auto type = toActivityType(kineto_type);
   cpu_trace_->emplace_activity(cpu_trace_->span, type, name);
   auto& act = libkineto::CpuTraceBuffer::toRef(cpu_trace_->activities.back());
   act.device = device_and_resource.device;

torch/csrc/profiler/kineto_shim.h

@@ -12,13 +12,14 @@
 #undef USE_KINETO
 #endif
 
+#include <ActivityType.h>
+
 #include <torch/csrc/Export.h>
 #include <torch/csrc/profiler/api.h>
 
 #ifdef USE_KINETO
 // Forward declarations so we don't have to include `libkineto.h` in a header.
 namespace libkineto {
-enum class ActivityType;
 struct CpuTraceBuffer;
 class ActivityTraceInterface;
 } // namespace libkineto
@@ -58,14 +59,6 @@ using trace_t = DummyTraceBuffer;
 using interface_trace_t = DummyTraceBuffer;
 #endif // USE_KINETO
 
-// Subset of `libkineto::ActivityType` for `addCPUActivity`.
-enum class KinetoActivityType : uint8_t {
-  CPU_OP = 0,
-  CPU_INSTANT_EVENT,
-  USER_ANNOTATION,
-  PYTHON_FUNCTION
-};
-
 using annotation_t = std::vector<std::pair<std::string, std::string>>;
 
 // Wraps: libkineto::CpuTraceBuffer
@@ -77,7 +70,7 @@ struct TraceWrapper {
   // The caller is expected to hold a mutex when calling `addCPUActivity`.
   void addCPUActivity(
       const std::string& name,
-      const KinetoActivityType kineto_type,
+      const libkineto::ActivityType type,
       const DeviceAndResource device_and_resource,
       const uint64_t correlation_id,
       const int64_t start_time,