nix-builder: add a hook to detect incorrect op registrations (#569)

* nix-builder: add a hook to detect incorrect op registrations

* ci: test that the hook fails on incorrect registrations

* doc fixes

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* Use full kernel build for torch-ops-check-hook

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

Author: danieldk

docs/source/builder/writing-kernels.md

@@ -351,6 +351,47 @@ def mykernel(x: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tenso
     return out
 ```
 
+## Registering Torch operators
+
+You may want to register Torch ops from your kernel's Python code or
+register fake ops for `torch.compile` support. It is important to register
+such ops in the namespace that kernel-builder makes for your kernel
+build. This is required for compliant kernels to ensure that multiple
+versions of the same kernel can be loaded at the same time without
+namespace conflicts.
+
+You can use the `add_op_namespace_prefix` to prefix an op name with the
+correct prefix. So for instance, replace
+
+```python
+@torch.library.register_fake("relu::relu_fwd")
+def relu_fwd_fake(input: torch.Tensor) -> torch.Tensor:
+    return torch.empty_like(input)
+```
+
+by
+
+```python
+from ._ops import add_op_namespace_prefix
+
+@torch.library.register_fake(add_op_namespace_prefix("relu_fwd"))
+def relu_fwd_fake(input: torch.Tensor) -> torch.Tensor:
+    return torch.empty_like(input)
+```
+
+As mentioned in the above, the `_ops` module is generated by kernel-builder.
+
+kernel-builder uses a hook to reject incorrect usage of Torch op registration
+functions. However, it can only catch direct use of certain `torch.library`
+decorators. For instance, the hook would not reject the following decorator,
+so it should be seen as a last-resort check if human review failed:
+
+```python
+@some_indirection_for_register_fake("relu::relu_fwd")
+def relu_fwd_fake(input: torch.Tensor) -> torch.Tensor:
+    return torch.empty_like(input)
+```
+
 ## Kernel tests
 
 Kernel tests are stored in the `tests` directory. Since running all

examples/kernels/silu-and-mul-bad-registration/CARD.md

@@ -0,0 +1,56 @@
+---
+library_name: kernels
+{% if license %}license: {{ license }}
+{% endif %}---
+
+This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated.
+
+## How to use
+{% if functions %}
+
+```python
+# make sure `kernels` is installed: `pip install -U kernels`
+from kernels import get_kernel
+
+kernel_module = get_kernel("{{ repo_id }}", version={{ version }})
+{{ functions[0] }} = kernel_module.{{ functions[0] }}
+
+{{ functions[0] }}(...)
+```
+{% else %}
+
+Usage example not available.
+{% endif %}
+
+## Available functions
+{% if functions %}
+{% for func in functions %}
+- `{{ func }}`
+{% endfor %}
+{% else %}
+
+Function list not available.
+{% endif %}
+{% if layers %}
+
+## Available layers
+{% for layer in layers %}
+- `{{ layer }}`
+{% endfor %}
+{% endif %}
+
+## Benchmarks
+{% if has_benchmark %}
+
+Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }} --version {{ version }}`.
+{% else %}
+
+No benchmark available yet.
+{% endif %}
+{% if upstream %}
+
+## Source code
+
+Source code of this kernel originally comes from {{ upstream }} and it was repurposed for compatibility with `kernels`.
+{% endif %}
+

examples/kernels/silu-and-mul-bad-registration/build.toml

@@ -0,0 +1,18 @@
+[general]
+name = "silu-and-mul-bad-registration"
+version = 1
+license = "apache-2.0"
+backends = [
+    "cpu",
+    "cuda",
+    "metal",
+    "rocm",
+    "xpu",
+]
+
+[general.hub]
+repo-id = "kernels-test/silu-and-mul-bad-registration"
+
+[torch-noarch]
+
+[kernel]

examples/kernels/silu-and-mul-bad-registration/flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for kernels tests";
+
+  inputs = {
+    kernel-builder.url = "path:../../..";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

examples/kernels/silu-and-mul-bad-registration/torch-ext/silu_and_mul_bad_registration/__init__.py

@@ -0,0 +1,12 @@
+import torch
+
+from ._ops import ops
+from .op import _silu_and_mul
+from . import layers
+
+
+def silu_and_mul(x: torch.Tensor) -> torch.Tensor:
+    return ops.silu_and_mul(x)
+
+
+__all__ = ["layers", "silu_and_mul"]

examples/kernels/silu-and-mul-bad-registration/torch-ext/silu_and_mul_bad_registration/layers.py

@@ -0,0 +1,20 @@
+import torch
+import torch.nn as nn
+
+from ._ops import ops
+
+
+class SiluAndMul(nn.Module):
+    """
+    Apply SiLU to one half of the array and use it as a multiplicative
+    gate for the other half.
+
+    Shapes:
+        x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
+        return: (num_tokens, d) or (batch_size, seq_len, d)
+    """
+
+    can_torch_compile: bool = True
+
+    def forward(self, x: torch.Tensor):
+        return ops.silu_and_mul(x)

examples/kernels/silu-and-mul-bad-registration/torch-ext/silu_and_mul_bad_registration/op.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn.functional as F
+
+
+@torch.library.custom_op("myns::silu_and_mul", mutates_args=())
+def _silu_and_mul(x: torch.Tensor) -> torch.Tensor:
+    d = x.shape[-1] // 2
+    return F.silu(x[..., :d]) * x[..., d:]
+
+
+def backward(ctx, grad_output):
+    x = ctx.saved_tensors[0]
+    d = x.shape[-1] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    sigmoid_x1 = torch.sigmoid(x1)
+    silu_x1 = F.silu(x1)
+    dsilu_dx1 = sigmoid_x1 + silu_x1 * (1 - sigmoid_x1)
+    dx1 = grad_output * x2 * dsilu_dx1
+    dx2 = grad_output * silu_x1
+    return torch.cat([dx1, dx2], dim=-1)
+
+
+def setup_context(ctx, inputs, output):
+    (x,) = inputs
+    ctx.save_for_backward(x)
+
+
+_silu_and_mul.register_autograd(backward, setup_context=setup_context)
+
+
+@_silu_and_mul.register_fake
+def _(x: torch.Tensor) -> torch.Tensor:
+    return x.new_empty(x.shape[0], x.shape[1] // 2)

flake.nix

@@ -201,7 +201,8 @@
       in
       rec {
         checks.default = pkgs.callPackage ./nix-builder/lib/checks.nix {
-          inherit buildSets;
+          inherit buildSets self;
+          inherit (self.lib) genKernelFlakeOutputs;
           build = buildPerSystem.${system};
         };
 

nix-builder/lib/checks.nix

@@ -1,21 +1,38 @@
 {
+  self,
   lib,
   runCommand,
+  testers,
+  python3,
+  stdenv,
 
   build,
   buildSets,
+  genKernelFlakeOutputs,
 }:
 
 let
   kernelBuildSets = build.applicableBuildSets {
     inherit buildSets;
     path = ../../examples/kernels/relu-torch-bounds;
   };
+
+  badRegistrationCheck = testers.testBuildFailure' {
+    drv =
+      (genKernelFlakeOutputs {
+        inherit self;
+        path = ../../examples/kernels/silu-and-mul-bad-registration;
+      }).packages.${stdenv.hostPlatform.system}.redistributable.torch-cuda;
+    expectedBuilderExitCode = 1;
+    expectedBuilderLogEntries = [
+      "Found Torch library registrations that do not use `add_op_namespace_prefix`:"
+    ];
+  };
 in
 assert lib.assertMsg (builtins.all (buildSet: buildSet.torch.version == "2.10.0") kernelBuildSets)
   ''
     Torch minver/maxver filtering does not work.
   '';
-runCommand "builder-nix-checks" { } ''
+runCommand "builder-nix-checks" { buildInputs = [ badRegistrationCheck ]; } ''
   touch $out
 ''

nix-builder/lib/extension/torch/arch.nix

@@ -16,6 +16,7 @@
   get-kernel-check,
   kernel-abi-check,
   kernel-layout-check,
+  torch-ops-check,
   ninja,
   python3,
   remove-bytecode-hook,
@@ -163,6 +164,7 @@ stdenv.mkDerivation (prevAttrs: {
     kernel-abi-check
     kernel-layout-check
     remove-bytecode-hook
+    torch-ops-check
   ]
   ++ lib.optionals doGetKernelCheck [
     (get-kernel-check.override { python3 = python3.withPackages (ps: dependencies); })