[RFC] Dynamic ActivityType Registration for Out-of-Tree Backends

[KarhouTam]

TL;DR

This RFC proposes a mechanism to register custom ActivityType values at runtime. This allows out-of-tree device backends to define their own profiling event categories without modifying Kineto or PyTorch source code. This change affects two repositories: Kineto (pytorch/kineto) and PyTorch (pytorch/pytorch).

Pain Point

IActivityProfiler::availableActivities() is constrained to libkineto::ActivityType, a compile-time enum with 26 fixed values. Only two are allocated to OOT use (PRIVATEUSE1_RUNTIME, PRIVATEUSE1_DRIVER). A typical accelerator needs distinct categories for such as compute kernels, synchronization, and hardware counters — far more than two.

Related Works

• [1/n] Add generalized event types and GPU Performance Monitoring counter event support #1212 is the primary long-term fix: device-agnostic generic types (RUNTIME, DRIVER, CONCURRENT_KERNEL, GPU_PM_COUNTER) replace the fragmented device-specific ones, with old names kept as deprecated aliases. This RFC addresses only the residual gap — truly vendor-specific semantics that cannot be generalized. The long-term direction is fewer static types, not more; this proposal adds a narrow dynamic layer for non-generalizable needs only.

• Add fine-grained activity type filtering to torch.profiler.profile pytorch#176351 lands fine-grained activity type filtering for static types. This RFC extends exactly that infrastructure to cover custom-registered types, so PR #176351 is a direct dependency for the PyTorch-side changes described here.

Proposed Solution

Add a custom ActivityType registry to libkineto, and extend the ActivityTypeMap / ActivityFilter infrastructure from PR pytorch/pytorch#176351 to cover custom-registered types on the PyTorch side.

Decision 1: Parallel type, not enum extension

Rather than modifying the existing ActivityType enum (which would be an ABI break), custom type IDs are cast directly to ActivityType via static_cast<ActivityType>(id). The existing enum is untouched. IDs are allocated in the range [kCustomActivityTypeBase, UINT8_MAX) — outside the 0–25 range of static enum values but round-trippable through static_cast.

Implementation:

• ActivityType.h introduces kCustomActivityTypeBase = 50; both ActivityType.h and Config.h now #include "CustomActivityTypeRegistry.h".
• CustomActivityTypeRegistry is a Meyers-singleton class providing:
  • registerType(plugin, type_name) → hash-based ID; idempotent.
  • lookup(type_name) → optional<int> ID if found.
  • entries() → all registered custom types for trace metadata embedding.
  • toCatString(id) → "plugin.type_name" string for trace display.

constexpr int kCustomActivityTypeBase = 50;

struct CustomActivityTypeEntry {
  int id;
  std::string plugin;
  std::string type_name;
};

class CustomActivityTypeRegistry {
 public:
  static CustomActivityTypeRegistry& instance();
  int registerType(const std::string& plugin, const std::string& type_name);
  std::optional<int> lookup(const std::string& type_name) const;
  std::string toCatString(int id) const;
  const std::vector<CustomActivityTypeEntry>& entries() const;
 private:
  CustomActivityTypeRegistry() = default;
  static int computeId(const std::string& plugin, const std::string& type_name);
  mutable std::mutex mutex_;
  std::vector<CustomActivityTypeEntry> entries_;
};

Decision 2: Hash-based ID allocation for cross-process stability

IDs are derived from a deterministic hash of <plugin>.<type_name>, mapped into [kCustomActivityTypeBase, UINT8_MAX). All ranks in a distributed job assign identical IDs without coordination.

Implementation: CustomActivityTypeRegistry.h uses FNV-1a hash with idempotent registration.

static int computeId(
      const std::string& plugin,
      const std::string& type_name) {
    const std::string key = plugin + "." + type_name;
    unsigned int hash = 216512456u;
    for (unsigned char c : key) {
      hash ^= c;
      hash *= 1234u;
    }
    const unsigned int range = UINT8_MAX - kCustomActivityTypeBase;
    return kCustomActivityTypeBase + static_cast<int>(hash % range);
  }

Decision 3: Extension of API to carry custom types

3.1 Mandatory trace metadata for self-describing traces

Custom ActivityType IDs are opaque integers; a trace file viewed on a machine without the OOT backend installed would show meaningless numbers. The mapping must travel with the trace.

toString(ActivityType t) in ActivityType.h is extended to handle IDs >= kCustomActivityTypeBase: it delegates to CustomActivityTypeRegistry::toCatString() which returns <plugin>.<type_name>. For IDs < kCustomActivityTypeBase, the existing constexpr array lookup is used.

Implementation: ActivityType.h extends toString(); CustomActivityTypeRegistry.h provides toCatString().

// ActivityType.h
inline const char* toString(ActivityType t) {
-  return _activityTypeNames[static_cast<int>(t)].name;
+  int type_id = static_cast<int>(t);
+  if (type_id < static_cast<int>(ActivityType::ENUM_COUNT)) {
+    return _activityTypeNames[type_id].name;
+  } else {
+    static thread_local std::string result;
+    const std::string cat = CustomActivityTypeRegistry::instance().toCatString(type_id);
+    result = cat.empty() ? std::to_string(type_id) : cat;
+    return result.c_str();
+  }
}

// CustomActivityTypeRegistry.h
std::string toCatString(int id) const {
    std::lock_guard<std::mutex> lock(mutex_);
    for (const auto& e : entries_) {
        if (e.id == id) {
        return e.plugin + "." + e.type_name;
        }
    }
    return "";
}

3.2 Reverse lookup: toActivityType()

ActivityType.cpp's toActivityType(string) now falls through to CustomActivityTypeRegistry::lookup() when no static enum name matches, enabling string→ID round-tripping for custom types.

ActivityType toActivityType(const std::string& str) {
  for (int i = 0; i < activityTypeCount; i++) {
    if (str == _activityTypeNames[i].name) {
      return _activityTypeNames[i].type;
    }
  }
+  auto custom_type_id = CustomActivityTypeRegistry::instance().lookup(str);
+  if (custom_type_id.has_value()) {
+    return static_cast<ActivityType>(*custom_type_id);
+  }
  throw std::invalid_argument(fmt::format("Invalid activity type: {}", str));
}

PyTorch-Side Integration (Building on PR #176351)

Registration API: privateuse1_profiler.h

registerPrivateUse1ActivityType(backend_name, type_name) forwards to CustomActivityTypeRegistry::instance().registerType() and returns the stable int ID. Registration is transparent to end users — when they import torch_<backend>, the C++ extension loads and all custom types are registered before any profile() call can occur.

Filtering: kineto_shim.cpp

Two changes extend the existing ActivityFilter plumbing:

1. filterActivities() — custom type ID propagation:

When a requested type name (e.g. "KERNEL") is absent from the static kPrivateUse1Types map, filterActivities() now:

• Looks up the name in CustomActivityTypeRegistry::instance().lookup()
• Adds the custom type ID itself to the result set (via static_cast<ActivityType>(*id))

std::unordered_set<libkineto::ActivityType> filterActivities(
    const ActivityTypeMap& defaults,
    const std::unordered_set<std::string>& requested) {
  std::unordered_set<libkineto::ActivityType> result;
  for (const auto& name : requested) {
    bool found = false;
    for (const auto& [type, type_name] : defaults) {
      if (type_name == name) {
        result.insert(type);
        found = true;
        break;
      }
    }
-    TORCH_CHECK(
-        found, "Unknown or non-member activity type name: '", name, "'");
+    if (!found) {
+      // Custom activity type requested — look up in the registry.
+      auto type_id = libkineto::CustomActivityTypeRegistry::instance().lookup(name);
+      TORCH_CHECK(
+          type_id.has_value(),
+          "Unknown or non-member activity type name: '",
+          name,
+          "'");
+
+      result.insert(static_cast<libkineto::ActivityType>(*type_id));
    }
  }
  return result;
}

Previously, the custom ID was validated but deliberately omitted from the result set, with a comment stating "out-of-range enum values may cause undefined behaviour in existing Kineto/backend code". In practice, static_cast<ActivityType>(custom_id) round-trips correctly through all Kineto code paths because ActivityType is a scoped enum with int underlying type, and all comparison/serialization logic operates on the integer value rather than assuming it falls within the static enum range.

2. Custom type insertion in prepareTrace():

After insertActivities(PrivateUse1, kPrivateUse1Types), the code now conditionally adds custom types:

• No filter specified: add all custom types from the registry (user gets every custom type).
• Filter specified: add only custom types whose type_name appears in the filter set. filterActivities() may have already added some; duplicates are skipped via k_activities.count(type_act).

Device classification: deviceTypeFromActivity()

deviceTypeFromActivity() in kineto_shim.cpp now checks activity_type >= kCustomActivityTypeBase and maps those IDs to c10::DeviceType::PrivateUse1, ensuring custom activities are correctly attributed without additional plumbing.

Registration Lifecycle

Registration is valid from process start until prepareTrace(). Between prepareTrace() and stopTrace(), calls log an error and return a sentinel {id: 0}. Between sessions, re-registration is idempotent (same input → same ID). This matches the lifecycle of REGISTER_PRIVATEUSE1_PROFILER.

Backwards Compatibility

• All existing ActivityType values and numeric assignments unchanged.
• Backends that register no custom types see no behavior change.

---

Looking forward to community feedback!

cc @divyanshk @briancoutinho @fffrog