feat: experimental kernel porting tool

[drbh]

Warning

This is an experiment/draft. The recipe language, the op set, and the CLI are all subject to change without notice.

this pr adds kernel-port, an experimental tool for porting kernel repos into the kernel-builder layout by running a recipe instead of doing it by hand.

the idea is that porting an existing kernel to the kernel-builder is a set of deterministic rewrites/restructuring. a recipe is a list of those operations, which has the benefit of being able to be checked into the repo and provide a way to deterministically reproduce given a specific upstream commit.

one of the difficulties of maintaining a port is that the upstream repo can drift and currently keeping the port in sync is a manual process.

recipes are a way to codify the porting process, so that if the upstream repo drifts we can simply bump the commit in the recipe and re-run the porting process. if the upstream repo has changed in a way that breaks the porting process, the recipe will fail to apply and we can fix it before continuing.

a port is a port.kdl recipe plus an overlay dir of checked in files. same pins + same recipe gives a byte identical tree every time, and every op hard fails on drift rather than silently porting the wrong thing.

recipes are kdl 2.0 documents, one node per op:

recipe version=1

source repo="https://github.com/org/repo" commit="<40 char sha>"
prune keep="csrc/**,pkg/**"
move from="csrc" to="kernel"
move from="pkg" to="torch-ext/pkg"
relativize_imports in="torch-ext/pkg/**" package_root="torch-ext/pkg" changes=6
kernel name="pkg" backend="cuda" src="kernel/**"
manifest name="pkg" backends="cuda" torch_src="torch-ext/*.cpp"

there are 15 ops (source, vendor, prune, delete, move, overlay, replace, strip_suffix, expect, convert_import, remap_module, relativize_imports, ensure_init, kernel, manifest). python rewrites go through libcst so comments and formatting are preserved byte for byte. build.toml is always generated by the manifest op, never overlaid.

the pins are the whole point. count=, files= and changes= are literals that have to match exactly, so a new file upstream cannot be rewritten without someone looking at it:

error: recipe line 17: relativize_imports: expected exactly 3 change(s) but made 4 - upstream drifted; review the new rewrites and update changes=

you can try an op without a checkout at all, -e takes the recipe inline and --file path=content builds the input tree in memory

cargo run -p kernel-port -- -e 'relativize_imports in="pkg/**" package_root="pkg" changes=1' \
    --file 'pkg/__init__.py=from pkg.ops import hello'

[line   1] relativize_imports  rewrote 1 import(s) in 1 file(s)
M pkg/__init__.py

FILE: pkg/__init__.py
from .ops import hello

the readme has a cookbook with one runnable command per op, plus the full arg list and failure modes for each. every example in it was run and pasted, not written by hand.

***NEXT STEPS are to target a JIT and AOT kernel in the kernels-community and experiment using this tool

[github-actions]

Coverage report — kernels/

Measured on: Python 3.10 / Torch 2.13.0.
Other CI configurations are not included in this number.
Hardware-gated code paths (ROCm/XPU/NPU/Darwin/Windows) are excluded or unreachable on the Linux+CUDA runner.

Total coverage: 84.3% — threshold: 80% — ✅

Per-file breakdown

Name	Stmts	Miss	Cover	Missing
src/kernels/__init__.py	13	0	100%
src/kernels/_system.py	6	1	83%	10
src/kernels/_versions.py	63	7	89%	46, 49, 52-53, 56-57, 100
src/kernels/backends.py	212	62	71%	40, 44, 48-51, 68, 90, 108, 117, 121, 125-127, 148, 157, 161, 165-167, 188, 199, 201, 208-211, 224, 228, 232-252, 260, 283-303
src/kernels/compat.py	8	1	88%	5
src/kernels/deps.py	58	4	93%	59-60, 101, 104
src/kernels/hf_hub.py	62	5	92%	18, 20, 114, 136-137
src/kernels/importer.py	48	3	94%	106, 110, 113
src/kernels/install.py	50	5	90%	78, 116-117, 155, 173
src/kernels/layer/__init__.py	6	0	100%
src/kernels/layer/_interval_tree.py	103	4	96%	23, 52, 147, 150
src/kernels/layer/device.py	48	14	71%	42, 47-49, 91, 96-98, 101, 149, 152, 155-157
src/kernels/layer/func.py	82	7	91%	83, 113, 185, 303, 309, 322, 340
src/kernels/layer/globals.py	5	0	100%
src/kernels/layer/kernelize.py	74	8	89%	255, 281, 289-290, 296, 300, 316-318
src/kernels/layer/layer.py	211	16	92%	169, 212, 218, 231, 339, 419-420, 432, 441, 449, 460, 489, 493, 506, 559, 589
src/kernels/layer/mode.py	14	0	100%
src/kernels/layer/repos.py	144	42	71%	27, 33, 36-43, 63-64, 70, 73-76, 90, 94, 103-104, 110, 113-116, 123-124, 130, 133-136, 143-144, 150, 153-156, 163-164, 170, 173-176, 257
src/kernels/load.py	74	9	88%	218, 224, 230-231, 252-264
src/kernels/locking.py	108	70	35%	42-109, 113-136, 144, 148-155, 159-169, 173-180
src/kernels/status.py	49	2	96%	23, 81
src/kernels/variants.py	278	22	92%	64, 95, 116, 146, 255-256, 298-301, 303, 387-394, 400-406, 437-443, 455-461, 611-613
src/kernels/verify.py	88	1	99%	32
TOTAL	1804	283	84%

Updated by the Test kernels workflow on commit f2c269d417a645b10e1f394ce98e7f74dcb8f7d8.

[sayakpaul]

an overlay dir of checked in files

We might want to elaborate on what overlay means in this context.

recipe version=1

Is the manual specification for version inspired by how we do it in build.toml of kernels?

the pins are the whole point. count=, files= and changes= are literals that have to match exactly, so a new file upstream cannot be rewritten without someone looking at it:

How are those pins derived?

you can try an op without a checkout at all, -e takes the recipe inline and --file path=content builds the input tree in memory

Do the users have to specify changes= manually like that? Feels cumbersome no?

More comments

• Yet another language 🥲
• For the code snippet provided in the description, I think we would want to first showcase the high-level structure of the kernel and then it'd be much easier to understand how those ops relate to one another in the context of the builder layout.
• I didn't find anything that registers the PyTorch ops as per the specifications from https://huggingface.co/docs/kernels/main/en/builder/writing-kernels#registering-torch-operators. Is that intentional? I think this should be an atomic op unless I am missing something.
• Do we have the possibility that overlay directory can grow big? If so, wouldn't it be a problem to PR that upstream where we essentially want to just PR with the KDL file along with a lightweight overlay folder if needed?

[drbh]

@sayakpaul

We might want to elaborate on what overlay means in this context.

the docs on the readme may already fill this need https://github.com/huggingface/kernels/blob/b7062eed8821f6a1235680844ad4bbc4711d7a7c/kernel-port/README.md#overlay

copied for reference

Note

overlay

overlay from="<dir>"

Copy a directory of checked-in files over the workspace, overwriting what is there. from is relative to the recipe file. This is where the files with no upstream equivalent live: torch_binding.cpp, flake.nix, a CARD.md. Keep it small - anything derivable from upstream should be an op, not an overlay file, so that upstream drift is detected rather than papered over.

**the readme in the pr has a "cookbook" that explains what each operator in the recipe does and what args are optional, etc..

Is the manual specification for version inspired by how we do it in build.toml of kernels?

the version is manually specified when writing the recipe and is only used so if the port tool has any breaking changes we can bump in the future

How are those pins derived?

these are added by the users (either by manually counting) or more realisitcally through an iterative process. the tool is designed to be run iteratively making changes to the recipe and checking the output. the tool errors if the count=, files= and changes= does not match exactly, and reports the count of mutations the op would make. this ensure that every change is accounted for and tracked (and not silently applied when migrating to a newer commit)

example of how it errors

error: recipe line 45: expect: expected "torch-ext/einops/**" to match exactly 10 file(s), found 17 - the upstream file set drifted; update the port definition (torch-ext/einops/__init__.py, torch-ext/einops/_backends.py, torch-ext/einops/_torch_specific.py, torch-ext/einops/array_api.py, torch-ext/einops/einops.py, torch-ext/einops/experimental/__init__.py, torch-ext/einops/layers/__init__.py, torch-ext/einops/layers/_einmix.py, torch-ext/einops/layers/flax.py, torch-ext/einops/layers/keras.py, torch-ext/einops/layers/oneflow.py, torch-ext/einops/layers/paddle.py, torch-ext/einops/layers/tensorflow.py, torch-ext/einops/layers/torch.py, torch-ext/einops/packing.py, torch-ext/einops/parsing.py, torch-ext/einops/py.typed)

Do the users have to specify changes= manually like that? Feels cumbersome no?

the changes= argument is optional in all cases https://github.com/huggingface/kernels/blob/b7062eed8821f6a1235680844ad4bbc4711d7a7c/kernel-port/README.md#rewriting-python-imports, but are intentionally verbose since pins (like changes) are included to make the number of mutations explicit and record.

in the readme under conventions we include the following (copied for reference) https://github.com/huggingface/kernels/blob/b7062eed8821f6a1235680844ad4bbc4711d7a7c/kernel-port/README.md#cookbook

Note

• Pins (count=, files=, changes=) are literal integers that must match exactly. They are the point: they turn "upstream changed" from a silent mis-port into a failed run.

so its recommended to include pins, however at the end of the day its fully optional

Yet another language 🥲

yea true, but at least its really tiny (only our 15 ops) and no control flow/etc so shouldn't be much effort to understand. I think the parts that can be confusing are concepts like pins but sufficient docs and good error messages should make this manageable

For the code snippet provided in the description, I think we would want to first showcase the high-level structure of the kernel and then it'd be much easier to understand how those ops relate to one another in the context of the builder layout.

added examples of a working port for einops for this reason in huggingface/kernels-community#1092 and how to update the pin in huggingface/kernels-community#1093

I didn't find anything that registers the PyTorch ops as per the specifications from https://huggingface.co/docs/kernels/main/en/builder/writing-kernels#registering-torch-operators. Is that intentional? I think this should be an atomic op unless I am missing something.

yes there are no recipe operations for torch ops specifically - this case is handled by the generic replace op, this way the exact string replacement for the ops is stored in the recipe (which makes them a bit longer), but the benefit of a generic op is that it works in all cases.

Do we have the possibility that overlay directory can grow big? If so, wouldn't it be a problem to PR that upstream where we essentially want to just PR with the KDL file along with a lightweight overlay folder if needed?

yes 100%, overlays will be big for kernels that are a combination of modified upstream and custom implementations we have in kernels-community. basically overlay is a way to inject files that follow kernel-builder conventions (no mutation needed). so in the case we are opening a pr for an external kernel, there is not really a case where we need to inject files (only the flake.nix) but if they dont already have a flake they can just add that to their root and we'll pull it in in the recipe.

tldr; in the pr case we'll just add two files port.kdl and flake.nix. the overlay is mostly a concept that caters to our situation where we have vendored kernels with non upstream additions.

[sayakpaul]

these are added by the users (either by manually counting) or more realisitcally through an iterative process. the tool is designed to be run iteratively making changes to the recipe and checking the output. the tool errors if the count=, files= and changes= does not match exactly, and reports the count of mutations the op would make. this ensure that every change is accounted for and tracked (and not silently applied when migrating to a newer commit)

Very good idea, IMO. But would we have a simple way to quickly test the correctness of its implementation? Not a blocker but I think we should strive for simplicity here.

so its recommended to include pins, however at the end of the day its fully optional

Maybe we need to distinguish between required pins and optional pins? count and filesfeel like they should be the required ones.

yes there are no recipe operations for torch ops specifically - this case is handled by the generic replace op, this way the exact string replacement for the ops is stored in the recipe (which makes them a bit longer), but the benefit of a generic op is that it works in all cases.

Cool, that works for me!

basically overlay is a way to inject files that follow kernel-builder conventions (no mutation needed). so in the case we are opening a pr for an external kernel, there is not really a case where we need to inject files (only the flake.nix) but if they dont already have a flake they can just add that to their root and we'll pull it in in the recipe.

Not sure if I fully understood it. Why would it differ for an external repo from kernels-community? Is it just because of the top-level flake?

tldr; in the pr case we'll just add two files port.kdl and flake.nix. the overlay is mostly a concept that caters to our situation where we have vendored kernels with non upstream additions.

Oh okay. I was under the impression that without the overlay folder (and with just the KDL recipe), it wouldn't build according to our compliance. So, the scope of overlay seemed a bit confusing to me. Possible to provide a bit more elaboration?

[sayakpaul]

Any reasoning behind making it a part of our nix-builder? For future CI?

[sayakpaul]

Thanks just left a bunch of comments. I think we are headed in a good direction. Once this takes a bit more shape, we could think about how we wire this in the CI, etc.