[CK_BUILDER] Unified type-erased interface for kernels #5015

Problem

Currently, CK-Builder implements a compile-time only abstraction over building kernels. ConvBuilder<SIG, ALGO>::Instance returns either a CK or CK Tile instance directly, it is not wrapped in anything. This is not a problem during CK-Builder testing, because its fine to only have the concrete instance.

For MIOpen, this is a problem though. Previously, the old CK factories returned a vector of DeviceOp<... /* equivalent of a SIGNATURE */>, a list of abstract kernel handles. This made it easy to work with vectors of different implementations at runtime, as it does not need to be known at compile-time. While its still possible to use the internal CK "base" device op class (for example, DeviceBatchedGemm or DeviceGroupedConvFwdMultipleABD) for this, it requires reaching for the backing implementation's internals, and thats precisely what we want to avoid with the Builder. See related merge request #4570 for how this is done now (basically, the backing virtual base device op is used temporarily). It would be useful if CK provided a canonical way to do this.

Suggestion

It seems that for MIOpen, it is still fine that the resulting operation is templated using the SIGNATURE, its just the ALGORITHM (the implementation's details) that needs to be type erased.

Idea

I'm not sure how useful it is to retain the current non-virtual setup, but here is an idea: Basically, we create a virtual Kernel<SIGNATURE> class. Each built instance then inherits from this class, which we can do easily by creating a ConcreteKernel<SIGNATURE, ALGORITHM> class that inherits from Kernel<SIGNATURE>.

This also needs integration with ckt::run(). I'm not entirely sure what the best way to go about that is: ckt::run() depends on ckt, so either we'd need to put Kernel and ConcreteKernel in the ckt namespace, or we need to decouple the run() function from the interface somehow. I think its easier to decouple the Kernel from ConvBuilder, though. This also makes sense, since the runtime interface is currently already only defined in ckt.

So in practise, the implementation would look something like this:

#include "ck_tile/builder/testing/testing.hpp" // or maybe the below is part of this header

namespace ck_tile::builder::test {

template <auto SIGNATURE_>
struct Kernel {
  constexpr static auto SIGNATURE = SIGNATURE_;

  virtual ~Kernel() = default;
  virtual RunResult run(
    const Args<SIGNATURE>& args,
    const Inputs<SIGNATURE>& inputs,
    const Outputs<SIGNATURE>& outputs,
    // Note: The StreamConfig is currently a template. This would need to be
    // replaced by a custom CK-Builder definition, but that is OK, because its
    // a leaky abstraction right now anyway.
    const StreamConfig& s_conf = {}
  ) const = 0;
};

template <auto SIGNATURE, auto ALGORITHM>
struct ConcreteKernel : Kernel<SIGNATURE> {
  // Or perhaps use typename Instance_ as type parameter of ConcreteKernel?
  using Instance = typename ConvBuilder<SIGNATURE, ALGORITHM>::Instance;

  RunResult run(
  const Args<SIGNATURE>& args,
    const Inputs<SIGNATURE>& inputs,
    const Outputs<SIGNATURE>& outputs,
    const StreamConfig& s_conf = {}
  ) const override {
    // Depending on if we want to put more than this run() function inside the
    // Kernel interface, its probably easier to specialize the ConcreteKernel
    // on the backing implementation (not sure how easy this is, since we'd need to mirror
    // the implementation in `make_conv_instance()`).

    auto instance = Instance{};
    return ::ck_tile::builder::testing::run(instance, inputs, outputs, s_conf);
  }
};

// If you don't decide to remove ckt::run entirely, and replace it by just the
// ConcreteKernel::run(), then provide an overload for the virtual interface
// to keep things unified.
template <auto SIGNATURE>
RunResult run(const Kernel<SIGNATURE>& instance, ...) {
  instance.run(...);
}

}

Alternative

An alternative idea may be to provide the runtime implementation directly in each factory struct. The ckt implementation can be decoupled by using a visitor pattern. The factory would return some wrapper struct instead of the actual instance which provides the visitor function, and ckt::run can then traverse that. Note that either way, there is going to be some functionality required that checks which the concrete CK/CK Tile/Reference implementation is, so the current implementation of ckt::run() will likely have some useful parts.