Skip to main content

Write a custom operator

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. The C samples abbreviate the api-table bootstrap that tutorial part 1 shows in full.

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.

mlp_block.cpp
#include <cstdint>
#include <cstdio>
#include <memory>

#include "ClikaRT/clika_rt.h"

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<MlpBlock> make(std::int64_t dim, std::int64_t hidden) {
        std::shared_ptr<MlpBlock> 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 {
        return ops::add(down_->forward(ops::relu(up_->forward(x))), x);
    }

 private:
    MlpBlock() = default;
    std::shared_ptr<Linear> up_;
    std::shared_ptr<Linear> down_;
};

int main() {
    const std::shared_ptr<MlpBlock> 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(), (int)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;
}
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. 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.

softclip_op.cpp
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <memory>
#include <vector>

#include "ClikaRT/clika_rt.h"

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<SoftClip> make() {
        return std::shared_ptr<SoftClip>(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<FakeTensor> output_shapes(
        ClikaRT::Span<const FakeTensor> 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<const Tensor> inputs,
                 ClikaRT::Span<Tensor> outputs) const override {
        const Tensor& in = inputs[0];
        const float* x = static_cast<const float*>(in.const_data_ptr());
        float* y = static_cast<float*>(outputs[0].mutable_data_ptr());
        for (std::int64_t i = 0; i < in.numel(); ++i)
            y[i] = x[i] / (1.0F + std::fabs(x[i]));
    }

 private:
    SoftClip() = default;
};

int main() {
    const std::shared_ptr<SoftClip> clip = SoftClip::make();

    std::vector<float> 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;
}
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.

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:

Stream s = this->stream();                                // the op's stream
auto cu = static_cast<cudaStream_t>(s.native_handle());   // the CUDA queue
const float* q = static_cast<const float*>(inputs[0].device_const_data_ptr());
float*       o = static_cast<float*>(outputs[0].device_mutable_data_ptr());
my_kernel<<<grid, block, smem, cu>>>(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.