feat: support rust cpu kernels

[drbh]

this pr adds support for cpu rust kernels

note this pr supersedes #724

[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: 87.2% — threshold: 80% — ✅

Per-file breakdown

Name	Stmts	Miss	Cover	Missing
src/kernels/__init__.py	14	0	100%
src/kernels/_system.py	6	1	83%	10
src/kernels/_versions.py	130	14	89%	53, 59-60, 63-64, 102, 165-170, 199, 219
src/kernels/archs.py	56	1	98%	94
src/kernels/backends.py	213	62	71%	42, 46, 50-53, 70, 92, 110, 119, 123, 127-129, 150, 159, 163, 167-169, 190, 201, 203, 210-213, 226, 230, 234-254, 262, 285-305
src/kernels/compat.py	9	1	89%	5
src/kernels/deps.py	70	1	99%	56
src/kernels/hf_hub.py	63	2	97%	21, 23
src/kernels/importer.py	44	5	89%	80, 84, 87, 101-102
src/kernels/install.py	21	7	67%	76-100
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	85	6	93%	90, 115, 191, 311, 338, 368
src/kernels/layer/globals.py	5	0	100%
src/kernels/layer/kernelize.py	80	8	90%	258, 293, 301-302, 308, 312, 328-330
src/kernels/layer/layer.py	215	14	93%	182, 229, 256, 390, 470-471, 492, 500, 511, 540, 544, 557, 610, 640
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	71	2	97%	338, 378
src/kernels/locking.py	89	64	28%	35-83, 91-98, 102-125, 137, 152-159, 165-175, 179-186
src/kernels/python_deps.py	58	6	90%	59-60, 64-65, 101, 104
src/kernels/resolver.py	156	2	99%	220, 226
src/kernels/status.py	50	2	96%	25, 79
src/kernels/validate.py	88	5	94%	9, 100, 167, 190-191
src/kernels/variants.py	278	17	94%	65, 96, 117, 147, 256-257, 299-302, 304, 388-394, 400-406, 455-461
src/kernels/verify.py	127	6	95%	46, 202-204, 318-319
TOTAL	2243	286	87%

Updated by the Test kernels workflow on commit c4d061ba813695b4d5accfd8f17e4bbf90bdcd42.

[danieldk]

Maybe it's possible to put cxx_flags/include and cargo_manifest into respectively Language::Cpp and Language::Rust? I think in general using typing to exclude impossible combinations is better than runtime checks, which can easily be forgotten. May make things a bit more nasty on the Nix side though.

If we create v6 of the config format, then it could have it structured like that (we have to experiment a bit what the TOML looks like) and for v5 we know that language is always C++ (and put the flags in the new nested structure).

[danieldk]

Not 100% sure if this is the direction we want to go, but it might be worth exploring.

[drbh]

this is a great suggestion, in the latest change I updated Kernel::Cpu to include a CpuLanguage enum and then check for invalid combinations when creating the CpuLanguage type. while we still error at parse time, it seemed better to move the check closer to the relevant type and remove Build::validate.

curious what you think! happy to revert or change in anyway

[danieldk]

Not super-happy with these global variables, but I also don't see a nice way around it.

[drbh]

yea totally agree

[danieldk]

Hmmm, but we'd have cxx_flags for CUDA, XPU, etc. as well. So I was hoping that we could have a Language enum that can be reused for all backends.

I think the tricky part here is that there is some coupling between the language and the backend. E.g., the language cpp needs cuda_flags, but only when the backend is cuda. But it is not true the other way around, because cuda could in the future also be Rust and cuda_flags might not apply when the language is some Rust framework that compiles to PTX.

Just a wild idea that I haven't thought much through yet. Currently we have:

pub enum Kernel {
  Cpu { ... },
  Cuda { ... },
  Xpu { ... },
}

Maybe we should instead have

pub enum Kernel {
  CppCpu { ... },
  RustCpu { ... },
  CppCuda { ... },
  CppXpu { ... },
}

And then create v6 (no changes to v5) with

#[serde(deny_unknown_fields, rename_all = "kebab-case", tag = "language")]
pub enum Kernel {
  // ...
}

So that e.g. a bit of the build.toml would look like this:

[kernel.relu]                                                                                                                                                                                                                        
language = "cpp-cuda"                                                                                                                                                                                                                     
depends = ["torch"]                                                                                                                                                                                                                  
src = ["relu_cuda/relu.cu"]
cxx-flags = ["..."]
cuda-flags = ["..."]

The Kernel enum could then add a backend method to get the backend.

With this approach we have the benefit that we can always nicely encode the flags that are unique to the language + backend combination.

What do you think?

(If this is the way to go, it would probably be better to create a separate PR with this change and then rebase this PR on top of it once merged.)