Etaler backend re-architecture

[marty1885]

The current backend architecture works well of out purpose. But it is not salable (in a programming sense) enough and will cause problems down the line. For the backend system to work, i.e. direct function calls to the correct backend depending using the CPU or GPU, we currently uses C++'s virtual class methods. While this is the obvious and correct solution, Adding new methods to a backend in is pain. First et::Backend and on of et::CPUBackend and et::GPUBackend have to be modified. Causing a total recompilation of the library (since almost everything depends on et::Backend). Furthermore, this changes the vtable layout and makes versions with different backend methods not binary compatible with each other.

The current design also does not allow runtime expansion of the backend. For example, adding linear algebra functions dynamically to a base backend is not allowed. Even with inheritance. It allows one to add new methods, Yet these methods won't be accessible from et::Backend* so no one can access them.

My proposal is to re-architect the backend that the et::Backend class essentially exposes two methods only, register_function and get_function. Where register_function is used to by the backend to register which methods are available from the backend and get_function returns a std::function (a callable object) to the method. This way the frontend could ask the beckend which functions are available at runtime and allows expansion of the backend by adding more methods after initialization.

Here's a tiny Proof of Concept

#include <iostream>
#include <string>
#include <vector>
#include <utility>
#include <functional>
#include <map>
#include <any>

struct BackendBase
{
        std::map<std::string, std::any> funcs_;
        template<typename Class, typename Return, typename ...Args>
        void expose_method(const std::string& name, Return(Class::*func)(Args...))
        {
                funcs_[name] = std::function<Return(Args...)>([this, func](Args ... args){
                        (static_cast<Class*>(this)->*func)(args...);
                });
        }
        template <typename Func>
        std::function<Func> get_method(const std::string& name)
        {
                return std::any_cast<std::function<Func>>(funcs_.at(name));
        }
};

struct Backend1 : public BackendBase
{
        void foo() { std::cout << "foo() from backend1" << std::endl; }
        Backend1() { expose_method("foo", &Backend1::foo); }
};

struct Backend2 : public BackendBase
{
        void foo() { std::cout << "foo() from backend2" << std::endl; }
        Backend2() { expose_method("foo", &Backend2::foo); }
};

int main()
{
        Backend1 b;
        std::cout << b.funcs_["foo"].type().name() << std::endl;
        b.get_method<void()>("foo")();
        Backend2 c;
        std::cout << c.funcs_["foo"].type().name() << std::endl;
        c.get_method<void()>("foo")();
}

Note for myself: Like how one of my old project is written. But better, no dynamic_cast and no more BoxedValues

---

Note: PyTorch uses a more monolithic approach. PyTorch doesn't have a 'backend` per say, but each tensor operator tries to support computing on different hardware. ex:

void add(const tensor& a, const tensor& b, tensor& out)
{
        if(a.on_cpu())
                add_on_cpu(...);
        else if(a.on_cuda())
                add_on_cuda(...);
}

This doesn't support expansions to be backend, but it may not matter. Just update the tensor operators to support more hardware. This have a lower overhead and doesn't have crazy syntax to access a function. But no expansions allowed.

[marty1885]

Quick update, this modified code is slightly faster than the above code. 1. Returns a reference to std::function and reads a reference of std::function from std::any. 2. Use a std::string_view as key instead of std::string. But this feature requires C++20

 struct BackendBase
 {
         std::map<std::string, std::any> funcs_;
 
         template<typename Class, typename Return, typename ...Args>
         void expose_method(const std::string& name, Return(Class::*func)(Args...))
         {   
                 funcs_[name] = std::function<Return(Args...)>([this, func](Args ... args) { 
                         (static_cast<Class*>(this)->*func)(args...);
                 }); 
         }   
 
         template <typename Func>
         const std::function<Func>& get_method(const std::string_view name)
         {   
                 return std::any_cast<const std::function<Func>&>(funcs_.at(name));
         }   
 };

This is not the final design tho. I still want to remove the use of std::function and use a custom, more specific class. Hopefully the method wrapper can ensure no allocation what so ever.

---

Update: The current lambda needs exactly 24 bytes of storage space on 64bit systems. Which is exactly how much space libc++'s std::function have for small functions. No heap allocations. We are good in that regards :)