opensplat: cleanup, fix rocm build (#538248)

This commit is contained in:
Gaétan Lepage
2026-07-04 09:11:47 +00:00
committed by GitHub

View File

@@ -20,20 +20,22 @@
fetchpatch2,
}:
let
version = "1.1.4";
torch = python3.pkgs.torch.override { inherit cudaSupport rocmSupport; };
# Using a normal stdenv with cuda torch gives
# ld: /nix/store/k1l7y96gv0nc685cg7i3g43i4icmddzk-python3.11-torch-2.2.1-lib/lib/libc10.so: undefined reference to `std::ios_base_library_init()@GLIBCXX_3.4.32'
stdenv' = if cudaSupport then cudaPackages.backendStdenv else stdenv;
in
stdenv'.mkDerivation {
stdenv'.mkDerivation (finalAttrs: {
pname = "opensplat";
inherit version;
version = "1.1.4";
__structuredAttrs = true;
strictDeps = true;
src = fetchFromGitHub {
owner = "pierotofy";
repo = "OpenSplat";
tag = "v${version}";
tag = "v${finalAttrs.version}";
hash = "sha256-u2UmD0O3sUWELYb4CjQE19i4HUjLMcaWqOinQH0PPTM=";
};
@@ -44,24 +46,34 @@ stdenv'.mkDerivation {
})
];
postPatch = lib.optionalString rocmSupport ''
# ROCm CMake targets must be available before find_package(Torch)
# because Torch's Caffe2Targets.cmake references them in torch_hip_library
substituteInPlace CMakeLists.txt \
--replace-fail "find_package(Torch REQUIRED)" \
"find_package(hip REQUIRED)
find_package(hiprtc REQUIRED)
find_package(hipblas REQUIRED)
find_package(hipfft REQUIRED)
find_package(hiprand REQUIRED)
find_package(hipsparse REQUIRED)
find_package(hipsolver REQUIRED)
find_package(hipblaslt REQUIRED)
find_package(rocblas REQUIRED)
find_package(rocsolver REQUIRED)
find_package(miopen REQUIRED)
find_package(Torch REQUIRED)"
'';
postPatch = lib.optionalString rocmSupport (
# Recent PyTorch HIP builds no longer hipify the caching allocator namespace:
# the symbol stays c10::cuda::CUDACachingAllocator instead of c10::hip::HIPCachingAllocator.
''
substituteInPlace model.cpp \
--replace-fail \
"c10::hip::HIPCachingAllocator::emptyCache();" \
"c10::cuda::CUDACachingAllocator::emptyCache();"
''
# ROCm CMake targets must be available before find_package(Torch)
# because Torch's Caffe2Targets.cmake references them in torch_hip_library
+ ''
substituteInPlace CMakeLists.txt \
--replace-fail "find_package(Torch REQUIRED)" \
"find_package(hip REQUIRED)
find_package(hiprtc REQUIRED)
find_package(hipblas REQUIRED)
find_package(hipfft REQUIRED)
find_package(hiprand REQUIRED)
find_package(hipsparse REQUIRED)
find_package(hipsolver REQUIRED)
find_package(hipblaslt REQUIRED)
find_package(rocblas REQUIRED)
find_package(rocsolver REQUIRED)
find_package(miopen REQUIRED)
find_package(Torch REQUIRED)"
''
);
nativeBuildInputs = [
cmake
@@ -107,7 +119,11 @@ stdenv'.mkDerivation {
NIX_LDFLAGS = "-L${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs"; # fixes -lcuda not found
}
// lib.optionalAttrs rocmSupport {
HIPFLAGS = "-I${lib.getInclude rocmPackages.rocthrust}/include -I${lib.getInclude rocmPackages.rocprim}/include";
PYTORCH_ROCM_ARCH = torch.gpuTargetString;
HIPFLAGS = toString [
"-I${lib.getInclude rocmPackages.rocprim}/include"
"-I${lib.getInclude rocmPackages.rocthrust}/include"
];
};
cmakeFlags = [
@@ -139,4 +155,4 @@ stdenv'.mkDerivation {
];
platforms = lib.platforms.linux ++ lib.optionals (!cudaSupport) lib.platforms.darwin;
};
}
})