Skip to content

Upgrade to 5090 machine and CUDA 12.8 #103

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Lantianyou wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Upgrade to 5090 machine and CUDA 12.8 #103

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
340e9de
feat: run this project on 5090 machine
Lantianyou Feb 20, 2026
af74680
feat: change build way
Lantianyou Feb 20, 2026
f79f887
feat: make sure run on 5090 machine
Zakary-D Feb 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,20 @@ pip install .
./download_models_and_data.sh
```

### Quick Start (Ubuntu + RTX 5090 + CUDA 12.8, no viewer)
Use the setup helper to create a fresh conda environment, install PyTorch `cu128`,
build DPVO CUDA extensions, and run a smoke test:
```bash
bash scripts/setup_ubuntu_5090_cuda131.sh
```

Optional flags (via env vars):
```bash
ENV_NAME=dpvo5090 RUN_DEMO=1 bash scripts/setup_ubuntu_5090_cuda131.sh
```

This path does not require `torch-scatter`.


### Recommended - Install the Pangolin Viewer
Note: You will need to have CUDA 11 and CuDNN installed on your system.
Expand Down
8 changes: 4 additions & 4 deletions dpvo/altcorr/correlation_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -208,7 +208,7 @@ std::vector<torch::Tensor> corr_cuda_forward(
auto opts = fmap1.options();
auto corr = torch::empty({B, M, D, D, H, W}, opts);

AT_DISPATCH_FLOATING_TYPES_AND_HALF(fmap1.type(), "corr_forward_kernel", ([&] {
AT_DISPATCH_FLOATING_TYPES_AND_HALF(fmap1.scalar_type(), "corr_forward_kernel", ([&] {
corr_forward_kernel<scalar_t><<<BLOCKS(B * M * H * W * D * D), THREADS>>>(radius,
fmap1.packed_accessor32<scalar_t,5,torch::RestrictPtrTraits>(),
fmap2.packed_accessor32<scalar_t,5,torch::RestrictPtrTraits>(),
Expand Down Expand Up @@ -270,7 +270,7 @@ std::vector<torch::Tensor> corr_cuda_backward(
auto fmap1_grad = torch::zeros_like(fmap1);
auto fmap2_grad = torch::zeros_like(fmap2);

AT_DISPATCH_FLOATING_TYPES_AND_HALF(fmap1.type(), "corr_backward_kernel", ([&] {
AT_DISPATCH_FLOATING_TYPES_AND_HALF(fmap1.scalar_type(), "corr_backward_kernel", ([&] {
corr_backward_kernel<scalar_t><<<BLOCKS(B * M * H * W * D * D), THREADS>>>(radius,
fmap1.packed_accessor32<scalar_t,5,torch::RestrictPtrTraits>(),
fmap2.packed_accessor32<scalar_t,5,torch::RestrictPtrTraits>(),
Expand All @@ -296,7 +296,7 @@ std::vector<torch::Tensor> patchify_cuda_forward(
auto opts = net.options();
auto patches = torch::zeros({B, M, C, D, D}, opts);

AT_DISPATCH_FLOATING_TYPES_AND_HALF(net.type(), "patchify_forward_kernel", ([&] {
AT_DISPATCH_FLOATING_TYPES_AND_HALF(net.scalar_type(), "patchify_forward_kernel", ([&] {
patchify_forward_kernel<scalar_t><<<BLOCKS(B * M * D * D), THREADS>>>(radius,
net.packed_accessor32<scalar_t,4,torch::RestrictPtrTraits>(),
coords.packed_accessor32<float,3,torch::RestrictPtrTraits>(),
Expand All @@ -322,7 +322,7 @@ std::vector<torch::Tensor> patchify_cuda_backward(

torch::Tensor net_gradient = torch::zeros_like(net);

AT_DISPATCH_FLOATING_TYPES_AND_HALF(net.type(), "patchify_backward_kernel", ([&] {
AT_DISPATCH_FLOATING_TYPES_AND_HALF(net.scalar_type(), "patchify_backward_kernel", ([&] {
patchify_backward_kernel<scalar_t><<<BLOCKS(B * M * D * D), THREADS>>>(radius,
gradient.packed_accessor32<scalar_t,5,torch::RestrictPtrTraits>(),
coords.packed_accessor32<float,3,torch::RestrictPtrTraits>(),
Expand Down
2 changes: 1 addition & 1 deletion dpvo/ba.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import torch
from torch_scatter import scatter_sum
from .scatter import scatter_sum

from . import fastba
from . import lietorch
Expand Down
10 changes: 5 additions & 5 deletions dpvo/blocks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
import torch.nn as nn
import torch.nn.functional as F

import torch_scatter
from .scatter import scatter_softmax, scatter_sum

class LayerNorm1D(nn.Module):
def __init__(self, dim):
Expand Down Expand Up @@ -39,8 +39,8 @@ def __init__(self, dim=512, expand=True):

def forward(self, x, ix):
_, jx = torch.unique(ix, return_inverse=True)
w = torch_scatter.scatter_softmax(self.g(x), jx, dim=1)
y = torch_scatter.scatter_sum(self.f(x) * w, jx, dim=1)
w = scatter_softmax(self.g(x), jx, dim=1)
y = scatter_sum(self.f(x) * w, jx, dim=1)

if self.expand:
return self.h(y)[:,jx]
Expand All @@ -58,8 +58,8 @@ def __init__(self, dim=512, expand=True):

def forward(self, x, ix):
_, jx = torch.unique(ix, return_inverse=True)
w = torch_scatter.scatter_softmax(self.g(x), jx, dim=1)
y = torch_scatter.scatter_sum(self.f(x) * w, jx, dim=1)
w = scatter_softmax(self.g(x), jx, dim=1)
y = scatter_sum(self.f(x) * w, jx, dim=1)

if self.expand:
return self.h(y)[:,jx]
Expand Down
5 changes: 1 addition & 4 deletions dpvo/lietorch/include/dispatch.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,7 @@

#define DISPATCH_GROUP_AND_FLOATING_TYPES(GROUP_INDEX, TYPE, NAME, ...) \
[&] { \
const auto& the_type = TYPE; \
/* don't use TYPE again in case it is an expensive or side-effect op */ \
at::ScalarType _st = ::detail::scalar_type(the_type); \
const at::ScalarType _st = TYPE; \
switch (_st) { \
PRIVATE_CASE_TYPE(GROUP_INDEX, at::ScalarType::Double, double, __VA_ARGS__) \
PRIVATE_CASE_TYPE(GROUP_INDEX, at::ScalarType::Float, float, __VA_ARGS__) \
Expand All @@ -45,4 +43,3 @@
}()

#endif

38 changes: 19 additions & 19 deletions dpvo/lietorch/src/lietorch_cpu.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -357,7 +357,7 @@ torch::Tensor exp_forward_cpu(int group_id, torch::Tensor a) {
int batch_size = a.size(0);
torch::Tensor X;

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.type(), "exp_forward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.scalar_type(), "exp_forward_kernel", ([&] {
X = torch::zeros({batch_size, group_t::N}, a.options());
exp_forward_kernel<group_t, scalar_t>(
a.data_ptr<scalar_t>(),
Expand All @@ -372,7 +372,7 @@ std::vector<torch::Tensor> exp_backward_cpu(int group_id, torch::Tensor grad, to
int batch_size = a.size(0);
torch::Tensor da = torch::zeros(a.sizes(), grad.options());

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.type(), "exp_backward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.scalar_type(), "exp_backward_kernel", ([&] {
exp_backward_kernel<group_t, scalar_t>(
grad.data_ptr<scalar_t>(),
a.data_ptr<scalar_t>(),
Expand All @@ -387,7 +387,7 @@ torch::Tensor log_forward_cpu(int group_id, torch::Tensor X) {
int batch_size = X.size(0);
torch::Tensor a;

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "log_forward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "log_forward_kernel", ([&] {
a = torch::zeros({batch_size, group_t::K}, X.options());
log_forward_kernel<group_t, scalar_t>(
X.data_ptr<scalar_t>(),
Expand All @@ -402,7 +402,7 @@ std::vector<torch::Tensor> log_backward_cpu(int group_id, torch::Tensor grad, to
int batch_size = X.size(0);
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "log_backward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "log_backward_kernel", ([&] {
log_backward_kernel<group_t, scalar_t>(
grad.data_ptr<scalar_t>(),
X.data_ptr<scalar_t>(),
Expand All @@ -417,7 +417,7 @@ torch::Tensor inv_forward_cpu(int group_id, torch::Tensor X) {
int batch_size = X.size(0);
torch::Tensor Y = torch::zeros_like(X);

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "inv_forward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "inv_forward_kernel", ([&] {
inv_forward_kernel<group_t, scalar_t>(
X.data_ptr<scalar_t>(),
Y.data_ptr<scalar_t>(),
Expand All @@ -431,7 +431,7 @@ std::vector<torch::Tensor> inv_backward_cpu(int group_id, torch::Tensor grad, to
int batch_size = X.size(0);
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "inv_backward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "inv_backward_kernel", ([&] {
inv_backward_kernel<group_t, scalar_t>(
grad.data_ptr<scalar_t>(),
X.data_ptr<scalar_t>(),
Expand All @@ -447,7 +447,7 @@ torch::Tensor mul_forward_cpu(int group_id, torch::Tensor X, torch::Tensor Y) {
int batch_size = X.size(0);
torch::Tensor Z = torch::zeros_like(X);

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "mul_forward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "mul_forward_kernel", ([&] {
mul_forward_kernel<group_t, scalar_t>(
X.data_ptr<scalar_t>(),
Y.data_ptr<scalar_t>(),
Expand All @@ -463,7 +463,7 @@ std::vector<torch::Tensor> mul_backward_cpu(int group_id, torch::Tensor grad, to
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
torch::Tensor dY = torch::zeros(Y.sizes(), grad.options());

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "mul_backward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "mul_backward_kernel", ([&] {
mul_backward_kernel<group_t, scalar_t>(
grad.data_ptr<scalar_t>(),
X.data_ptr<scalar_t>(),
Expand All @@ -480,7 +480,7 @@ torch::Tensor adj_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
int batch_size = X.size(0);
torch::Tensor b = torch::zeros(a.sizes(), a.options());

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "adj_forward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adj_forward_kernel", ([&] {
adj_forward_kernel<group_t, scalar_t>(
X.data_ptr<scalar_t>(),
a.data_ptr<scalar_t>(),
Expand All @@ -496,7 +496,7 @@ std::vector<torch::Tensor> adj_backward_cpu(int group_id, torch::Tensor grad, to
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
torch::Tensor da = torch::zeros(a.sizes(), grad.options());

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "adj_backward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adj_backward_kernel", ([&] {
adj_backward_kernel<group_t, scalar_t>(
grad.data_ptr<scalar_t>(),
X.data_ptr<scalar_t>(),
Expand All @@ -514,7 +514,7 @@ torch::Tensor adjT_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
int batch_size = X.size(0);
torch::Tensor b = torch::zeros(a.sizes(), a.options());

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "adjT_forward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adjT_forward_kernel", ([&] {
adjT_forward_kernel<group_t, scalar_t>(
X.data_ptr<scalar_t>(),
a.data_ptr<scalar_t>(),
Expand All @@ -530,7 +530,7 @@ std::vector<torch::Tensor> adjT_backward_cpu(int group_id, torch::Tensor grad, t
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
torch::Tensor da = torch::zeros(a.sizes(), grad.options());

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "adjT_backward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adjT_backward_kernel", ([&] {
adjT_backward_kernel<group_t, scalar_t>(
grad.data_ptr<scalar_t>(),
X.data_ptr<scalar_t>(),
Expand All @@ -548,7 +548,7 @@ torch::Tensor act_forward_cpu(int group_id, torch::Tensor X, torch::Tensor p) {
int batch_size = X.size(0);
torch::Tensor q = torch::zeros(p.sizes(), p.options());

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "act_forward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act_forward_kernel", ([&] {
act_forward_kernel<group_t, scalar_t>(
X.data_ptr<scalar_t>(),
p.data_ptr<scalar_t>(),
Expand All @@ -564,7 +564,7 @@ std::vector<torch::Tensor> act_backward_cpu(int group_id, torch::Tensor grad, to
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
torch::Tensor dp = torch::zeros(p.sizes(), grad.options());

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "act_backward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act_backward_kernel", ([&] {
act_backward_kernel<group_t, scalar_t>(
grad.data_ptr<scalar_t>(),
X.data_ptr<scalar_t>(),
Expand All @@ -582,7 +582,7 @@ torch::Tensor act4_forward_cpu(int group_id, torch::Tensor X, torch::Tensor p) {
int batch_size = X.size(0);
torch::Tensor q = torch::zeros(p.sizes(), p.options());

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "act4_forward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act4_forward_kernel", ([&] {
act4_forward_kernel<group_t, scalar_t>(
X.data_ptr<scalar_t>(),
p.data_ptr<scalar_t>(),
Expand All @@ -598,7 +598,7 @@ std::vector<torch::Tensor> act4_backward_cpu(int group_id, torch::Tensor grad, t
torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
torch::Tensor dp = torch::zeros(p.sizes(), grad.options());

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "act4_backward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act4_backward_kernel", ([&] {
act4_backward_kernel<group_t, scalar_t>(
grad.data_ptr<scalar_t>(),
X.data_ptr<scalar_t>(),
Expand All @@ -616,7 +616,7 @@ torch::Tensor as_matrix_forward_cpu(int group_id, torch::Tensor X) {
int batch_size = X.size(0);
torch::Tensor T4x4 = torch::zeros({X.size(0), 4, 4}, X.options());

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "as_matrix_forward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "as_matrix_forward_kernel", ([&] {
as_matrix_forward_kernel<group_t, scalar_t>(
X.data_ptr<scalar_t>(),
T4x4.data_ptr<scalar_t>(),
Expand All @@ -631,7 +631,7 @@ torch::Tensor orthogonal_projector_cpu(int group_id, torch::Tensor X) {
int batch_size = X.size(0);
torch::Tensor P;

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "orthogonal_projector_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "orthogonal_projector_kernel", ([&] {
P = torch::zeros({X.size(0), group_t::N, group_t::N}, X.options());
orthogonal_projector_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(), P.data_ptr<scalar_t>(), batch_size);
}));
Expand All @@ -645,7 +645,7 @@ torch::Tensor jleft_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a)
int batch_size = X.size(0);
torch::Tensor b = torch::zeros(a.sizes(), a.options());

DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "jleft_forward_kernel", ([&] {
DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "jleft_forward_kernel", ([&] {
jleft_forward_kernel<group_t, scalar_t>(
X.data_ptr<scalar_t>(),
a.data_ptr<scalar_t>(),
Expand Down
Loading