Extended kernel flavours

[mortendahl]

This is a discussion issue.

We currently have three kernel flavours: hybrid, transparent, and concrete. We are looking at getting rid of hybrid since the use of Into/TryInto is causing some complexity. However, some kernels do not naturally fit with neither transparent nor concrete. One example is sharing, which can only take a symbolic host tensor but ideally returns a concrete replicated tensor.

modelled_kernel! {
    PlacementShare::share, RepShareOp,
    [
        (ReplicatedPlacement, (HostRing64Tensor) -> ReplicatedRing64Tensor => [hybrid] Self::ring_kernel),
        (ReplicatedPlacement, (HostRing128Tensor) -> ReplicatedRing128Tensor => [hybrid] Self::ring_kernel),
        (ReplicatedPlacement, (HostBitTensor) -> ReplicatedBitTensor => [hybrid] Self::ring_kernel),
        (ReplicatedPlacement, (HostBitArray64) -> ReplicatedBitArray64 => [concrete] Self::array_kernel),
        ...
    ]
}

One approach might be to introduce more flavours, say concrete_out (eg for sharing) and concrete_in (eg for revealing).

However, given that this does not scale well with arity, it might be more interesting to move away from a combined flavour and instead annotate inputs and outputs as in the following:

modelled_kernel! {
    PlacementShare::share, RepShareOp,
    [
        (ReplicatedPlacement, ([transparent] HostRing64Tensor) -> [concrete] ReplicatedRing64Tensor => Self::ring_kernel),
        (ReplicatedPlacement, ([transparent] HostRing128Tensor) -> [concrete] ReplicatedRing128Tensor => Self::ring_kernel),
        (ReplicatedPlacement, ([transparent] HostBitTensor) -> [concrete] ReplicatedBitTensor => Self::ring_kernel),
        (ReplicatedPlacement, ([concrete] HostBitArray64) -> [concrete] ReplicatedBitArray64 => Self::array_kernel),
        ...
    ]
}

Note that the above would implement the PlacementShare trait twice: once for purely symbolic, and once matching the annotations.

[mortendahl]

cc @voronaam

[voronaam]

My objection to this would be the fact that the kernel's flavour is part of "what" to do with an op and it gets into the left-side "when" do it. For example, the [concrete] is not part of the type signature. Perhaps having a separate place for it would make sense. Something like

modelled_kernel! {
    PlacementShare::share, RepShareOp,
    [
        (ReplicatedPlacement, (HostRing64Tensor) -> ReplicatedRing64Tensor => [transparent -> concrete] => Self::ring_kernel),
        (ReplicatedPlacement, (HostRing128Tensor) -> ReplicatedRing128Tensor => [transparent -> concrete] => Self::ring_kernel),
        (ReplicatedPlacement, (HostBitTensor) -> ReplicatedBitTensor => [transparent -> concrete] => Self::ring_kernel),
        (ReplicatedPlacement, (HostBitArray64) -> ReplicatedBitArray64 => [transparent -> concrete] => Self::array_kernel),
        ...
    ]
}

I was doing something similar to get rid of custom kernels with something like

modelled_kernel! {
    PlacementFill::fill, FillOp{value: Constant},
    [
        (ReplicatedPlacement, (ReplicatedShape) -> ReplicatedRing64Tensor => Self::pre_process_u64 => Self::ring_kernel),
        (ReplicatedPlacement, (ReplicatedShape) -> ReplicatedRing128Tensor => Self::pre_process_u128 => Self::ring_kernel),
        (ReplicatedPlacement, (ReplicatedShape) -> ReplicatedBitTensor => Self::pre_process_bit => Self::ring_kernel),

It feels like a custom syntax in the middle describing what kind of gluing code we want may be beneficial. Still looking for its syntax though.