pederbs/just-nixpkgs
0
2
Fork 0
Code Issues Pull Requests Releases Activity

overlays.withCudaOrRocm WIP

Browse Source
This commit is contained in:
Peder Bergebakken Sundt
2025-04-24 19:38:13 +02:00
parent 4157beb068
commit cf228eda6e
1 changed files with 69 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
overlays.nix
View File

@@ -148,6 +148,75 @@ let
}
);
# TODO: this way of overriding does not work
# usage: nix-build ./dev.nix -A <package>.withCuda
# usage: nix-build ./dev.nix -A <package>.withRocm
overlays.withCudaOrRocm = final: prev: {
lib = prev.lib.extend (
finalLib: prevLib: {
# makeOverridable :: (AttrSet -> a) -> AttrSet -> a
makeOverridable =
f: args:
let
fArgs = lib.functionArgs f;
# Creates a functor with the same arguments as f
mirrorArgs = lib.mirrorFunctionArgs f;
fWithCuda = mirrorArgs (
origArgs:
f' (
origArgs
// lib.optionalAttrs (fArgs ? cudaSupport) {
cudaSupport = true;
}
// lib.optionalAttrs (fArgs ? config) {
config = fArgs.config // {
cudaSupport = true;
};
}
)
);
fWithRocm = mirrorArgs (
origArgs:
f' (
origArgs
// lib.optionalAttrs (fArgs ? rocmSupport) {
rocmSupport = true;
}
// lib.optionalAttrs (fArgs ? config) {
config = fArgs.config // {
rocmSupport = true;
};
}
)
);
f' = mirrorArgs (
origArgs:
let
result = f origArgs;
in
if lib.isAttrs result then
if result ? overrideAttrs then
result.overrideAttrs (old: {
passthru = {
withCuda = fWithCuda origArgs;
withRocm = fWithRocm origArgs;
} // old.passthru or { };
})
else
{
withCuda = fWithCuda origArgs;
withRocm = fWithRocm origArgs;
}
// result
else
result
);
in
prevLib.makeOverridable f' args;
}
);
};
# very hacky, not guaranteed to work, but may save a lot of .whl rebuilds
# usage: nix-build ./dev.nix -A python3Packages.<package>.twostage
# usage: nix-build ./dev.nix -A python3Packages.<package>.twostage.first