Kernel Registry

In the previous example, we saw how to use wisdom files by creating a WisdomKernel object. This object will compile the kernel code on the first call and then keep the kernel loaded as long as the object exists. Typically, one would define the WisdomKernel object as part of a class or as a global variable.

However, in certain scenarios, it is inconvenient or impractical to store WisdomKernel objects. In these cases, it is possible to use the KernelRegistry that essentially acts like a global table of compiled kernel instances.

Source code

Consider the following code snippet:

#include "kernel_launcher.h"

// Namespace alias.
namespace kl = kernel_launcher;

class VectorAddDescriptor: kl::IKernelDescriptor {
public:
    template <typename T>
    static VectorAddDescriptor for_type() {
        return VectorAddDescriptor(kl::type_of<T>());
    }

    VectorAddDescriptor(kl::TypeInfo t): element_type(t) {}

    kl::KernelBuilder build() const override {
        kl::KernelBuilder builder("vector_add", "vector_add.cu");

        auto threads_per_block = builder.tune("block_size", {32, 64, 128, 256, 512, 1024});
        auto elements_per_thread = builder.tune("elements_per_thread", {1, 2, 4, 8});
        auto elements_per_block = threads_per_block * elements_per_thread;

        builder
            .tuning_key("vector_add_" + this->element_type.name())
            .problem_size(kl::arg0)
            .block_size(threads_per_block)
            .grid_divisors(threads_per_block * elements_per_thread)
            .template_args(element_type)
            .define("ELEMENTS_PER_THREAD", elements_per_thread);

        return builder;
    }

    bool equals(const IKernelDescriptor& other) const override {
        if (auto p = dynamic_cast<const VectorAddDescriptor*>(&other)) {
            return this->element_type == p->element_type;
        }

        return false;
    }

    private:
        kl::TypeInfo element_type;
};

int main() {
    kl::set_global_wisdom_directory("wisdom/");
    kl::set_global_capture_directory("captures/");

    // Initialize CUDA memory. This is outside the scope of kernel_launcher.
    unsigned int n = 1000000;
    float *dev_A, *dev_B, *dev_C;
    /* cudaMalloc, cudaMemcpy, ... */

    // Launch the kernel!
    kl::default_registry()
        .lookup(VectorAddDescriptor::for_type<float>())
        .launch(n, dev_C, dev_A, dev_B);

    // Or use the short equivalent syntax:
    kl::launch(VectorAddDescriptor::for_type<float>(), n, dev_C, dev_A, dev_B);

    return 0;
}

Code Explanation

The code example consists of two parts. In the first part, a class VectorAddDescriptor is defined. In the second part, this class is searched in the global kernel registry.

Defining a kernel descriptor

class VectorAddDescriptor: kl::IKernelDescriptor {
public:
    template <typename T>
    static VectorAddDescriptor for_type() {
        return VectorAddDescriptor(kl::type_of<T>());
    }

    VectorAddDescriptor(kl::TypeInfo t): element_type(t) {}

    kl::KernelBuilder build() const override {
        kl::KernelBuilder builder("vector_add", "vector_add.cu");

        auto threads_per_block = builder.tune("block_size", {32, 64, 128, 256, 512, 1024});
        auto elements_per_thread = builder.tune("elements_per_thread", {1, 2, 4, 8});
        auto elements_per_block = threads_per_block * elements_per_thread;

        builder
            .tuning_key("vector_add_" + this->element_type.name())
            .problem_size(kl::arg0)
            .block_size(threads_per_block)
            .grid_divisors(threads_per_block * elements_per_thread)
            .template_args(element_type)
            .define("ELEMENTS_PER_THREAD", elements_per_thread);

        return builder;
    }

    bool equals(const IKernelDescriptor& other) const override {
        if (auto p = dynamic_cast<const VectorAddDescriptor*>(&other)) {
            return this->element_type == p->element_type;
        }

        return false;
    }

    private:
        kl::TypeInfo element_type;
};

This part of the code defines an IKernelDescriptor: a class that encapsulates the information required to compile a kernel. This class should override two methods:

• build to instantiate a KernelBuilder,

• equals to check for equality with another IKernelDescriptor.

The last method is required since a kernel registry is essentially a hash table that maps IKernelDescriptor objects to kernel objects. The equals method is used to check if two descriptors (i.e., keys in the hash table) are equivalent.

Using the KernelRegistry

    // Launch the kernel!
    kl::default_registry()
        .lookup(VectorAddDescriptor::for_type<float>())
        .launch(n, dev_C, dev_A, dev_B);

Here, the vector-add kernel is searched in the registry and launched with the given arguments. It is important to note that this code can be called multiple times from different functions of a program, but the kernel is only compiled once and stored in the registry.

    // Or use the short equivalent syntax:
    kl::launch(VectorAddDescriptor::for_type<float>(), n, dev_C, dev_A, dev_B);

Alternatively, it is possible to use the above short-hand syntax. This syntax also makes it easy to replace the element type float with some other type such as int:

kl::launch(VectorAddDescriptor::for_type<int>(), n, dev_C, dev_A, dev_B);

It is even possible to define a templated function that passes type T on to VectorAddDescriptor, for some extra template magic:

template <typename T>
void launch_vector_add(T* C, const T* A, const T* B) {
    kl::launch(VectorAddDescriptor::for_type<T>(), n, C, A, B);
}

Instead of using the global kernel registry, it is also possible to create a local registry by creating a KernelRegistry instance.