Write one line of a + b, hit Enter, and you have a result. But between that line and the instructions that actually perform the addition sit a set of contracts, a stretch of code generation, and a runtime dispatcher.
The previous piece covered where tensor memory comes from, how it is shared, and how the version is guarded. This one covers how operators get running: you call an operator, and how the framework finds the kernel that should do the work. Put simply: the common logic is handed to the machine to generate, the differences between operators are written into the contract, and at run time the dispatch table knows the way. The person writing an operator only has to care about that line of math; parameter validation, device alignment, type promotion, gradient recording, version updates — that scutwork, the framework handles on the side.
Why Not Write Everything by Hand
The difference between pointwise operators is, in essence, one line of mathematical formula. But a usable operator is far more than that one line: before execution it has to validate the argument count, confirm that all tensors sit on the same device, do type promotion when needed, decide whether to record the autograd graph, and for in-place computation also update the version counter — while conveniently leaving a probe point for the Profiler along the way.
If all of that were handwritten, adding one relu would mean keeping four places in sync: declaring the method on the Tensor class, registering the kernel with the dispatcher, exposing the Python interface in the binding layer, and filling in the derivative formula in the automatic differentiation layer. Changing that many places at once is easy to miss — and this kind of bug does not crash. It just quietly calls the wrong kernel on some backend, and the answer comes out wrong.
TensorPlay hands the common part of those four places to the code generator, writes the differences into a single contract, and lets the dispatch table complete the run-time kernel routing on its own. To add an operator, there are only two places where the pen touches the paper: define the function signature in the contract file, then write one line of math for the matching hardware kernel.
The Contract: Two YAML Files Call the Shots
TensorPlay stores the information for all operators in two YAML files. The interface, the kernel dispatch, the gradient logic all grow out of these two files, with nothing handwritten alongside — one fewer place that has to be “remembered and kept in sync” is one fewer place to forget.
The first is config/native_functions.yaml, the operator’s “identity file”, spelling out the signature, the variants, and the hardware dispatch:
# config/native_functions.yaml
- func: convolution(Tensor input, Tensor weight, Tensor? bias, int[] stride,
int[] padding, int[] dilation, bool transposed,
int[] output_padding, int groups) -> Tensor
variants: function
dispatch:
CPU: convolution_cpu
CUDA: convolution_cuda
tags: core
The three fields each handle one thing. func is the complete signature; the C++ interface, the Python bindings, and the dispatch key table are all assembled from it. variants controls the interface form: function corresponds to a static factory method (such as Tensor::zeros), method to an instance method (such as tensor.view). dispatch names the kernel for each hardware backend; an operator with no CUDA/Vulkan configured falls back to the general-purpose Composite implementation by default.
The second, config/derivatives.yaml, handles gradients. It does not repeat the full signature; it only links the operator name and gives the derivative formula:
# config/derivatives.yaml
- name: add(Tensor self, Tensor other, Scalar alpha=1) -> Tensor
self: grad
other: grad * alpha
self and other are input names, grad is the gradient coming back from the reverse pass, and the right-hand side is the formula on that edge. The generator does not do string substitution — it compiles the formula as an expression, tokenizes it and runs precedence climbing to parse it into an AST, and one pass tells it which intermediates the backward node has to save. The node structure, the constructor arguments, and the grad_fn call site are therefore all derived from one single derivation, and cannot drift apart.
Change one line of YAML and the C++ interface, the Python bindings, the type hints, the gradient nodes, and the autocast registration all change along with it. This one file holds close to three thousand - func: records.
Note: every generated artifact begins with
// Generated by tools/codegen/main.py -- DO NOT EDIT. When the build complains, what you should edit is the YAML or the generator, not the generated file — otherwise the next build wipes out your hand edits.
The Generator: One Context, Many Outputs
The entry point of the generator is CodegenContext, holding the function list, the derivative table, the autocast set, and the output directory. Every generation module shares the same context and reads the same data, so parameter names, default values, and dispatch keys have a hard time disagreeing:
# tools/codegen/main.py
@dataclass
class CodegenContext:
funcs: list
native_by_opname: dict
derivatives: dict
autocast_ops: set[str]
autograd_ops: dict[str, str]
out_dir: str
pkg_out: str | None = None
...
The whole generation logic is split into a dozen-odd target modules, each owning one piece: TensorMethods writes the Tensor method declarations and implementations, Redispatch writes the redispatch header, AutogradNodes writes the backward nodes, Bindings writes the bindings, Autocast writes the autocast registration. One run drops a dozen-odd artifacts into the output directory.
# tools/codegen/main.py
DEFAULT_TARGETS = ["TensorMethods", "Redispatch", "AutogradNodes", "TPXOps",
"AutogradRegistration", "Bindings", "Autocast",
"InplaceOrView", "Structured", "PythonFunctional",
"PythonCAPI", "PythonDispatch"]
There is also an unassuming little function, _write_if_changed, which decides how many translation units actually get recompiled when you change a YAML file:
# tools/codegen/main.py
def _write_if_changed(path, content):
try:
with open(path) as fh:
if fh.read() == content:
print(f'Unchanged "{path}"')
return
except FileNotFoundError:
pass
with open(path, "w") as fh:
fh.write(content)
print(f'Generated "{path}"')
If the content has not changed it is not rewritten, the original timestamp stays, and only the modified .cpp files are recompiled — no need to disturb every artifact each time.
Layering: A Thin Tensor Plus Mixed-In Methods
The Tensor class is kept as thin as possible. Underneath, a homegrown intrusive reference-counted pointer, intrusive_ptr<TensorImpl> (the counter lives inside the object itself, not in a std::shared_ptr), holds one block of tensor memory. There are only a dozen or so handwritten methods such as dim() and numel(); every other operator method is “mixed into” the class by the generated header file, saving several hundred lines of method declarations:
// p10/include/Tensor.h
class P10_API Tensor {
// ... dim()、numel()、device() 等手写基础方法 ...
#include "tensorplay/ops/TensorGenerated.h"
};
Because the #include lands inside the class definition, static Tensor zeros(...) in the generated header becomes a static member of Tensor, and Tensor add(const Tensor& other) const becomes a const instance method. The generated .cpp is compiled separately. The framework does conflict detection along the way: C++ cannot overload on the return type or default arguments alone, so plan_members compares parameter type prefixes — if two members share the same prefix and every parameter after that prefix has a default, they count as a conflict (a call passing only the prefix is valid for both), and such variants are marked “do not emit” ahead of time, saving you a stream of “ambiguous overload” errors at compile time.
The Dispatcher: Priority by Key Value
The routing logic for operator execution rests on a single enum, DispatchKey. The framework encodes run-time attributes such as device type, gradient mode, autocast, and batching into dispatch keys with fixed numeric values, and the larger the value the higher the priority:
// p10/include/DispatchKey.h
enum class DispatchKey : uint8_t {
CPU = 0, CUDA = 1, Vulkan = 2, Sparse = 3, Meta = 4, // 硬件后端
Python = 5,
DynamicLayerBackMode = 6,
AutogradCPU = 7, ... AutogradMeta = 11, // 自动微分（固定偏移+7）
AutocastCPU = 12, ... AutocastMeta = 16, // 自动类型转换（固定偏移+12）
VmapCPU = 17, ... VmapMeta = 21, // 批处理（固定偏移+17）
Composite = 22,
VmapMode = 23, DynamicLayerFrontMode = 24,
EndOfKeys = 25
};
constexpr size_t kAutogradKeyOffset = 7;
constexpr size_t kAutocastKeyOffset = 12;
constexpr size_t kVmapKeyOffset = 17;
This ordering buys two convenient things. A derived key comes out of arithmetic alone — add a fixed offset to the base key, and toAutogradKey(CPU) is just CPU + 7, no table lookup. And finding the current highest priority does not mean comparing one by one: the key set is a 32-bit bitmap, so highest_priority_key() just finds the highest bit that is set.
The key set is not stored in a field; it is computed on every call:
// p10/include/TensorImpl.h
DispatchKeySet key_set() const {
const DispatchKey backend =
sparse_state_ ? DispatchKey::Sparse : computeDispatchKey(device_);
DispatchKeySet ks;
ks.add(backend); // 稀疏张量优先走稀疏键
if (is_batched()) {
ks.add(toVmapKey(backend)); // 批处理场景仅启用Vmap键，跳过梯度
return ks;
}
if (!inference_tensor_) ks.add(toAutogradKey(backend)); // 推理模式关闭梯度
return ks;
}
Not storing a field means one less piece of storage and one less thing that has to be kept consistent; what it costs is one device query plus two bit operations — cheaper than atomically loading a bitmap out of a field.
The Dispatch Table: Lock-Free Lookup, With a Composite Fallback
Each operator gets its own independent dispatch table. Inside the table, indexed by dispatch key, sit two arrays of atomics: one holding kernel function pointers, one holding signature types. The lookup is lock-free end to end:
// p10/include/Dispatcher.h
using KernelFunction = void*;
struct P10_API DispatchTable {
std::string name;
std::array<std::atomic<KernelFunction>, kDispatchKeyCount> kernels;
std::array<std::atomic<const std::type_info*>, kDispatchKeyCount> signatures;
};
Registration writes with release, lookup reads with acquire. The extra signatures array is there to backstop type safety: each kernel records its exact function pointer type at registration time, and if the caller casts the slot back to the stub signature and it does not match, that is undefined behavior; having recorded this copy means you can at least catch that kind of mistake at run time, instead of quietly executing some wrong memory.
The core kernel lookup carries a built-in Composite fallback: if a backend has not registered a dedicated kernel, the general composite implementation steps in; once that backend does register a specialized kernel, it simply overwrites, with no need to deregister first:
// p10/include/Dispatcher.h
KernelFunction getKernel(DispatchKey key) const noexcept {
if (!table_ || dispatchKeyIndex(key) >= kDispatchKeyCount) return nullptr;
auto kernel = table_->kernels[dispatchKeyIndex(key)].load(std::memory_order_acquire);
if (!kernel && is_backend_key(key)) {
constexpr auto kCompositeIdx = dispatchKeyIndex(DispatchKey::Composite);
kernel = table_->kernels[kCompositeIdx].load(std::memory_order_acquire);
}
return kernel;
}
Put plainly: the composite kernel lets every backend run at all; the specialized kernel makes it run a bit faster. If a + b does not yet have a specialized kernel for CUDA, it goes through the Composite implementation — what comes out is the same thing, only possibly slower.
Type Safety: The Cast Is Confined to One Place
The generated code does not cast on its own at the call site. Type restoration happens in exactly one place, DispatchStub::call; the type erasure of the whole dispatch is collected inside this one template:
// p10/include/Dispatcher.h
template<typename Return, typename... Args>
class DispatchStub {
public:
static Return call(const OperatorHandle& handle, DispatchKey key, Args... args) {
const auto local = impl::tls_local_dispatch_key_set();
auto keys = DispatchKeySet::make(key);
keys = (keys | local.included) - local.excluded;
DispatchKey selected = keys.highest_priority_key();
// 处理 Vmap 层、Python 层、变换层的分流
key = selected;
if (is_autocast_key(key) && !autocast::is_enabled(key)) [[unlikely]] {
TP_THROW(NotImplementedError, "Autocast kernel dispatched while autocast is disabled");
}
auto kernel_void = handle.getKernel(key);
if (!kernel_void) {
if (getenv("TP_TRACE_KERNEL_MISS") != nullptr) {
// 打印算子、分发键、32 层栈回溯，定位内核缺失
}
TP_THROW(NotImplementedError, "Kernel not found for op on backend");
}
using FuncType = Return(*)(Args...);
auto kernel = reinterpret_cast<FuncType>(kernel_void);
return kernel(std::forward<Args>(args)...);
}
};
Two spots are worth a mention. autocast::is_enabled is forward declared, so that Tensor.h does not get pulled into every single consumer; even though the upper-level dispatch has already skipped the autocast keys, it checks once more here — the convention that “autocast disabled counts as a miss” is exactly what this backs up. The branch is marked [[unlikely]], so the normal path predicts accurately. When a kernel cannot be found, set TP_TRACE_KERNEL_MISS and stderr will print the operator name, the key, and a 32-frame stack backtrace — chasing down “how come this backend never registered” will most likely rely on it.
Redispatch: Probing, Shape Sampling, and Version Bumps All Live Here
Every operator execution has to pass through this generated code, so the probe only has to be done in this one place:
// build/generated/tensorplay/ops/TensorGenerated.cpp （示意）
TENSORPLAY_API Tensor redispatch_relu_function(const Tensor& self) {
tensorplay::prof::OpRecord __tp_prof_rec("relu");
if (tensorplay::prof::g_capture_shapes.load(std::memory_order_acquire)) {
// 按需采集张量形状、数据类型，供性能分析用
}
static const OperatorHandle op_handle = Dispatcher::singleton().findHandle("relu");
DispatchKey __tp_arg_key = dispatchKeyForTensorArgs(self);
DispatchKey dispatch_key = (is_vmap_key(__tp_arg_key) ? __tp_arg_key : toBackendKey(__tp_arg_key));
auto&& __tp_result = DispatchStub<Tensor, const Tensor&>::call(op_handle, dispatch_key, self);
return std::move(__tp_result);
}
This stretch of code takes care of several things along the way: one OpRecord probe per operator, which when inactive is just a single acquire load; shape sampling branches only when the session really asks for it, so the second load lands inside the active branch; static const OperatorHandle is resolved once, after which the hot path is two acquire loads plus one branch prediction. In-place operators also run if (!InferenceMode::is_enabled()) self.unsafeGetTensorImpl()->bump_version(); to bump the version counter — that is exactly the version counter from the previous piece.
Registration: From One Static Object Per Operator to One Block Per Key
Early registration was one static object per operator per key, every macro expanding into a global variable, the symbol table bloating, and breakpoints scattered across hundreds of places. The bulk macro funnels dozens of registrations under the same key into a single static initialization:
// p10/include/Dispatcher.h
#define TENSORPLAY_LIBRARY_IMPL(KEY, NAME) \
static void TP_CONCAT(tensorplay_library_init_, NAME)(::tensorplay::Library&); \
static struct TP_CONCAT(TensorPlayLibraryInit_, NAME) { \
TP_CONCAT(TensorPlayLibraryInit_, NAME)() { \
::tensorplay::Library lib(::tensorplay::DispatchKey::KEY); \
TP_CONCAT(tensorplay_library_init_, NAME)(lib); \
} \
} TP_CONCAT(tensorplay_library_init_instance_, NAME); \
static void TP_CONCAT(tensorplay_library_init_, NAME)(::tensorplay::Library& m)
In use it is just a group of m.impl(...):
// p10/src/backend/composite/Pointwise.cpp
TENSORPLAY_LIBRARY_IMPL(Composite, PointwiseComposite) {
m.impl("xlogy", xlogy_native);
m.impl("xlogy_", xlogy__native);
m.impl("xlog1py", xlog1py_native);
// 批量注册更多通用算子
}
The benefit is that the symbol count drops from operator level to library level, and Debugging gets easier too — set one breakpoint on registerKernel, run it once, and you have fished out every bulk registration under that key, without going one by one to set points on hundreds of scattered static objects.
End to End: One Addition
The whole chain for tp.add(a, b): a and b are CPU float tensors, a has gradients enabled, autocast is off, there is no batching, and it is not inference mode.
- The Python binding layer unpacks the arguments with a fast argument array and calls the generated operator wrapper;
key_set()computes the key set{CPU, AutogradCPU}on the spot, and the gradient key has the higher priority;AutogradCPUis hit first, hitting the gradient wrapper and recording the backward nodeAddBackward;- after
remove_autograd, a second dispatch lands on the baseCPUkey; - it matches the
add_cpunative kernel and performs the real addition.
This trip only touches two of the twenty-odd slots. Switch to inference mode, the tensor needs no gradient, the key set is just {CPU}, one dispatch reaches the kernel, and the chain gets noticeably shorter.
The Python Binding Layer: Handwritten Comes First, Generated Fills In
The generated Python entry point goes through the METH_FASTCALL argument array and does not build a temporary tuple. It coexists with the handwritten pybind11 path under the rule “handwritten comes first, generated fills in”: at install time, if some name is already occupied by a handwritten binding, the generated one is skipped and will not overwrite the handwritten semantics. The two paths meet on the bridge, and an operator name keeps only one binding.
Design Comparison: The Trade-offs Against Mainstream Frameworks
Mainstream frameworks broadly share the same thinking on operator dispatch; only the trade-offs differ.
Mainstream frameworks hand type promotion, alias analysis, and meta tensors to the generator together, so there are more artifacts and more keys; TensorPlay does only the core path, with low-frequency axes like sparse and quantization kept out of the core dispatch, at the cost of a few null pointers, or one composite key.
For dispatch keys, both sides use keys for run-time routing, pick the kernel by numeric magnitude, and fall back to a composite implementation. Mainstream frameworks store the key set as a bitmap, and once there are many keys you need bit operations; TensorPlay encodes the keys at fixed offsets and computes the key set on the fly — a few branch tests, cheaper than one atomic bitmap load.
For type safety, mainstream frameworks have a whole system of type erasure; TensorPlay folds the type erasure into the single DispatchStub template, the code makes no bare casts, and relies on the signatures array to check the signature at run time.
References
- Repository: github.com/lexing-2026/TensorPlay (Apache 2.0)
- Documentation and tutorials: tensorplay.cn
- Source code guide: deepwiki.com/bluemoon-o2/TensorPlay