cuda-modules: surface capability-category helpers into _cuda.db and _cuda.lib#532065
cuda-modules: surface capability-category helpers into _cuda.db and _cuda.lib#532065jfroche wants to merge 2 commits into
Conversation
We want to categorize CUDA capabilities according to various criteria (e.g. whether they are architecture-specific, family-specific, or Jetson-specific) in order to enable more fine-grained filtering of capabilities in the backend and elsewhere. Currently, the logic for determining these categories is implemented as inline let-bindings in backendStdenv/default.nix, which makes it difficult to reuse this categorization logic elsewhere (e.g. in out-of-tree consumers). This change moves the the three inline filter computations that were local let-bindings in backendStdenv/default.nix into _cuda.db, alongside allSortedCudaCapabilities and cudaArchNameToCapabilities. Also add the corresponding _- prefixed functions to _cuda.lib for internal use within cuda-modules.
…uda.db Enable the previously added functions to be used in the public API of _cuda, so they can be safely called from lib or from any code that receives _cuda as an argument, without risking infinite recursion or relying on unstable, internal predicates. Add stable, supported names for the two predicates consumers are most likely to need, plus the list primitive they are built on: - cudaCapabilitiesJetsonSubset :: [CudaCapability] -> [CudaCapability] - cudaCapabilitiesAreJetson :: [CudaCapability] -> Bool - cudaCapabilitiesAreJetsonArch :: String -> [CudaCapability] -> Bool Also add db.jetsonArchNames :: [String] so callers can enumerate or validate arch names (e.g. "Ampere", "Blackwell") without hard-coding strings.
All of the custom I think the de-facto public bit for "is jetson" is: ]$ nix eval -f "<nixpkgs>" --argstr system aarch64-linux --arg config '{cudaCapabilities = [ "8.7"]; }' cudaPackages.flags.isJetsonBuild
trueThat said...
...still sounds like a regression
Heads-up: none of |
|
Sorry for the delay, I'm hoping to get around to this today. |
ConnorBaker
left a comment
There was a problem hiding this comment.
Sorry it took so long for me to get a chance to review this!
| let | ||
| jetsonSubset = | ||
| cudaCapabilities: lib.intersectLists _cuda.db.jetsonCudaCapabilities cudaCapabilities; | ||
| jetsonArchSubset = | ||
| archName: cudaCapabilities: | ||
| lib.intersectLists (_cuda.db.cudaArchNameToJetsonCapabilities.${archName} or [ ]) cudaCapabilities; | ||
| in |
There was a problem hiding this comment.
Inline and indirect through _cuda.lib so overrides/extensions function as expected.
|
|
||
| : The list of CUDA capabilities to check | ||
| */ | ||
| _cudaCapabilitiesAreArchitectureSpecific = |
There was a problem hiding this comment.
Name needs work -- does not reflect the fact that the function is the join (any satisfies), not the meet (all satisfy), which is implied from Are in the function name. This holds for all the other functions introduced here.
|
|
||
| : The list of CUDA capabilities to filter | ||
| */ | ||
| cudaCapabilitiesJetsonSubset = jetsonSubset; |
There was a problem hiding this comment.
This should be called something like getJetsonCudaCapabilities or filterJetsonCudaCapabilities.
Three capability-category filters (
architectureSpecificCudaCapabilities,familySpecificCudaCapabilities,jetsonCudaCapabilities) were computed as inline let-bindings inbackendStdenv/default.nix.The logic was not reusable outside
backendStdenv: downstream consumers that need "is this a Jetson build?" cannot readcudaPackages.backendStdenv.hasJetsonCudaCapabilityfrom the lib fixpoint without infinite recursion — it pulls the whole cudaPackages scope. They were forced to re-derive the logic from_cuda.dbthemselves.Commit 1: moves the three category lists into
_cuda.dbalongsideallSortedCudaCapabilitiesandcudaArchNameToCapabilities, adds _-prefixed predicates to _cuda.lib, and updates backendStdenv to inherit from_cuda.db. No behaviour change.Commit 2: adds stable public API to
_cuda.libfor out-of-tree consumers, alongside_cuda.db.jetsonArchNames