// aten/src/ATen/core/dispatch/Dispatcher.h, 第 112 行
C10_ALWAYS_INLINE static Dispatcher& singleton() {
#if !defined C10_MOBILE
    static Dispatcher& s = realSingleton();
    return s;
#else
    return realSingleton();
#endif
}

// aten/src/ATen/core/dispatch/Dispatcher.cpp, 第 87 行
C10_EXPORT Dispatcher& Dispatcher::realSingleton() {
    static Dispatcher _singleton;
    return _singleton;
}

// aten/src/ATen/core/dispatch/Dispatcher.h, 第 400-435 行
class Dispatcher final {
    // 所有已注册算子的链表
    std::list<OperatorDef> operators_;

    // 算子名称 → OperatorHandle 的并发安全映射
    LeftRight<ska::flat_hash_map<OperatorName, OperatorHandle>>
        operatorLookupTable_;

    // 每个 DispatchKey 的后备内核
    std::array<impl::AnnotatedKernel, num_runtime_entries>
        backendFallbackKernels_;

    // 已注册的库命名空间
    ska::flat_hash_map<std::string, std::string> libraries_;

    // 注册监听器列表
    std::unique_ptr<detail::RegistrationListenerList> listeners_;

    // 多解释器同步用的条件变量
    std::condition_variable cond_var_;

    // 生命周期守卫
    std::shared_ptr<Guard> guard_;
};

// aten/src/ATen/core/dispatch/Dispatcher.h, 第 76 行
struct OperatorDef final {
    impl::OperatorEntry op;   // 真正的算子注册表
    size_t def_count = 0;     // schema 注册计数
    size_t def_and_impl_count = 0;  // schema + impl 注册计数
};

// aten/src/ATen/core/dispatch/Dispatcher.h, 第 569-591 行
class OperatorHandle {
    // 直接指针，避免迭代器解引用的额外指令
    Dispatcher::OperatorDef* operatorDef_;

    // 保留迭代器，用于 cleanup() 时的快速删除
    std::list<Dispatcher::OperatorDef>::iterator operatorIterator_;
};

// aten/src/ATen/core/dispatch/Dispatcher.h, 第 605-633 行
template <class Return, class... Args>
class TypedOperatorHandle<Return(Args...)> final : public OperatorHandle {
    C10_ALWAYS_INLINE Return call(Args... args) const {
        return c10::Dispatcher::singleton().call<Return, Args...>(
            *this, std::forward<Args>(args)...);
    }

    C10_ALWAYS_INLINE Return
    redispatch(DispatchKeySet currentDispatchKeySet, Args... args) const {
        return c10::Dispatcher::singleton().redispatch<Return, Args...>(
            *this, currentDispatchKeySet, std::forward<Args>(args)...);
    }
};

// 方式一：按名称查找，未找到则抛异常
OperatorHandle op = Dispatcher::singleton()
    .findSchemaOrThrow("aten::add", "Tensor");

// 方式二：按名称查找，返回 optional
auto op = Dispatcher::singleton()
    .findOp(OperatorName("aten::add", "Tensor"));

// aten/src/ATen/core/dispatch/OperatorEntry.h, 第 232-283 行
class OperatorEntry final {
    OperatorName name_;                    // 算子名称
    std::optional<AnnotatedSchema> schema_; // 算子 schema（可能延迟注册）

    // === 热路径数据 ===
    std::array<KernelFunction, num_runtime_entries> dispatchTable_;
        // 分发表：DispatchKey 索引 → KernelFunction
    DispatchKeyExtractor dispatchKeyExtractor_;
        // 键提取器：从参数中提取 DispatchKeySet

    // === 冷路径数据 ===
    ska::flat_hash_map<DispatchKey, std::list<AnnotatedKernel>> kernels_;
        // 已注册内核的完整记录（支持覆盖与卸载）
    PyHandleCache py_cache_;
        // Python 端 torch.ops 对象的缓存
};

// aten/src/ATen/core/dispatch/OperatorEntry.h, 第 37 行
struct AnnotatedKernel final {
    KernelFunction kernel;                              // 实际的内核函数
    std::unique_ptr<FunctionSchema> inferred_function_schema; // 推断的签名
    std::string debug;                                  // 调试信息（注册位置）
};

// aten/src/ATen/core/dispatch/Dispatcher.h, 第 776 行
template <class Return, class... Args>
C10_ALWAYS_INLINE_UNLESS_MOBILE Return Dispatcher::call(
    const TypedOperatorHandle<Return(Args...)>& op,
    Args... args) const {
  // Step 1: 提取 DispatchKeySet
  auto dispatchKeySet =
      op.operatorDef_->op.dispatchKeyExtractor()
          .template getDispatchKeySetUnboxed<Args...>(args...);

  // Step 2: 在分发表中查找内核
  const KernelFunction& kernel = op.operatorDef_->op.lookup(dispatchKeySet);

  // Step 3: 调用内核
  return kernel.template call<Return, Args...>(
      op, dispatchKeySet, std::forward<Args>(args)...);
}

// aten/src/ATen/core/dispatch/DispatchKeyExtractor.h, 第 260-277 行
struct DispatchKeyExtractor final {
    // 位集合：标记哪些参数是 Tensor（反序存储，用于栈操作）
    c10::utils::bitset dispatch_arg_indices_reverse_;

    // 掩码：排除 fallthrough 的功能键
    DispatchKeySet nonFallthroughKeys_;

    // 每个后端独立的 fallthrough 掩码（慢路径）
    std::array<DispatchKeySet, num_backends> nonFallthroughKeysPerBackend_;

    // 是否需要使用 per-backend 掩码
    bool requiresBitsetPerBackend_{false};
};

// aten/src/ATen/core/dispatch/DispatchKeyExtractor.h, 第 186 行
template <class... Args>
DispatchKeySet getDispatchKeySetUnboxed(const Args&... args) const {
    // (1) 从所有 Tensor 参数中 OR 合并 key_set
    auto ks = detail::multi_dispatch_key_set(args...);
    // (2) 结合 TLS 并应用 fallthrough 掩码
    return impl::computeDispatchKeySet(ks, nonFallthroughKeys_);
}

// aten/src/ATen/core/dispatch/DispatchKeyExtractor.h, 第 54 行
struct MultiDispatchKeySet : at::IterArgs<MultiDispatchKeySet> {
    DispatchKeySet ts;
    void operator()(const at::Tensor& x) {
        ts = ts | x.key_set();    // Tensor: 合并 key_set
    }
    void operator()(at::ArrayRef<at::Tensor> xs) {
        for (const auto& x : xs)  // TensorList: 遍历合并
            ts = ts | x.key_set();
    }
    template <typename T>
    void operator()(const T&) {}   // 非 Tensor: 忽略
};

// aten/src/ATen/core/dispatch/DispatchKeyExtractor.h, 第 24 行
inline DispatchKeySet computeDispatchKeySet(
    DispatchKeySet ks,
    DispatchKeySet key_mask) {
  LocalDispatchKeySet local = tls_local_dispatch_key_set();
  return (((ks | local.included_) - local.excluded_) & key_mask);
}

// aten/src/ATen/core/dispatch/OperatorEntry.h, 第 182 行
const KernelFunction& lookup(DispatchKeySet ks) const {
    const auto idx = ks.getDispatchTableIndexForDispatchKeySet();
    if (C10_UNLIKELY(idx == -1)) {
        reportError(ks.highestPriorityTypeId());
    }
    const auto& kernel = dispatchTable_[idx];
    if (C10_UNLIKELY(!kernel.isValidUnboxed())) {
        if (!kernel.isValid()) {
            reportError(ks.highestPriorityTypeId());
        }
    }
    return kernel;
}

// c10/core/DispatchKeySet.h, 第 488 行
int getDispatchTableIndexForDispatchKeySet() const {
    // (1) 提取最高优先级的 functionality 索引
    auto functionality_idx =
        DispatchKeySet(repr_ >> num_backends).indexOfHighestBit();

    // (2) 从预计算表中获取 offset 和 mask
    auto offset_and_mask = offsetsAndMasks()[functionality_idx];

    // (3) 提取 backend 索引
    auto backend_idx =
        DispatchKeySet((repr_ & offset_and_mask.mask) >> 1)
            .indexOfHighestBit();

    // (4) 最终索引 = offset + backend_idx
    return offset_and_mask.offset + backend_idx;
}

// aten/src/ATen/core/dispatch/Dispatcher.h, 第 835 行
template <class Return, class... Args>
inline Return Dispatcher::redispatch(
    const TypedOperatorHandle<Return(Args...)>& op,
    DispatchKeySet currentDispatchKeySet,
    Args... args) const {
  const KernelFunction& kernel =
      op.operatorDef_->op.lookup(currentDispatchKeySet);
  return kernel.template call<Return, Args...>(
      op, currentDispatchKeySet, std::forward<Args>(args)...);
}

// 由 tools/autograd/gen_variable_type.py 代码生成
at::Tensor autograd_mm(
    c10::DispatchKeySet ks,
    const at::Tensor& self,
    const at::Tensor& mat2) {
  // 保存梯度信息
  auto grad_fn = std::make_shared<MmBackward>();
  grad_fn->self_ = self;
  grad_fn->mat2_ = mat2;

  // 移除 autograd 键，继续分发
  auto result = at::redispatch::mm(
      ks & c10::after_autograd_keyset,  // 排除 Autograd 系列键
      self, mat2);

  // 将 grad_fn 附加到输出
  set_history(result, grad_fn);
  return result;
}

// aten/src/ATen/core/dispatch/Dispatcher.h, 第 853 行
inline void Dispatcher::callBoxed(
    const OperatorHandle& op, Stack* stack) const {
  const auto& entry = op.operatorDef_->op;
  // 从栈上提取 DispatchKeySet
  auto dispatchKeySet =
      entry.dispatchKeyExtractor().getDispatchKeySetBoxed(stack);
  // 查找内核
  const auto& kernel = entry.lookup(dispatchKeySet);
  // 通过 boxed 接口调用内核
  kernel.callBoxed(op, dispatchKeySet, stack);
}

// aten/src/ATen/core/dispatch/Dispatcher.h, 第 898 行
inline void Dispatcher::callBoxedForDispatchKey(
    const OperatorHandle& op, DispatchKey dk, Stack* stack) const {
  const auto& entry = op.operatorDef_->op;
  auto dispatchKeySet =
      entry.dispatchKeyExtractor().getDispatchKeySetBoxed(stack);
  const auto& kernel = ([&]() {
    if (op.hasKernelForDispatchKey(dk)) {
      return entry.kernelForDispatchKey(dk);
    } else {
      auto idx = getDispatchTableIndexForDispatchKey(dk);
      return backendFallbackKernels_[idx].kernel;
    }
  })();
  kernel.callBoxed(op, dispatchKeySet, stack);
}

// aten/src/ATen/core/dispatch/Dispatcher.h, 第 921 行
inline void Dispatcher::redispatchBoxed(
    const OperatorHandle& op,
    DispatchKeySet dispatchKeySet,
    Stack* stack) const {
  const auto& entry = op.operatorDef_->op;
  const auto& kernel = entry.lookup(dispatchKeySet);
  kernel.callBoxed(op, dispatchKeySet, stack);
}

import os
os.environ["TORCH_SHOW_DISPATCH_TRACE"] = "1"

import torch
a = torch.randn(3, 3, requires_grad=True)
b = torch.randn(3, 3, requires_grad=True)
c = a + b