boxing相关的概念在这篇博客里也讲的比较详细：https://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/
在具体看boxing/unboxing实现之前，需要先看一下为什么会有这两种调用方式，这两种调用方式的示例，以及为什么需要支持相互调用
这里有一个概念要分清，就是博客中提到的API和kernel，我们会有box/unbox api，以及box/unbox kernel
- Unbox api -> unbox kernel就是性能最快的路径，没有额外开销，类型信息都是静态的
-
通过Box api调用，比如调用custom op等，目的是为任何算子都生效，所以不会假设有unbox路径的代码（比如codegen出来的）
- 调用box kernel/unbox kernel取决于具体的算子
- 调用到box kernel，比如fallback kernel，处理任意类型的输入
claude给总结了一下torch中常见的调用链路，API层面的：
| 你在 Python 里这么写 | 走的是哪条路 |
|---|---|
torch.add(x, y)、x + y、x.relu()、F.linear(...) 等所有"正常"API |
Unboxed |
torch.ops.aten.add(x, y) / torch.ops.mylib.foo(...) |
Boxed |
__torch_dispatch__ 拦截下来再 dispatch |
Boxed |
vmap / functorch / functionalize 包住的 op |
Boxed（fallback 是 boxed） |
torch.compile / Dynamo / AOTAutograd 在 trace 期间 |
Boxed（fake tensor 那一套 + Python dispatch） |
torch.compile 编出来的 Inductor 代码运行时 |
Unboxed（codegen 的 C++/Triton 调用） |
上面说的主要是API，然后从kernel层面看，绝大多数的kernel都是unboxed，比如codegen出来的aten的算子们。对于boxed kernel，主要就是两种情况：
- Fallback kernel，上面已经提过了，需要处理任意类型的输入，不是静态的
-
python层的kernel，是运行时的产物，比如
@torch.library.impl("mylib::foo", "CUDA") def foo_cuda_impl(x): return x.sin()
对应下表，看一下api/kernel的调用示例：
| 调用方式 | kernel 是 unboxed | kernel 是 boxed |
|---|---|---|
C++ at::add(...) (unboxed) |
直调 unboxed | unboxed→boxed 转换 (BoxedKernelWrapper) |
torch.add(x, y) (Python，但 codegen 走 unboxed) |
直调 unboxed | unboxed→boxed 转换 |
torch.ops.aten.add(...) (boxed) |
boxed→unboxed 转换 (make_boxed_from_unboxed_functor) |
直调 boxed |
注册链路
整体算子的注册链路在上一篇文章已经写了，然后这里主要是看一下boxing/unboxing相关的处理。然后针对unboxing，看一下是怎么在编译期做的这个事情
整体上看分为这几层，上次主要关注的是registerImpl之后，dispatch table相关的
```C++
┌─────────────────────────────────────────────────────────────────┐
│ 用户代码层 │
│ TORCH_LIBRARY / TORCH_LIBRARY_IMPL │
│ m.def(...) / m.impl(...) / m.fallback(...) │
└──────────────────────────────┬──────────────────────────────────┘
▼
┌─────────────────────────────────────────────────────────────────┐
│ torch::CppFunction 层 (torch/library.h) │
│ 构造函数把不同形式的"函数"包成统一的 KernelFunction │
│ 不同重载决定走 makeFromUnboxed* 还是 makeFromBoxed* │
└──────────────────────────────┬──────────────────────────────────┘
▼
┌─────────────────────────────────────────────────────────────────┐
│ c10::KernelFunction 层 (KernelFunction.h) │
│ 持有 boxed_kernel_func_ + unboxed_kernel_func_ │
│ makeFromUnboxed* → 自动生成 boxed wrapper │
│ makeFromBoxed* → unboxed 指针保持 nullptr │
└──────────────────────────────┬──────────────────────────────────┘
▼
┌─────────────────────────────────────────────────────────────────┐
│ Library::_impl / Library::_fallback (library.cpp) │
│ → Dispatcher::registerImpl / registerFallback │
└──────────────────────────────┬──────────────────────────────────┘
▼
┌─────────────────────────────────────────────────────────────────┐
│ OperatorEntry::registerKernel (OperatorEntry.cpp) │
│ 把 KernelFunction 塞进 kernels_[dispatch_key] │
│ 触发 updateDispatchTable_() 重建 dispatch table │
└─────────────────────────────────────────────────────────────────┘
<pre><code class="line-numbers"><br>
在CppFunction层：
- 对于unboxed kernel来说，会根据CppFunction构造函数的不同重载来决定是走那条链路，比如makeFromUnboxedFunction/makeFromUnboxedLambda等
- 对于boxed kernel来说，需要主要调用CppFunction::makeFromBoxedFunctor等函数，或者是通过fallthrough自动构造
<br>
最终都会调用到makeFromUnboxedFunctor和makeFromBoxedKernel这两个位置上
- makeFromUnboxedFunctor会自动生成boxed call。
- 而makeFromBoxedKernel则只会保存传入的boxed kernel，unboxed kernel则是空
## makeFromUnboxedFunctor
看调用路径：
```C++
m.impl("foo", &foo_cpu) ← 普通函数指针
→ CppFunction(Func*)
→ makeFromUnboxedRuntimeFunction
→ WrapFunctionIntoRuntimeFunctor ← 把函数指针包成 functor
↘
m.impl("foo", TORCH_FN(foo_cpu)) ← 编译期函数指针 (codegen 用这个)
→ CppFunction(FuncPtr)
→ makeFromUnboxedFunction
→ WrapFunctionIntoFunctor ← 把编译期函数指针包成 functor
↘
m.impl("foo", [](Tensor x) {...}) ← lambda
→ CppFunction(Lambda&&)
→ makeFromUnboxedLambda
→ WrapFunctionIntoRuntimeFunctor ← lambda 已经是 functor，再包一层
↘
makeFromUnboxedFunctor ← ★ 这里 ★
↗
m.impl("foo", std::make_unique<MyFunctor>()) ← 手写 functor (罕见)
↗
进到makeFromUnboxedFunctor里面，会wrap一下这里的KernelFunctor
```C++
auto* unboxed_fn = &impl::wrap_kernel_functor_unboxed<KernelFunctor>::call;
void* void_unboxed_fn = reinterpret_cast<void*>(unboxed_fn);
<pre><code class="line-numbers">这里传入的functor是OperatorKernel类型的，真正调用的时候会根据Functor的类型转化回去，包括需要的args
即这里的unboxed\_fn的作用是做一下static\_cast，把上面wrap的OperatorKernel转回真实的类型，然后调用 operator()
然后unboxed\_fn的类型再被擦除掉成为void\*。
理解这里的核心点就是：把void\*转化回真实kernel并调用的这个操作，也是在编译器生成的函数，后续只需要保存这个函数指针即可。
<br>
通过unbox生成box kernel的路径：make\_boxed\_from\_unboxed\_functor，核心点就是：
- 参数类型对齐BoxedKernel::InternalBoxedKernelFunction
- 也是类型擦除后的，是一个通用类型
- call\_functor\_with\_args\_from\_stack生成从stack中解析参数的代码，然后调用到原本的Functor中
- 这一段代码同样是per kernel预先codegen出来的。所以后续调用的时候也不需要有具体的类型。
<br>
这里再提一下，unbox代码在有类型擦除的情况下，还能做到unbox的原理就是根据每个类型的kernel生成void\*到kernel type的cast代码，运行时保存这个cast的代码即可。
<br>
```C++
KernelFunction
├── boxed_kernel_func_ : BoxedKernel
│ ├── functor_ : intrusive_ptr<OperatorKernel> ──→ KernelFunctor 实例
│ └── boxed_kernel_func_ : InternalBoxedKernelFunction*
│ ──→ make_boxed_from_unboxed_functor<KernelFunctor>::call
│ (自动生成的 boxed wrapper)
├── unboxed_kernel_func_ : void*
│ ──→ wrap_kernel_functor_unboxed<KernelFunctor>::call
│ (强类型 unboxed wrapper，已 reinterpret_cast)
└── sym_unboxed_kernel_func_ : void* (要么这个填，要么上面填，二者必有一空)
wrap好的kernel（lambda，函数指针等），都是保存在OperatorKernel中， 都放到了BoxedKernel下面
然后有不同链路的裸指针，boxed_kernel_func_/unboxed_kernel_func_等，对应上面说的cast的代码，最终调用到OperatorKernel里
makeFromBoxedKernel
实现上没啥可以看的，主要就是看用户传来的boxed function的构造流程。这个可以看后面的case study
调用链路
其中，boxed api调用unbox kernel在上面已经有展示了，就是在注册的时候生成好了从stack中解析参数的代码，是在编译时发生的。
运行时通过callBoxed来调用
对于unbox api调用boxed kernel，直观理解就是把unbox api中的参数写到栈中，给boxed kernel使用。
具体的实现就在KernelFunction::call中，在unbox api调用时，如果发现没有unbox kernel，就会调用到Boxed fallback中
BoxedKernelWrapper帮忙处理了这块，把参数套到了栈中，然后再调用的callBoxed
额外提一嘴就是这里的栈就是vector
- IValue就是一个tag标记类型，一个union保存各种类型的数值
- 数值有三种，POD类型，就是int什么的。然后有一些堆上的复杂结构，通过指针存，比如string，tuple等。还有就是tensor会单独处理，应该是为了做一些性能优化，细节我这里也没看
然后还有一个点就是unbox api是静态的，所以unbox api fallback到box kernel这段逻辑也是静态的，所以把参数push到栈中的这段代码是编译期生成的，没有解释的开销
CaseStudy: CustomOp
然后这里来看一下常用的custom op是怎么注册进来，以及怎么调用的
用户 Python 代码 Python 侧
┌───────────────────────────────────────────────────────────┐
│ @torch.library.custom_op("mylib::foo", mutates_args=()) │
│ def foo(x: Tensor) -> Tensor: ... │
│ │
│ 或更底层: │
│ lib = torch.library.Library("mylib", "FRAGMENT") │
│ lib.impl("foo", foo_cpu_fn, "CPU") │
└───────────────────────┬───────────────────────────────────┘
│
▼
┌───────────────────────────────────────────────────────────┐
│ Library.impl (torch/library.py:425) │
│ self.m.impl(name, dispatch_key, fn, with_keyset) │
│ self.m = torch._C._dispatch_library(...) 返回的 C++ obj │
└───────────────────────┬───────────────────────────────────┘
│ pybind 边界
▼
┌───────────────────────────────────────────────────────────┐ C++ 侧
│ python_dispatch.cpp 的 "impl" lambda (370 行) │
│ if func is fallthrough_kernel: │
│ lib.impl(name, dispatch(dk, makeFallthrough())) │
│ else: │
│ lib.impl(name, dispatch(dk, │
│ makeFromBoxedFunctor( │
│ unique_ptr<PythonKernelHolder>(func, dk, ...)) │
│ )) │
└───────────────────────┬───────────────────────────────────┘
python层的注册，调用到lib.impl的时候，这里会有一个判断，对于非fallthrough kernel，会通过PythonKernelHolder来保存python层的kernel，然后生成BoxedFunctor，后面就是registerImpl那些了
```C++
class PythonKernelHolder : public c10::OperatorKernel {
c10::SafePyObject func_;
c10::DispatchKey dispatch_key_;
// If "with_keyset", then we expect a keyset as the first arg.
bool with_keyset_;
// If "with_op", then we expect the op as first arg (or second if keyset)
bool with_op_;
}
with\_keyset/with\_op主要是函数的调用规约。这里核心就是SafePyObject func\_对应python层的对象。
PythonKernelHolder的operator()就对应了kernel的调用，这里会：
- pop栈上的ivalue
- 然后获取GIL（需要在GIL里构造对象）
- 代码中的`py::gil_scoped_acquire g`
- parseIValuesToPyArgsKwargs，把IValue翻译成python的args, kwargs
- 调用python层的函数，如果有输出的话，会转化成IValue保存到栈中
<br>
然后看调用链路：
```C++
torch.ops.mylib.foo(x) # Python
│
▼
OpOverload.__call__ # torch/_ops.py:870
│ self._op(*args, **kwargs)
▼
pybind C++ binding (invokeOperatorFromPython) # pybind_utils.cpp:888
│ 1. getOpWithStack(): 解析 args 成 Stack<IValue>
│ 2. found_op->getOperation()(stack) ← BOXED 调用
▼
Operator::getOperation 返回的 lambda
│ → Dispatcher::callBoxed(op, stack) # Dispatcher.h:851
▼
getDispatchKeySetBoxed(stack) # DispatchKeyExtractor.h:146
│ 从 stack peek tensor 计算 dispatch key set
▼
OperatorEntry::lookup(ks) # 命中 PythonKernelHolder 那一格
▼
KernelFunction::callBoxed # KernelFunction_impl.h:87
│ boxed_kernel_func_.callBoxed(op, ks, stack)
▼
BoxedKernel::callBoxed # BoxedKernel_impl.h:41
│ (*boxed_kernel_func_)(functor_.get(), op, ks, stack)
│ ↑ 这里 functor_ 是 PythonKernelHolder
▼
BoxedKernel::makeFromFunctor 时生成的那个 lambda
│ (*static_cast<PythonKernelHolder*>(kernel))(op, ks, stack)
▼
PythonKernelHolder::operator() # python_dispatch.cpp:100
│
├─ 1. 检查 TorchDispatchMode 栈 → 命中则走 mode interpreter
├─ 2. 检查输入 tensor 是否带 Python key → 命中则走 tensor 的 interpreter
└─ 3. 直接路径：
pop stack → parseIValuesToPyArgsKwargs
→ py::gil_scoped_acquire
→ func_(*args, **kwargs) ← 调用用户的 Python foo()
→ pushPyOutToStack
</code></pre>
cpp/python层的分界，调用到invokeOperatorFromPython，把参数放到stack中，然后到dispatcher层。后面的就都一样了。
<br>
再来看看py层发生了什么：
```C++
torch.ops.mylib.foo(x) ── Python 表达式
│
│ (Python 属性查找，前两次创建缓存)
│
├── torch.ops ── _Ops 实例
│ __getattr__("mylib")
│
├── torch.ops.mylib ── _OpNamespace 实例
│ __getattr__("foo")
│ │
│ └── torch._C._jit_get_operation("mylib::foo")
│ │
│ └── pybind C++:
│ - getAllSortedOperatorsFor(symbol)
│ - 返回 py::cpp_function(闭包持有 operators)
│
├── torch.ops.mylib.foo ── OpOverloadPacket 实例
│ __call__(x)
│ │
│ └── self._op(x, **kwargs) ── 调用步骤里的 cpp_function
│ │
└── pybind 边界 ─────────────────────────────────────────
▼
┌──────────────────────────────────────────────────────┐
│ Lambda 闭包 (init.cpp:1823) │
│ ToIValueAllowNumbersAsTensors guard │
│ _get_operation_for_overload_or_packet(...) │
└──────────────────────────────┬───────────────────────┘
▼
┌──────────────────────────────────────────────────────┐
│ _get_operation_for_overload_or_packet │
│ (pybind_utils.cpp:992) │
│ if 有 __torch_function__: 走 override │
│ else: invokeOperatorFromPython │
└──────────────────────────────┬───────────────────────┘
▼
┌──────────────────────────────────────────────────────┐
│ invokeOperatorFromPython (pybind_utils.cpp:888) │
│ 1. getOpWithStack: │
│ - 按 args/kwargs 匹配 schema 选 overload │
│ - 把 args/kwargs 转成 Stack ← boxing │
│ 2. py::gil_scoped_release │
│ 3. found_op->getOperation()(stack) │
│ 4. createPyObjectForStack(stack) ← unbox 返回值
└──────────────────────────────┬───────────────────────┘
▼
┌──────────────────────────────────────────────────────┐
│ Operator::getOperation()(stack) │
│ (这是 JIT 包装层，真正调到 dispatcher) │
│ → Dispatcher::callBoxed(op, &stack) │
└──────────────────────────────┬───────────────────────┘
▼
剩下进入 dispatcher boxed 路径
（我们之前看过的链路）
```
getAllSortedOperatorsFor会把所有的operator都取出来，用来处理重载。
在getOpWithStack里面，会根据schema选出对应的operator，再调用到dispatcher层
不过这里还处理了_torch_function_的链路，有点复杂，有一个直观的概念就好
文章评论