From e328dd36b06a3fbce5942d6c988e75c5cf7fcec6 Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Sat, 30 May 2026 13:40:34 -0700 Subject: [PATCH 1/3] Some fixes for rocm. --- openequivariance/CMakeLists.txt | 2 +- openequivariance/openequivariance/_torch/extlib/__init__.py | 3 +-- .../openequivariance/extension/libtorch_tp_jit_stable.cpp | 3 --- 3 files changed, 2 insertions(+), 6 deletions(-) diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index 232233f..f0fac08 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -158,8 +158,8 @@ if(hip_FOUND) ) set(HIP_LINK_LIBS - hiprtc hip_stub_lib + hip::host ) add_stable_extension(torch_stable_hip HIP_BACKEND "${HIP_LINK_LIBS}") endif() diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index ef17724..3812655 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -89,7 +89,6 @@ def load_jit_extension(): extra_cflags.append("-DCUDA_BACKEND") elif torch.version.hip: - extra_link_args.extend(["-lhiprtc"]) torch_libs = library_paths("cuda")[0] extra_link_args.append("-Wl,-rpath," + torch_libs) extra_cflags.append("-DHIP_BACKEND") @@ -207,4 +206,4 @@ def DeviceProp(*args, **kwargs): _raise_import_error_helper("DeviceProp") def GPUTimer(*args, **kwargs): - _raise_import_error_helper("GPUTimer") + _raise_import_error_helper("GPUTimer") \ No newline at end of file diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp index 440b764..abece92 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp @@ -11,9 +11,6 @@ #include #include -#ifdef HIP_BACKEND - #include -#endif using Tensor = torch::stable::Tensor; using Dtype = torch::headeronly::ScalarType; From 140e04b428e6ba8c586730dd2ec5804540330a05 Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Sat, 30 May 2026 13:42:28 -0700 Subject: [PATCH 2/3] More minor fixes. --- openequivariance/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index f0fac08..e7fac87 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -159,7 +159,7 @@ if(hip_FOUND) set(HIP_LINK_LIBS hip_stub_lib - hip::host + hip::host ) add_stable_extension(torch_stable_hip HIP_BACKEND "${HIP_LINK_LIBS}") endif() From 2e378990e144052692ad7a5ef0ef557df835f386 Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Sat, 30 May 2026 13:42:46 -0700 Subject: [PATCH 3/3] Ruff. --- openequivariance/openequivariance/_torch/extlib/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index 3812655..7326230 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -206,4 +206,4 @@ def DeviceProp(*args, **kwargs): _raise_import_error_helper("DeviceProp") def GPUTimer(*args, **kwargs): - _raise_import_error_helper("GPUTimer") \ No newline at end of file + _raise_import_error_helper("GPUTimer")