--- title: Write a custom operator description: "Subclass nn::Module two ways: compose built-ins into a reusable block, or run your own kernel through the runtime with output_shapes and compute." --- The built-in `ops::` library does not have to be the end of the line. `nn::Module` is both the parameter container (PyTorch-style dotted names, so HuggingFace checkpoints bind by name) and the custom-op entry point: a leaf module that overrides `output_shapes()` and `compute()` runs its own kernel through the runtime, eager and tracing alike. This guide builds one of each. ## Compose built-ins into a module A composite module subclasses `nn::Module`, registers its children in the constructor, and defines its own `forward` composing `ops::` and child modules. Registration is what buys the dotted names: the block below enumerates `up.weight` and `down.weight`, which is exactly how a checkpoint refers to them. ```cpp title="mlp_block.cpp" #include #include #include #include using ClikaRT::DataType; using ClikaRT::nn::Linear; using ClikaRT::NamedTensors; using ClikaRT::Tensor; namespace nn = ClikaRT::nn; namespace ops = ClikaRT::ops; // A residual MLP block: y = down(relu(up(x))) + x. class MlpBlock final : public nn::Module { public: static std::shared_ptr make(std::int64_t dim, std::int64_t hidden) { std::shared_ptr m(new MlpBlock()); m->up_ = Linear::make(dim, hidden); m->down_ = Linear::make(hidden, dim); m->register_module("up", m->up_); m->register_module("down", m->down_); return m; } Tensor forward(const Tensor& x) const { Tensor h = this->up_->forward(x); ops::relu_(h); // fresh from the projection: nothing else sees it h = this->down_->forward(h); ops::add_(h, x); // the residual add writes the fresh down output; x stays a read return h; } private: MlpBlock() = default; std::shared_ptr up_; std::shared_ptr down_; }; int main() { const std::shared_ptr block = MlpBlock::make(4, 8); std::printf("parameters (dotted, checkpoint-shaped):\n"); for (const auto& [name, t] : block->named_parameters()) std::printf(" %-12s fake=%d\n", name.c_str(), static_cast(t.is_fake())); // Bind weights by name, the way a checkpoint would. NamedTensors ckpt; ckpt.set("up.weight", Tensor::full({8, 4}, 0.1, DataType::Float32)); ckpt.set("down.weight", Tensor::full({4, 8}, 0.1, DataType::Float32)); block->load_state_dict(ckpt); const Tensor x = Tensor::ones({1, 4}, DataType::Float32); std::printf("y = %s\n", block->forward(x).to_string().c_str()); return 0; } ``` ```python title="mlp_block.py" import clika_runtime as crt import clika_runtime.nn as nn class MlpBlock(nn.Module): """A residual MLP block: y = down(relu(up(x))) + x. A model is a Module subclass: assigning a module REGISTERS it: load_state_dict, named_parameters and repr all see `up` and `down` automatically. The first linear fuses its ReLU as an activation epilogue, named by its string spelling; without the epilogue the same block writes F.relu(self.up(x)) with `import clika_runtime.nn.functional as F`. Bias defaults to on, so these layers opt out explicitly.""" def __init__(self, dim: int, hidden: int) -> None: super().__init__() self.up = nn.Linear(dim, hidden, bias=False, activation="relu") self.down = nn.Linear(hidden, dim, bias=False) def forward(self, x: crt.Tensor) -> crt.Tensor: return self.down(self.up(x)) + x def main() -> None: block = MlpBlock(4, 8) print(block) # the module tree, registered names included print("parameters (dotted, checkpoint-shaped):") for name, t in block.named_parameters(): print(f" {name}") # Weights arrive as a plain dict; dotted keys route to the registered # submodules. block.load_state_dict({ "up.weight": crt.full((8, 4), 0.1), "down.weight": crt.full((4, 8), 0.1), }) y = block(crt.ones(1, 4)) # calling the module runs forward print(f"y = {y}") if __name__ == "__main__": main() ``` ```text parameters (dotted, checkpoint-shaped): up.weight fake=1 down.weight fake=1 y = Tensor(shape=[1, 4], dtype=Float32, device=CPU, numel=4, data=[1.32, 1.32, 1.32, 1.32]) ``` The math checks out by hand: `up` maps ones to 0.4 per unit, relu passes it, `down` sums 8 of them times 0.1 to 0.32, and the residual adds the input back, 1.32. Python prints a tensor as its values alone, so its last line reads `y = tensor([[1.3200, 1.3200, 1.3200, 1.3200]])`. `Linear::make` declares storage-free slots, so the parameters read as fake until `load_state_dict` binds them; `block->to(device)` moves the whole tree. ## Run your own kernel through the runtime A leaf custom op overrides two virtuals. `output_shapes()` is the shape rule: it reads the inputs' metadata (`FakeTensor`: symbolic dims, dtype, placement) and returns the outputs' metadata. `compute()` is the kernel, with one accessor pair per residence: for a plain host loop over CPU-resident tensors, inputs read through `const_data_ptr()` and the pre-allocated outputs write through `mutable_data_ptr()` (both wait for the data and refuse a device tensor); a kernel on the op's device reads `device_const_data_ptr()` and writes `device_mutable_data_ptr()`, with no host wait and the pointer ordered on the op's stream. An exception thrown inside `compute()` surfaces as a typed error in the caller's runtime instead of crossing the library boundary. `forward` hands both to `dispatch()`, which runs the op through the runtime: eager mode calls `compute()` now, a tracing scope records the op from the shape rule alone. A leaf that dispatches must be owned by a `shared_ptr`. ```cpp title="softclip_op.cpp" #include #include #include #include #include #include using ClikaRT::DataType; using ClikaRT::FakeTensor; using ClikaRT::Tensor; namespace nn = ClikaRT::nn; // Elementwise soft clip: y = x / (1 + |x|). class SoftClip final : public nn::Module { public: static std::shared_ptr make() { return std::shared_ptr(new SoftClip()); } Tensor forward(const Tensor& x) const { return dispatch({&x, 1})[0]; } // Shape rule: one output, same shape and dtype as the input. std::vector output_shapes( ClikaRT::Span inputs) const override { return {FakeTensor(inputs[0].shape(), inputs[0].dtype(), inputs[0].stream())}; } // Kernel: runs on the op's stream; here the CPU path, a plain loop over // host pointers. A device kernel (a CUDA arm, say) reads // `device_const_data_ptr()`, writes `device_mutable_data_ptr()`, and // launches on `this->stream().native_handle()`; the host accessors below // wait for the data and refuse a device tensor. void compute(ClikaRT::Span inputs, ClikaRT::Span outputs) const override { const Tensor& in = inputs[0]; const float* x = static_cast(in.const_data_ptr()); float* y = static_cast(outputs[0].mutable_data_ptr()); const std::int64_t n = in.numel(); for (std::int64_t i = 0; i < n; ++i) y[i] = x[i] / (1.0F + std::fabs(x[i])); } private: SoftClip() = default; }; int main() { const std::shared_ptr clip = SoftClip::make(); std::vector v = {-9.0F, -1.0F, 0.0F, 1.0F, 9.0F}; const Tensor x = Tensor::from_data(v.data(), {5}, DataType::Float32); std::printf("y = %s\n", clip->forward(x).to_string().c_str()); return 0; } ``` ```python title="softclip_op.py" import clika_runtime as crt import clika_runtime.nn as nn class SoftClip(nn.Module): """Elementwise soft clip: y = x / (1 + |x|), composed from the built-in operations; a custom module needs nothing beyond a forward. Writing a custom KERNEL (your own shape rule and compute) is done through the C++ API; the C++ tab walks it.""" def forward(self, x: crt.Tensor) -> crt.Tensor: return x / (1 + x.abs()) def main() -> None: clip = SoftClip() x = crt.tensor([-9.0, -1.0, 0.0, 1.0, 9.0]) print(f"y = {clip(x)}") if __name__ == "__main__": main() ``` ```text y = Tensor(shape=[5], dtype=Float32, device=CPU, numel=5, data=[-0.9, -0.5, 0, 0.5, 0.9]) ``` The runtime allocated the output, placed it beside the input, and ran the kernel; the op composes with everything else (`ops::` calls before and after, `on_complete`, the scopes) because it went through `dispatch` like a built-in. Python's line reads `y = tensor([-0.9000, -0.5000, 0.0000, 0.5000, 0.9000])`: the same values in its own tensor form. ## Launching on an accelerator `compute()` runs on whatever device its tensors live on. Inside it, `this->stream()` is the op's stream: `stream().device()` says where you are, and `stream().native_handle()` is the backend's queue as an opaque pointer, `cudaStream_t` on CUDA (`nullptr` on the CPU backend, which has no device queue). The launch pattern, from the `nn/module.h` contract: ```text Stream s = this->stream(); // the op's stream auto cu = static_cast(s.native_handle()); // the CUDA queue const float* q = static_cast(inputs[0].device_const_data_ptr()); float* o = static_cast(outputs[0].device_mutable_data_ptr()); my_kernel<<>>(q, ..., o); // your kernel ``` The device accessors return the device address with no host wait, ordered on the op's stream, so the launch above is valid as written; the host accessors (`const_data_ptr()`, `mutable_data_ptr()`) are for CPU-resident tensors and would refuse here. Size the launch from `get_device_properties(stream().device())`, and take stream-ordered scratch from `stream().allocate(shape, dtype)`, which recycles safely on that stream only. The handle is owned by the runtime: never destroy it, and keep the `Stream` alive while using it. The bundle's `flash_attention` example is the complete worked case, a hand-written CUDA attention kernel dispatched through this exact interface. A rule of thumb for choosing the shape: compose built-ins when the math decomposes into `ops::` (the runtime already fuses and places them); write a leaf when you have a kernel the library does not, and keep its `output_shapes` honest, because tracing trusts it without running `compute`.