// aten/src/ATen/core/dispatch/Dispatcher.cpp:18
bool show_dispatch_trace() {
  static auto envar = c10::utils::get_env("TORCH_SHOW_DISPATCH_TRACE");
  if (envar.has_value()) {
    if (envar == "0") return false;
    if (envar == "1") return true;
    TORCH_WARN("ignoring invalid value for TORCH_SHOW_DISPATCH_TRACE: ...");
  }
  return false;
}

// aten/src/ATen/core/dispatch/Dispatcher.cpp:70
void _print_dispatch_trace(const std::string& label,
                           const std::string& op_name,
                           const DispatchKeySet& dispatchKeySet) {
  auto nesting_value = dispatch_trace_nesting_value();
  for (int64_t i = 0; i < nesting_value; ++i) std::cerr << ' ';
  std::cerr << label << " op=[" << op_name
            << "], key=[" << toString(dispatchKeySet.highestPriorityTypeId())
            << ']' << std::endl;
}

// aten/src/ATen/core/dispatch/Dispatcher.h:32
struct DispatchTraceNestingGuard {
  DispatchTraceNestingGuard() { dispatch_trace_nesting_incr(); }
  ~DispatchTraceNestingGuard() { dispatch_trace_nesting_decr(); }
};

import torch

# 列出 CPU dispatch key 上注册的所有算子
cpu_ops = torch._C._dispatch_get_registrations_for_dispatch_key("CPU")
print(f"CPU 上注册了 {len(cpu_ops)} 个算子")
print(cpu_ops[:5])
# 输出示例:
# CPU 上注册了 1842 个算子
# ['aten::abs', 'aten::abs.out', 'aten::acos', 'aten::acos.out', 'aten::add.Tensor']

# 检查 aten::mm 是否有 CPU 内核
has_cpu = torch._C._dispatch_has_kernel_for_dispatch_key(
    "aten::mm", "CPU"
)
print(f"aten::mm 有 CPU 内核: {has_cpu}")  # True

# 检查 aten::mm 是否有 Meta 内核
has_meta = torch._C._dispatch_has_kernel_for_dispatch_key(
    "aten::mm", "Meta"
)
print(f"aten::mm 有 Meta 内核: {has_meta}")  # True

# 打印 aten::mm 的注册状态
print(torch._C._dispatch_dump("aten::mm"))

print(torch._C._dispatch_dump_table("aten::mm"))

# 查看当前线程的 include/exclude dispatch key 集合
include_set = torch._C._dispatch_tls_local_include_set()
exclude_set = torch._C._dispatch_tls_local_exclude_set()
print(f"TLS include: {include_set}")
print(f"TLS exclude: {exclude_set}")

# 在 no_grad 上下文中观察 TLS 变化
with torch.no_grad():
    exclude_set_no_grad = torch._C._dispatch_tls_local_exclude_set()
    print(f"no_grad 下 TLS exclude: {exclude_set_no_grad}")
    # AutogradCPU, AutogradCUDA 等会被加入 exclude 集合

a = torch.randn(2, 3)
print(torch._C._dispatch_keys(a))
# 输出: DispatchKeySet(CPU, AutogradCPU)

b = torch.randn(2, 3, device="cuda")
print(torch._C._dispatch_keys(b))
# 输出: DispatchKeySet(CUDA, AutogradCUDA)

import torch
a = torch.randn(2, 3, device="cuda", requires_grad=True)
b = torch.randn(3, 4, device="cuda")
c = torch.mm(a, b)  # 调用入口

// torch/csrc/autograd/generated/VariableType_4.cpp (生成代码，简化版)
at::Tensor mm(c10::DispatchKeySet ks, const at::Tensor & self, const at::Tensor & mat2) {
  // 1. 检查是否需要计算梯度
  auto _any_requires_grad = compute_requires_grad(self, mat2);

  // 2. 创建 grad_fn 并保存反向传播所需的输入
  std::shared_ptr<MmBackward0> grad_fn;
  if (_any_requires_grad) {
    grad_fn = std::shared_ptr<MmBackward0>(new MmBackward0(), deleteNode);
    grad_fn->self_ = SavedVariable(self, false);
    grad_fn->mat2_ = SavedVariable(mat2, false);
  }

  // 3. 排除 autograd key，redispatch 到下一层
  auto _tmp = ([&]() {
    at::AutoDispatchBelowADInplaceOrView guard;
    return at::redispatch::mm(ks & c10::after_autograd_keyset, self, mat2);
  })();

  auto result = std::move(_tmp);

  // 4. 将 grad_fn 关联到输出
  if (grad_fn) {
    set_history(flatten_tensor_args(result), grad_fn);
  }

  return result;
}

// aten/src/ATen/core/LegacyTypeDispatch.h:61
struct TORCH_API AutoDispatchBelowAutograd {
  AutoDispatchBelowAutograd() :
    autograd_guard_(c10::autograd_dispatch_keyset) {}
  c10::impl::ExcludeDispatchKeyGuard autograd_guard_;
};

// c10/core/DispatchKeySet.h:731
constexpr DispatchKeySet after_autograd_keyset =
    DispatchKeySet(DispatchKeySet::FULL_AFTER, c10::DispatchKey::AutogradOther);

import torch
from torch.library import Library, impl

# 定义算子 schema
lib = Library("myops", "DEF")
lib.define("my_add(Tensor x, Tensor y) -> Tensor")

# 注册 CPU 内核
@impl(lib, "my_add", "CPU")
def my_add_cpu(x, y):
    print("  [my_add_cpu] 进入 CPU 内核")
    return x + y

# 注册 Autograd 内核（使用 CompositeImplicitAutograd 或 AutogradCPU）
@impl(lib, "my_add", "AutogradCPU")
def my_add_autograd(x, y):
    print("  [my_add_autograd] 进入 AutogradCPU 内核")
    # 必须 redispatch，否则会无限递归
    return torch.ops.myops.my_add(x, y)

# 检查分发表
print(torch._C._dispatch_dump("myops::my_add"))

# 检查各 key 是否有内核
for key in ["CPU", "CUDA", "AutogradCPU", "AutogradCUDA", "Meta"]:
    has = torch._C._dispatch_has_kernel_for_dispatch_key("myops::my_add", key)
    print(f"  {key}: {'已注册' if has else '未注册'}")

# 错误: 只注册了 CPU，但在 CUDA tensor 上调用
a_cuda = torch.randn(3, device="cuda")
b_cuda = torch.randn(3, device="cuda")
torch.ops.myops.my_add(a_cuda, b_cuda)
# RuntimeError: Could not run 'myops::my_add' with arguments from the 'CUDA' backend.

# 第 1 步：确认算子名称
print(torch._C._dispatch_has_kernel("aten::xxx"))

# 第 2 步：查看已注册的 key
print(torch._C._dispatch_dump("aten::xxx"))

# 第 3 步：检查具体缺失的 key
torch._C._dispatch_has_kernel_for_dispatch_key("aten::xxx", "YYY")

# 检查当前 TLS 状态
include = torch._C._dispatch_tls_local_include_set()
exclude = torch._C._dispatch_tls_local_exclude_set()
print(f"include: {include}")
print(f"exclude: {exclude}")

# 如果 exclude 集合中意外包含了某个 key，
# 说明可能有 RAII guard 没有正确释放

# 验证 autocast key 是否在分发链中
with torch.autocast(device_type="cuda"):
    include = torch._C._dispatch_tls_local_include_set()
    print(f"autocast 下 include: {include}")
    # 应该看到 AutocastCUDA 在 include 集合中