Skip to content

Commit

Permalink
Merge pull request #487 from johnnynunez/master
Browse files Browse the repository at this point in the history
initial support blackwell
  • Loading branch information
Tom94 authored Jan 25, 2025
2 parents 1e38444 + e66aa8c commit 394bfd2
Show file tree
Hide file tree
Showing 4 changed files with 23 additions and 11 deletions.
10 changes: 10 additions & 0 deletions .github/workflows/main.yml
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,12 @@ jobs:
strategy:
matrix:
include:
- os: ubuntu-24.04
cuda: "12.8"
arch: 120
- os: ubuntu-24.04
cuda: "12.8"
arch: 100
- os: ubuntu-24.04
cuda: "12.6"
arch: 89
Expand Down Expand Up @@ -72,6 +78,10 @@ jobs:
strategy:
matrix:
include:
- os: windows-2025
visual_studio: "Visual Studio 17 2022"
cuda: "12.8.0"
arch: 120
- os: windows-2025
visual_studio: "Visual Studio 17 2022"
cuda: "12.6.3"
Expand Down
4 changes: 3 additions & 1 deletion bindings/torch/setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,10 @@ def max_supported_compute_capability(cuda_version):
return 80
elif cuda_version < parse_version("11.8"):
return 86
else:
elif cuda_version < parse_version("12.8"):
return 90
else:
return 120

# Find version of tinycudann by scraping CMakeLists.txt
with open(os.path.join(ROOT_DIR, "CMakeLists.txt"), "r") as cmakelists:
Expand Down
2 changes: 1 addition & 1 deletion bindings/torch/tinycudann/modules.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@

import torch

ALL_COMPUTE_CAPABILITIES = [20, 21, 30, 35, 37, 50, 52, 53, 60, 61, 62, 70, 72, 75, 80, 86, 87, 89, 90]
ALL_COMPUTE_CAPABILITIES = [20, 21, 30, 35, 37, 50, 52, 53, 60, 61, 62, 70, 72, 75, 80, 86, 87, 89, 90, 100, 101, 120]

if not torch.cuda.is_available():
raise EnvironmentError("Unknown compute capability. Ensure PyTorch with CUDA support is installed.")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

# Dictionary of known cuda versions and thier download URLS, which do not follow a consistent pattern :(
$CUDA_KNOWN_URLS = @{
"8.0.44" = "http://developer.nvidia.com/compute/cuda/8.0/Prod/network_installers/cuda_8.0.44_win10_network-exe";
"8.0.44" = "http://developer.nvidia.com/compute/cuda/8.0/Prod/network_installers/cuda_8.0.44_win10_network-exe";
"8.0.61" = "http://developer.nvidia.com/compute/cuda/8.0/Prod2/network_installers/cuda_8.0.61_win10_network-exe";
"9.0.176" = "http://developer.nvidia.com/compute/cuda/9.0/Prod/network_installers/cuda_9.0.176_win10_network-exe";
"9.1.85" = "http://developer.nvidia.com/compute/cuda/9.1/Prod/network_installers/cuda_9.1.85_win10_network";
Expand All @@ -25,10 +25,11 @@ $CUDA_KNOWN_URLS = @{
"11.3.0" = "https://developer.download.nvidia.com/compute/cuda/11.3.0/network_installers/cuda_11.3.0_win10_network.exe";
"11.3.1" = "https://developer.download.nvidia.com/compute/cuda/11.3.1/network_installers/cuda_11.3.1_win10_network.exe";
"11.5.0" = "https://developer.download.nvidia.com/compute/cuda/11.5.0/network_installers/cuda_11.5.0_win10_network.exe";
"11.5.1" = "https://developer.download.nvidia.com/compute/cuda/11.5.1/network_installers/cuda_11.5.1_windows_network.exe";
"11.5.1" = "https://developer.download.nvidia.com/compute/cuda/11.5.1/network_installers/cuda_11.5.1_windows_network.exe";
"11.8.0" = "https://developer.download.nvidia.com/compute/cuda/11.8.0/network_installers/cuda_11.8.0_windows_network.exe";
"12.5.0" = "https://developer.download.nvidia.com/compute/cuda/12.5.0/network_installers/cuda_12.5.0_windows_network.exe";
"12.6.3" = "https://developer.download.nvidia.com/compute/cuda/12.6.3/network_installers/cuda_12.6.3_windows_network.exe";
"12.6.3" = "https://developer.download.nvidia.com/compute/cuda/12.6.3/network_installers/cuda_12.6.3_windows_network.exe";
"12.8.0" = "https://developer.download.nvidia.com/compute/cuda/12.8.0/network_installers/cuda_12.8.0_windows_network.exe";
}

# @todo - change this to be based on _MSC_VER intead, or invert it to be CUDA keyed instead?
Expand Down Expand Up @@ -74,7 +75,7 @@ $CUDA_PATCH=$Matches.patch
# Exit if visual studio is too new for the cuda version.
$VISUAL_STUDIO = $env:visual_studio.trim()
if ($VISUAL_STUDIO.length -ge 4) {
$VISUAL_STUDIO_YEAR = $VISUAL_STUDIO.Substring($VISUAL_STUDIO.Length-4)
$VISUAL_STUDIO_YEAR = $VISUAL_STUDIO.Substring($VISUAL_STUDIO.Length-4)
if ($VISUAL_STUDIO_YEAR.length -eq 4 -and $VISUAL_STUDIO_MIN_CUDA.containsKey($VISUAL_STUDIO_YEAR)){
$MINIMUM_CUDA_VERSION = $VISUAL_STUDIO_MIN_CUDA[$VISUAL_STUDIO_YEAR]
if ([version]$CUDA_VERSION_FULL -lt [version]$MINIMUM_CUDA_VERSION) {
Expand All @@ -99,15 +100,14 @@ $CUDA_PACKAGES = ""
# }
# }

Foreach ($package in $CUDA_PACKAGES_IN) {
foreach ($package in $CUDA_PACKAGES_IN) {
# Make sure the correct package name is used for nvcc.
if($package -eq "nvcc" -and [version]$CUDA_VERSION_FULL -lt [version]"9.1"){
$package="compiler"
} elseif($package -eq "compiler" -and [version]$CUDA_VERSION_FULL -ge [version]"9.1") {
$package="nvcc"
}
$CUDA_PACKAGES += " $($package)_$($CUDA_MAJOR).$($CUDA_MINOR)"

}
echo "$($CUDA_PACKAGES)"
## -----------------
Expand All @@ -116,9 +116,9 @@ echo "$($CUDA_PACKAGES)"

# Select the download link if known, otherwise have a guess.
$CUDA_REPO_PKG_REMOTE=""
if($CUDA_KNOWN_URLS.containsKey($CUDA_VERSION_FULL)){
if ($CUDA_KNOWN_URLS.containsKey($CUDA_VERSION_FULL)){
$CUDA_REPO_PKG_REMOTE=$CUDA_KNOWN_URLS[$CUDA_VERSION_FULL]
} else{
} else {
# Guess what the url is given the most recent pattern (at the time of writing, 10.1)
Write-Output "note: URL for CUDA ${$CUDA_VERSION_FULL} not known, estimating."
$CUDA_REPO_PKG_REMOTE="http://developer.download.nvidia.com/compute/cuda/$($CUDA_MAJOR).$($CUDA_MINOR)/Prod/network_installers/cuda_$($CUDA_VERSION_FULL)_win10_network.exe"
Expand All @@ -133,7 +133,7 @@ $CUDA_REPO_PKG_LOCAL="cuda_$($CUDA_VERSION_FULL)_win10_network.exe"
# Get CUDA network installer
Write-Output "Downloading CUDA Network Installer for $($CUDA_VERSION_FULL) from: $($CUDA_REPO_PKG_REMOTE)"
Invoke-WebRequest $CUDA_REPO_PKG_REMOTE -OutFile $CUDA_REPO_PKG_LOCAL | Out-Null
if(Test-Path -Path $CUDA_REPO_PKG_LOCAL){
if (Test-Path -Path $CUDA_REPO_PKG_LOCAL){
Write-Output "Downloading Complete"
} else {
Write-Output "Error: Failed to download $($CUDA_REPO_PKG_LOCAL) from $($CUDA_REPO_PKG_REMOTE)"
Expand Down

0 comments on commit 394bfd2

Please sign in to comment.