// c10/core/TensorImpl.h
struct C10_API TensorImpl : public c10::intrusive_ptr_target {
    // ...
    DispatchKeySet key_set_;  // 该 Tensor 的调度键集合
    // ...
};

// c10/core/DispatchKey.h (line 65-103)
enum class BackendComponent : uint8_t {
    InvalidBit = 0,
    CPUBit,
    CUDABit,
    HIPBit,
    XLABit,
    MPSBit,
    IPUBit,
    XPUBit,
    HPUBit,
    VEBit,
    LazyBit,
    MTIABit,
    MAIABit,
    PrivateUse1Bit,
    PrivateUse2Bit,
    PrivateUse3Bit,
    MetaBit,

    EndOfBackendKeys = MetaBit,
};

// c10/core/DispatchKey.h (line 567-570)
static_assert(
    static_cast<uint8_t>(BackendComponent::EndOfBackendKeys) <= 16,
    "BackendComponent currently only supports <= 16 backends.");

// c10/core/DispatchKey.h (line 57-63)
#define C10_FORALL_FUNCTIONALITY_KEYS(_) \
  _(Dense, )                             \
  _(Quantized, Quantized)                \
  _(Sparse, Sparse)                      \
  _(SparseCsr, SparseCsr)                \
  _(NestedTensor, NestedTensor)          \
  _(AutogradFunctionality, Autograd)

// c10/core/DispatchKey.h (line 222-421) 摘选
BackendSelect,       // 在无 tensor 参数时选择正确后端
Python,              // Python 层调度
Fake,                // Fake Tensor 支持
Functionalize,       // 别名和变异消除
ADInplaceOrView,     // inplace/view 操作的自动求导准备
FuncTorchBatched,    // vmap 批处理
FuncTorchGradWrapper,// functorch 梯度
AutocastCPU,         // CPU 自动混合精度
AutocastCUDA,        // CUDA 自动混合精度
PythonDispatcher,    // 跳过 C++ 调度器

// c10/core/DispatchKey.h (line 457-481)
Autograd,                                // 映射到所有 Autograd* 运行时 key
CompositeImplicitAutograd,               // 默认 autograd 公式
CompositeExplicitAutograd,               // 显式指定 autograd 的组合算子
CompositeExplicitAutogradNonFunctional,  // 非函数化版本
FuncTorchBatchedDecomposition,           // functorch 批处理分解

// c10/core/DispatchKeySet.h (line 167, 516)
class DispatchKeySet final {
 private:
    uint64_t repr_ = 0;
};

// c10/core/DispatchKeySet.h (line 206-250) 简化
constexpr explicit DispatchKeySet(DispatchKey k) {
    if (k == DispatchKey::Undefined) {
        // Case 1: 未定义 -> 空集合
        repr_ = 0;
    } else if (k <= DispatchKey::EndOfFunctionalityKeys) {
        // Case 2: 纯功能键 -> 只设置功能位
        repr_ = 1ULL << (num_backends + static_cast<uint8_t>(k) - 1);
    } else if (k <= DispatchKey::EndOfRuntimeBackendKeys) {
        // Case 3: 运行时 per-backend key -> 同时设置功能位和后端位
        auto functionality_k = toFunctionalityKey(k);
        uint64_t functionality_val = 1ULL << (num_backends + ...);
        auto backend_k = toBackendComponent(k);
        uint64_t backend_val = 1ULL << (static_cast<uint8_t>(backend_k) - 1);
        repr_ = functionality_val + backend_val;
    }
}

// c10/core/DispatchKeySet.h (line 338-339)
constexpr DispatchKeySet operator-(DispatchKeySet other) const {
    return DispatchKeySet(repr_ & (full_backend_mask | ~other.repr_));
}

// c10/core/DispatchKeySet.h (line 407-441)

// 获取最高优先级的功能键
DispatchKey highestFunctionalityKey() const {
    auto functionality_idx = indexOfHighestBit();
    if (functionality_idx < num_backends)
        return DispatchKey::Undefined;
    return static_cast<DispatchKey>(functionality_idx - num_backends);
}

// 获取最高优先级的后端组件
BackendComponent highestBackendKey() const {
    auto backend_idx =
        DispatchKeySet(repr_ & full_backend_mask).indexOfHighestBit();
    if (backend_idx == 0)
        return BackendComponent::InvalidBit;
    return static_cast<BackendComponent>(backend_idx);
}

// 获取最高优先级的运行时 DispatchKey
DispatchKey highestPriorityTypeId() const {
    auto functionality_k = highestFunctionalityKey();
    if (isPerBackendFunctionalityKey(functionality_k)) {
        return toRuntimePerBackendFunctionalityKey(
            functionality_k, highestBackendKey());
    }
    return functionality_k;
}

// c10/core/impl/LocalDispatchKeySet.h (line 31-51)
struct C10_API PODLocalDispatchKeySet {
    uint64_t included_;   // 额外包含的 key（XOR 编码）
    uint64_t excluded_;   // 额外排除的 key（XOR 编码）

    DispatchKeySet included() const {
        return DispatchKeySet(DispatchKeySet::RAW, included_) ^
            c10::default_included_set;
    }
    DispatchKeySet excluded() const {
        return DispatchKeySet(DispatchKeySet::RAW, excluded_) ^
            c10::default_excluded_set;
    }
};

// 1. IncludeDispatchKeyGuard - 临时包含一个 key
{
    c10::impl::IncludeDispatchKeyGuard guard(DispatchKey::AutocastCUDA);
    // 在此范围内，AutocastCUDA 被包含在 TLS included 中
    // 退出时自动恢复
}

// 2. ExcludeDispatchKeyGuard - 临时排除一个 key（最常用）
{
    c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::AutogradCPU);
    // Autograd 处理完后用此 guard 排除自身，避免无限递归
    // 退出时自动恢复
}

// 3. ForceDispatchKeyGuard - 强制设置整个 TLS 状态
{
    c10::impl::ForceDispatchKeyGuard guard(include_set, exclude_set);
    // 完全覆盖当前 TLS，退出时恢复之前的状态
}

// c10/core/DispatchKey.h (line 687-722)
constexpr DispatchKey toRuntimePerBackendFunctionalityKey(
    DispatchKey functionality_k,
    BackendComponent backend_k) {
    if (functionality_k == DispatchKey::Dense) {
        return static_cast<DispatchKey>(
            static_cast<uint8_t>(DispatchKey::StartOfDenseBackends) +
            static_cast<uint8_t>(backend_k));
    }
    // Sparse, SparseCsr, Quantized, NestedTensor, AutogradFunctionality
    // 类似处理...
}

// c10/core/DispatchKey.h (line 617-657)
constexpr BackendComponent toBackendComponent(DispatchKey k) {
    if (k >= DispatchKey::StartOfDenseBackends &&
        k <= DispatchKey::EndOfDenseBackends) {
        return static_cast<BackendComponent>(
            static_cast<uint8_t>(k) -
            static_cast<uint8_t>(DispatchKey::StartOfDenseBackends));
    }
    // 对 Quantized, Sparse, SparseCsr, NestedTensor,
    // AutogradFunctionality 各区间做相同处理...
}

// c10/core/DispatchKey.h (line 659-677)
constexpr DispatchKey toFunctionalityKey(DispatchKey k) {
    if (k <= DispatchKey::EndOfFunctionalityKeys) return k;
    else if (k <= DispatchKey::EndOfDenseBackends) return DispatchKey::Dense;
    else if (k <= DispatchKey::EndOfQuantizedBackends) return DispatchKey::Quantized;
    else if (k <= DispatchKey::EndOfSparseBackends) return DispatchKey::Sparse;
    // ...
}

// c10/core/DispatchKey.h (line 585-586)
constexpr uint16_t num_runtime_entries = num_functionality_keys +
    (numPerBackendFunctionalityKeys() * (num_backends - 1));

// c10/core/DispatchKeySet.h (line 22-35)
struct FunctionalityOffsetAndMask {
    uint16_t offset;  // 在操作符表中的起始偏移
    uint16_t mask;    // 后端位掩码
};

inline DispatchKeySet computeDispatchKeySet(
    DispatchKeySet ks,
    DispatchKeySet key_mask) {
  c10::impl::LocalDispatchKeySet local =
      c10::impl::tls_local_dispatch_key_set();
  return (((ks | local.included_) - local.excluded_) & key_mask);
}

import torch

# CPU 普通 Tensor
t_cpu = torch.randn(3, 4)
print(t_cpu._dispatch_key_set())
# DispatchKeySet(CPU, Dense)

# CUDA Tensor
if torch.cuda.is_available():
    t_cuda = t_cpu.cuda()
    print(t_cuda._dispatch_key_set())
    # DispatchKeySet(CUDA, Dense)

# 需要梯度的 Tensor
t_grad = torch.randn(3, 4, requires_grad=True)
print(t_grad._dispatch_key_set())
# DispatchKeySet(CPU, AutogradCPU, Dense)

# 稀疏 Tensor
t_sparse = t_cpu.to_sparse()
print(t_sparse._dispatch_key_set())
# DispatchKeySet(CPU, SparseCPU, Sparse)

import torch

# 查看某个 dispatch key 下注册的所有操作
registrations = torch._C._dispatch_get_registrations_for_dispatch_key("CPU")
print(f"CPU 上注册了 {len(registrations)} 个操作")
for reg in list(registrations)[:5]:
    print(f"  {reg}")

# 检查特定操作是否在特定 key 上有内核
has_cpu = torch._C._dispatch_has_kernel_for_dispatch_key(
    "aten::add.Tensor", "CPU")
print(f"aten::add.Tensor 在 CPU 上有内核: {has_cpu}")

has_autograd = torch._C._dispatch_has_kernel_for_dispatch_key(
    "aten::add.Tensor", "AutogradCPU")
print(f"aten::add.Tensor 在 AutogradCPU 上有内核: {has_autograd}")

import torch

# 打印 aten::add.Tensor 的完整调度表
print(torch._C._dispatch_dump("aten::add.Tensor"))