Generate C++ entry points

[dcrc2]

Draft PR: this would currently break the embedded C++ examples added in #877. I'd prefer to merge #877 first and do the fixes here, unless there is still much more work to do on #877.
Now merged with #877.

The purpose of this PR is to reimplement the functionality currently provided by the function with_ks_allocator. That is, for each ks function that we wish to provide python bindings for, we produce a second C++ function which wraps it. This second C++ function is the one that actually gets used by the pybind module. Previously the wrapper function was defined using C++ template trickery in with_ks_allocator; this PR instead generates code for the wrapper in python.

Example code produced for relu3:

#include "knossos-entry-points.h"

namespace ks {
namespace entry_points {
namespace generated {

ks::tensor<1, double> entry(ks::tensor<1, double> arg0) {
    if (g_logging) {
        std::cerr << "vrelu3$aT1f(" << arg0 << ") =" << std::endl;
        auto ret = ks::vrelu3$aT1f(&g_alloc, arg0);
        std::cerr << ret << std::endl;
        return ret;
    } else {
        return ks::vrelu3$aT1f(&g_alloc, arg0);
    }
}

ks::tensor<1, double> entry_vjp(ks::tensor<1, double> arg0, ks::tensor<1, double> arg1) {
    if (g_logging) {
        std::cerr << "sufrev$vrelu3$aT1f(" << arg0 << ", "  << arg1 << ") =" << std::endl;
        auto ret = ks::sufrev$vrelu3$aT1f(&g_alloc, arg0, arg1);
        std::cerr << ret << std::endl;
        return ret;
    } else {
        return ks::sufrev$vrelu3$aT1f(&g_alloc, arg0, arg1);
    }
}

}
}
}


#include "knossos-pybind.h"

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("reset_allocator", &ks::entry_points::reset_allocator);
    m.def("allocator_top", &ks::entry_points::allocator_top);
    m.def("allocator_peak", &ks::entry_points::allocator_peak);
    m.def("logging", &ks::entry_points::logging);

    declare_tensor_1<double>(m, "Tensor_1_Float");
    declare_tensor_2<double>(m, "Tensor_2_Float");
    declare_tensor_2<int>(m, "Tensor_2_Integer");


        m.def("entry", &ks::entry_points::generated::entry);


        m.def("entry_vjp", &ks::entry_points::generated::entry_vjp);

}

#include "knossos-entry-points.cpp"

[dcrc2]

(The aim is to switch these generated wrappers to use torch::tensor rather than ks::tensor, but I would prefer to make that a separate change.)

[awf]

Yes, looking forward to the version with torch tensor.

[awf]

I think this will be more readable as an f string.

[awf]

Hmm, let me see if I agree with myself:

   args=", ".join(f"{ks_cpp_type(arg_type)} arg{i}" for i, arg_type in enumerate(arg_types));
   cpp_declaration = f"{ks_cpp_type(return_type)} {cpp_function_name}({args})"

I think I do.

[dcrc2]

OK, done

[dcrc2]

(I've rebased this onto master in order to use ks::Float from #924)

[dcrc2]

In order to merge with #877, I've included the following code in each of the embedded C++ examples:

ks::tensor<1, ks::Float> entry(ks::tensor<1, ks::Float> t) {
    return ks::vrelu3(&ks::entry_points::g_alloc, t);
}
ks::tensor<1, ks::Float> entry_vjp(ks::tensor<1, ks::Float> t, ks::tensor<1, ks::Float> dret) {
    return ks::sufrev_vrelu3(&ks::entry_points::g_alloc, t, dret);
}

It's not really possible to generate this code automatically (replicating what this PR does when compiling a ks file), because it's hard to generate the code without knowing the function signatures. But if we want to reproduce what these wrapper functions should look like, it's always possible to run one of the ks benchmarks, and then copy the wrapper functions from the generated C++ file.

[dcrc2]

(Rebased onto #944 and #947)

[toelli-msft]

This seems fine. Generating the wrappers ourselves seems to afford more flexibility than doing so with C++ template trickery.

On the other hand there's a bunch of code that duplicates Cgen (see the N.B. comment). An alternative would be to have Cgen output a ks::without_allocator::... function for every function in the module. For example, if my .ks file defines f : Integer -> Integer then as well as Cgen generating ks::Integer ks::f(allocator *, ks::Integer) it could also generate ks::Integer ks::without_allocator::f(ks::Integer). That way we don't need any special logic on the Python side.

What do you think?

[toelli-msft]

I wonder how useful this is if ret is a size one million tensor.

[dcrc2]

Ah, maybe this code is past its sell-by date: it's originally from #762, when we needed to debug issues to do with python object lifetimes. So, it was intended for debugging rather than routine logging. Though we're making changes in this area at the moment so maybe it's worth keeping a little while longer.

[toelli-msft]

Sure, fine by me.

[toelli-msft]

I'm always uneasy with zip because I wonder whether trailing elements are being silently discarded. Perhaps the following instead?

arg_types_names = ((arg.type_, f"arg{i}") for (i, arg) in enumerate(decl.args))
...
    for arg_type, arg_name in arg_types_names

(Even inlining arg_types_names would probably be fine.)

[dcrc2]

I like this but need to do something about the other uses of arg_names in this function. I've pushed a change with one way of doing this - what do you think?

[toelli-msft]

Thanks, that version looks good. I think if I were doing it I would do

arg_types_names = [(arg.type_, f"arg{i}") for (i, arg) in enumerate(decl.args)]
...
    for arg_type, arg_name in arg_types_names
...
arg_names = [arg_name for (_, arg_name) in arg_types_names]

but that's not necessarily to everyone's taste.

[dcrc2]

I think the first line has to be arg_types = ... because we need to do de-oneargification as well. (I probably should write the function multiargify rather than inlining it here, but I don't think that solves the problem.) I tried enumerate(multiargify(...)) but that seemed to be doing too much in one line of code. Anyway, it's looking a bit better now - happy to come back to this but I'll merge here.

[toelli-msft]

OK, fine by me

[toelli-msft]

N.B. Overlaps with Cgen

[dcrc2]

This seems fine. Generating the wrappers ourselves seems to afford more flexibility than doing so with C++ template trickery.

On the other hand there's a bunch of code that duplicates Cgen (see the N.B. comment). An alternative would be to have Cgen output a ks::without_allocator::... function for every function in the module. For example, if my .ks file defines f : Integer -> Integer then as well as Cgen generating ks::Integer ks::f(allocator *, ks::Integer) it could also generate ks::Integer ks::without_allocator::f(ks::Integer). That way we don't need any special logic on the Python side.

What do you think?

As things stand in this PR, I think you're right that this would fit better in ksc.

I'm not sure whether this will remain true when we add more functionality to the wrappers. The next PR (#931, currently draft) changes the wrappers to take torch::Tensor arguments, handling the conversion between torch types and ks types. So if this was part of ksc, it would mean ksc taking on some of the responsibility of knowing how we will call into ks code from PyTorch. The declaration of the pybind11 module is generated in python code, so this would create some overlap in responsibilities between ksc and python code.

If we wanted ksc to handle all of the binding code (including the declaration of the pybind11 module), that seems more consistent to me. But then we're probably committing to having ksc know how to generate bindings for all of the languages that we support (PyTorch, pure Python, and others that might be added in future).

I'd suggest we wait to see what these wrappers look like in their final version, as there are plenty of issues left to resolve. I don't think the amount of duplication with ksc is going to grow much further.

[toelli-msft]

I'd suggest we wait to see what these wrappers look like in their final version, as there are plenty of issues left to resolve. I don't think the amount of duplication with ksc is going to grow much further.

Makes sense, thanks!