Support functions as layers (#188)

This change adds two types of new functionality. First of all, it
introduces the `(Locked|Local)?FuncRepo` classes these can be used
to extend a layer with a kernel function. For instance, a layer
like

```
@use_kernel_forward_from_hub("SiluAndMul")
class SiluAndMul(nn.Module):
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        d = input.shape[-1] // 2
        return F.silu(input[..., :d]) * input[..., d:]
```

can now also be kernelized using a function `silu_and_mul` from the
Hub:

```
with use_kernel_mapping({
    "SiluAndMul": {
        "cuda": FuncRepository(
            repo_id="kernels-community/activation",
            func_name="silu_and_mul",
        ),
    }
}):
    kernelize(...)
```

This makes it easier to kernelize pure layers (layers that do not use
module state), since the Hub kernel does not have to provide a `layers`
Python module with wrappers.

Secondly, we introduce a decorator `use_kernel_func_from_hub` that turns
functions into layers that can be kernelized. For example:

```
@use_kernel_forward_from_hub("silu_and_mul")
def silu_and_mul(x: torch.Tensor) -> torch.Tensor:
    d = x.shape[-1] // 2
    return F.silu(x[..., :d]) * x[..., d:]
```

will implicitly create an instance of the following class:

```
class Func(nn.Module):
    # We add some magic to preserve the function's signature.
    def forward(self, *args, **kwargs):
        return silu_and_mul(*args, **kwargs)
```

Due to the `__call__` implementation of `nn.Module`, the instance
still behaves as a function:

```
out = silu_and_mul(x)
```

However, when the function is used as a member of an `nn.Module`,
it will be kernelized:

```
class FeedForward(nn.Module):
  def __init__(self, in_features: int, out_features: int):
      self.linear = nn.Linear(in_features, out_features)
      # Note: silu_and_mul is a Torch module.
      self.silu_and_mul = silu_and_mul

  def forward(self, x: torch.Tensor) -> torch.Tensor:
      return self.silu_and_mul(self.linear(x))
```

Co-authored-by: Mohamed Mekkouri <[email protected]>

Author: danieldk

docs/source/api/layers.md

@@ -6,6 +6,10 @@
 
 [[autodoc]] kernels.use_kernel_forward_from_hub
 
+### use_kernel_func_from_hub
+
+[[autodoc]] kernels.use_kernel_func_from_hub
+
 ### replace_kernel_forward_from_hub
 
 [[autodoc]] kernels.replace_kernel_forward_from_hub
@@ -36,14 +40,26 @@
 
 [[autodoc]] kernels.Mode
 
+### FuncRepository
+
+[[autodoc]] kernels.FuncRepository
+
 ### LayerRepository
 
 [[autodoc]] kernels.LayerRepository
 
+### LocalFuncRepository
+
+[[autodoc]] kernels.LocalFuncRepository
+
 ### LocalLayerRepository
 
 [[autodoc]] kernels.LocalLayerRepository
 
+### LockedFuncRepository
+
+[[autodoc]] kernels.LockedFuncRepository
+
 ### LockedLayerRepository
 
 [[autodoc]] kernels.LockedLayerRepository

docs/source/layers.md

@@ -43,6 +43,36 @@ replace_kernel_forward_from_hub(SiluAndMul, "SiluAndMul")
 it signifies that the maintainer intends to keep the `forward` signature
 compatible with layers from the hub.
 
+### Using a function as a layer
+
+Sometimes it can be useful to make a function extensible, for example
+because the function cannot be replaced by a layer. In such cases, you
+can annotate the function with the `use_kernel_func_from_hub` decorator:
+
+```python
+@use_kernel_func_from_hub("silu_and_mul")
+def silu_and_mul(x: torch.Tensor) -> torch.Tensor:
+    d = x.shape[-1] // 2
+    return F.silu(x[..., :d]) * x[..., d:]
+```
+
+This will replace the function by an instantiated `torch.nn.Module`
+(singleton) that calls the function itself in its forward method.
+
+**Note:** for kernelization to see the function, it must be a member of
+another `torch.nn.Module` that is part of the model. For example:
+
+```python
+class FeedForward(nn.Module):
+  def __init__(self, in_features: int, out_features: int):
+      self.linear = nn.Linear(in_features, out_features)
+      # Note: silu_and_mul is a Torch module.
+      self.silu_and_mul = silu_and_mul
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+      return self.silu_and_mul(self.linear(x))
+```
+
 ## Kernelizing a model
 
 A model will not use Hub kernels by default, even if it contains extensible
@@ -157,6 +187,21 @@ with use_kernel_mapping(kernel_layer_mapping):
 This ensures that the mapping is not active anymore outside the
 `with`-scope.
 
+If the layer is stateless (it does not use member variables in its forward _or_ it was
+originally a function that was converted into a kernel layer with
+`use_kernel_func_from_hub`), it can also be mapped to a kernel function:
+
+```python
+kernel_layer_mapping = {
+    "SiluAndMul": {
+        "cuda": FuncRepository(
+            repo_id="kernels-community/activation",
+            func_name="silu_and_mul",
+        ),
+    }
+}
+```
+
 ### Using version bounds
 
 Kernels are versioned using tags of the form `v<major>.<minor>.<patch>`.

src/kernels/__init__.py

@@ -2,15 +2,22 @@
 
 __version__ = importlib.metadata.version("kernels")
 
-from kernels.layer import Device, CUDAProperties
-from kernels.layer import kernelize, register_kernel_mapping, use_kernel_mapping
-from kernels.layer import Mode
 from kernels.layer import (
+    CUDAProperties,
+    Device,
+    FuncRepository,
     LayerRepository,
+    LocalFuncRepository,
     LocalLayerRepository,
+    LockedFuncRepository,
     LockedLayerRepository,
+    Mode,
+    kernelize,
+    register_kernel_mapping,
     replace_kernel_forward_from_hub,
     use_kernel_forward_from_hub,
+    use_kernel_func_from_hub,
+    use_kernel_mapping,
 )
 from kernels.utils import (
     get_kernel,
@@ -25,8 +32,11 @@
     "__version__",
     "CUDAProperties",
     "Device",
+    "FuncRepository",
     "LayerRepository",
+    "LocalFuncRepository",
     "LocalLayerRepository",
+    "LockedFuncRepository",
     "LockedLayerRepository",
     "Mode",
     "get_kernel",
@@ -38,7 +48,6 @@
     "load_kernel",
     "register_kernel_mapping",
     "replace_kernel_forward_from_hub",
-    "replace_kernel_func_from_hub",
     "use_kernel_forward_from_hub",
     "use_kernel_func_from_hub",
     "use_kernel_mapping",

src/kernels/layer/__init__.py

@@ -1,4 +1,10 @@
-from .device import Device, CUDAProperties
+from .device import CUDAProperties, Device
+from .func import (
+    FuncRepository,
+    LocalFuncRepository,
+    LockedFuncRepository,
+    use_kernel_func_from_hub,
+)
 from .kernelize import (
     kernelize,
     register_kernel_mapping,
@@ -16,13 +22,17 @@
 __all__ = [
     "CUDAProperties",
     "Device",
+    "FuncRepository",
     "LayerRepository",
+    "LocalFuncRepository",
     "LocalLayerRepository",
+    "LockedFuncRepository",
     "LockedLayerRepository",
     "Mode",
     "kernelize",
     "register_kernel_mapping",
     "replace_kernel_forward_from_hub",
     "use_kernel_forward_from_hub",
+    "use_kernel_func_from_hub",
     "use_kernel_mapping",
 ]