AMD-Ecosystem · awelling2801 · Aug 21, 2026 · Aug 21, 2026 · Aug 21, 2026
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -576,7 +576,7 @@ if (BUILD_SPARSE)
         GPU_TARGETS=${CMAKE_HIP_ARCHITECTURES}
         CMAKE_C_COMPILER=${CMAKE_C_COMPILER}
         CMAKE_CXX_COMPILER=${CMAKE_CXX_COMPILER}
-        CMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH}
+        CMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH}"
         CMAKE_HIP_ARCHITECTURES=${CMAKE_HIP_ARCHITECTURES}
         BINDIR=${CMAKE_CURRENT_BINARY_DIR}
         INCLUDEDIR="${DGL_INCLUDE_DIRS}"
@@ -641,7 +641,6 @@ if (BUILD_GRAPHBOLT)
         PYTORCH_ROCM_ARCH=${CMAKE_HIP_ARCHITECTURES}
         CMAKE_C_COMPILER=${CMAKE_C_COMPILER}
         CMAKE_CXX_COMPILER=${CMAKE_CXX_COMPILER}
-        CMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH}
         GPU_TARGETS=${CMAKE_HIP_ARCHITECTURES}
         CMAKE_HIP_ARCHITECTURES=${CMAKE_HIP_ARCHITECTURES}
         USE_HIP=${USE_HIP}

diff --git a/docker/Dockerfile.ci_gpu_rocm b/docker/Dockerfile.ci_gpu_rocm
@@ -1,3 +1,6 @@
+# Deprecated: use docker/Dockerfile.rocm instead.
+# This filename is kept for older build scripts.
+#
 # "Copyright Advanced Micro Devices, Inc.
 # Licensed under the Apache License Version 2.0"
 
@@ -13,16 +16,16 @@ ENV MAX_JOBS=${ARG_MAX_JOBS:-24}
 
 # Install basic tools
 RUN apt-get update \
-    && apt install -y clang-format-15 ninja-build tree 
+    && apt install -y clang-format-15 git ninja-build patch tree
 
-#Updating to latest version of Cmake for Ubuntu 22.04
+# Updating to latest version of Cmake for Ubuntu 24.04
 RUN apt-get update && apt-get install -y wget build-essential \
  && wget -qO- https://cmake.org/files/v3.28/cmake-3.28.3-linux-x86_64.tar.gz \
  | tar --strip-components=1 -xz -C /usr/local
 
 # Use bash shell for all RUN commands
 SHELL ["/bin/bash", "--login", "-c"]
-    
+
 # Copy Sources
 ENV DGL_SRC_DIR="/src/dgl"
 RUN mkdir -p ${DGL_SRC_DIR}
@@ -35,14 +38,16 @@ ENV GPU_BUILD_TARGETS=${ARG_GPU_BUILD_TARGETS}
 ENV DGL_BUILD_DIR="${DGL_SRC_DIR}/build"
 ARG ARG_DGL_ARTIFACTS_DIR="/artifacts"
 ENV DGL_ARTIFACTS_DIR="${ARG_DGL_ARTIFACTS_DIR}"
+ENV DGL_LIBRARY_PATH="${DGL_BUILD_DIR}"
+ENV PYTHONPATH="${DGL_SRC_DIR}/tests:${DGL_SRC_DIR}/python"
 
 # Install GraphBolt dependencies
 RUN mkdir -p /tmp/deps && cd /tmp/deps && bash ${DGL_SRC_DIR}/script/install_graphbolt_deps.sh
 
 # Configure and build DGL
 WORKDIR ${DGL_SRC_DIR}
-RUN cmake --preset rocm -DCMAKE_HIP_ARCHITECTURES=${ARG_GPU_BUILD_TARGETS} -DGPU_TARGETS=${ARG_GPU_BUILD_TARGETS}  
-RUN cmake --build ${DGL_BUILD_DIR} --parallel ${MAX_JOBS}  
+RUN cmake --preset rocm -DCMAKE_HIP_ARCHITECTURES=${ARG_GPU_BUILD_TARGETS} -DGPU_TARGETS=${ARG_GPU_BUILD_TARGETS}
+RUN cmake --build ${DGL_BUILD_DIR} --parallel ${MAX_JOBS}
 
 WORKDIR ${DGL_SRC_DIR}/python
 # Make sure we have wheel and build installed
@@ -66,7 +71,7 @@ Container manifest:
 /src/dgl
 ├── README.md
 ├── benchmarks # Benchmark scripts
-├── ... 
+├── ...
 ├── build      # C/C++ build directory
 ├── tests      # Test scripts
 

diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm
@@ -0,0 +1,83 @@
+# "Copyright Advanced Micro Devices, Inc.
+# Licensed under the Apache License Version 2.0"
+
+#############################################################################
+ARG BASE_IMAGE=rocm/pytorch:rocm7.0.2_ubuntu24.04_py3.12_pytorch_release_2.6.0
+FROM ${BASE_IMAGE} AS dgl_build
+
+# NOTE: This dockerfile **assumes** that BASE_IMAGE comes with the appropriate
+# python executable in the path.
+
+ARG ARG_MAX_JOBS=24
+ENV MAX_JOBS=${ARG_MAX_JOBS:-24}
+
+# Install basic tools
+RUN apt-get update \
+    && apt install -y clang-format-15 git ninja-build patch tree
+
+# Updating to latest version of Cmake for Ubuntu 24.04
+RUN apt-get update && apt-get install -y wget build-essential \
+ && wget -qO- https://cmake.org/files/v3.28/cmake-3.28.3-linux-x86_64.tar.gz \
+ | tar --strip-components=1 -xz -C /usr/local
+
+# Use bash shell for all RUN commands
+SHELL ["/bin/bash", "--login", "-c"]
+
+# Copy Sources
+ENV DGL_SRC_DIR="/src/dgl"
+RUN mkdir -p ${DGL_SRC_DIR}
+COPY . ${DGL_SRC_DIR}
+
+# Set GPU build targets
+ARG ARG_GPU_BUILD_TARGETS="gfx90a,gfx942"
+ENV GPU_BUILD_TARGETS=${ARG_GPU_BUILD_TARGETS}
+
+ENV DGL_BUILD_DIR="${DGL_SRC_DIR}/build"
+ARG ARG_DGL_ARTIFACTS_DIR="/artifacts"
+ENV DGL_ARTIFACTS_DIR="${ARG_DGL_ARTIFACTS_DIR}"
+ENV DGL_LIBRARY_PATH="${DGL_BUILD_DIR}"
+ENV PYTHONPATH="${DGL_SRC_DIR}/tests:${DGL_SRC_DIR}/python"
+
+# Install GraphBolt dependencies
+RUN mkdir -p /tmp/deps && cd /tmp/deps && bash ${DGL_SRC_DIR}/script/install_graphbolt_deps.sh
+
+# Configure and build DGL
+WORKDIR ${DGL_SRC_DIR}
+RUN cmake --preset rocm -DCMAKE_HIP_ARCHITECTURES=${ARG_GPU_BUILD_TARGETS} -DGPU_TARGETS=${ARG_GPU_BUILD_TARGETS}
+RUN cmake --build ${DGL_BUILD_DIR} --parallel ${MAX_JOBS}
+
+WORKDIR ${DGL_SRC_DIR}/python
+# Make sure we have wheel and build installed
+RUN python -m pip install build wheel
+# Install an editable version of the package so we can make source changes and test them
+RUN python -m pip install -e .
+# Build the wheel file for packaging (won't update automatically with source changes)
+RUN python -m build --wheel
+
+
+# Store the artifacts
+RUN mkdir -p ${DGL_ARTIFACTS_DIR}
+RUN cp dist/*.whl ${DGL_ARTIFACTS_DIR}
+
+# Setup for the user
+RUN cat <<'EOT' >> ~/.bashrc
+echo "
+Welcome to the DGL CI container!
+
+Container manifest:
+/src/dgl
+├── README.md
+├── benchmarks # Benchmark scripts
+├── ...
+├── build      # C/C++ build directory
+├── tests      # Test scripts
+
+Run the tests with:
+    cd /src/dgl/
+    bash tests/scripts/task_unit_test_rocm.sh pytorch gpu
+"
+EOT
+
+# Set the default command to an interactive bash shell.
+WORKDIR ${DGL_SRC_DIR}
+CMD ["/bin/bash"]
diff --git a/graphbolt/include/graphbolt/async.h b/graphbolt/include/graphbolt/async.h
@@ -43,10 +43,11 @@
 #include <c10/cuda/CUDAStream.h>
 #include <torch/csrc/api/include/torch/cuda.h>
 #elif defined(GRAPHBOLT_USE_HIP)
-#include <dgl/hip/cuda_to_hip.h>
+#include <ATen/hip/HIPContext.h>
 #include <ATen/hip/HIPEvent.h>
 #include <ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h>
 #include <ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h>
+#include <c10/hip/HIPGuard.h>
 #include <torch/csrc/api/include/torch/cuda.h>
 
 namespace c10::cuda {

diff --git a/graphbolt/src/cuda/common.h b/graphbolt/src/cuda/common.h
@@ -26,17 +26,21 @@
 // checks if this is defined rather than checking the value.
 #undef CUDART_VERSION
 
+#include <ATen/hip/HIPContext.h>
 #include <ATen/hip/HIPEvent.h>
 #include <ATen/hip/impl/HIPCachingAllocatorMasqueradingAsCUDA.h>
 #include <ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h>
 #include <c10/hip/HIPException.h>
-#include <dgl/hip/cuda_to_hip.h>
 #include <hip/hip_runtime.h>
+#include <dgl/hip/cuda_to_hip.h>
 
 using namespace c10::hip;
 using GPUStream_t = at::hip::HIPStreamMasqueradingAsCUDA;
 #define THRUST_BACKEND thrust::hip
 #define GET_CURRENT_GPU_STREAM getCurrentHIPStreamMasqueradingAsCUDA
+#define GRAPHBOLT_CACHING_ALLOCATOR c10::hip::HIPCachingAllocator
+#define GRAPHBOLT_C10_CHECK C10_HIP_CHECK
+#define GRAPHBOLT_C10_KERNEL_LAUNCH_CHECK C10_HIP_KERNEL_LAUNCH_CHECK
 
 #else // we're using CUDA
 
@@ -50,6 +54,9 @@
 using GPUStream_t = at::cuda::CUDAStream;
 #define THRUST_BACKEND thrust::cuda
 #define GET_CURRENT_GPU_STREAM getCurrentCUDAStream
+#define GRAPHBOLT_CACHING_ALLOCATOR c10::cuda::CUDACachingAllocator
+#define GRAPHBOLT_C10_CHECK C10_CUDA_CHECK
+#define GRAPHBOLT_C10_KERNEL_LAUNCH_CHECK C10_CUDA_KERNEL_LAUNCH_CHECK
 
 #endif // DGL_USE_HIP
 
@@ -97,13 +104,13 @@
   CUDAWorkspaceAllocator& operator=(const CUDAWorkspaceAllocator&) = default;
 
   void operator()(void* ptr) const {
-   CUDACachingAllocator::raw_delete(ptr);
+   GRAPHBOLT_CACHING_ALLOCATOR::raw_delete(ptr);
   }
 
   // Required by thrust to satisfy allocator requirements.
   value_type* allocate(std::ptrdiff_t size) const {
     return reinterpret_cast<value_type*>(
-       CUDACachingAllocator::raw_alloc(size * sizeof(value_type)));
+       GRAPHBOLT_CACHING_ALLOCATOR::raw_alloc(size * sizeof(value_type)));
   }
 
   // Required by thrust to satisfy allocator requirements.
@@ -114,7 +121,7 @@
       std::size_t size) const {
     return std::unique_ptr<T, CUDAWorkspaceAllocator>(
         reinterpret_cast<T*>(
-           CUDACachingAllocator::raw_alloc(sizeof(T) * size)),
+           GRAPHBOLT_CACHING_ALLOCATOR::raw_alloc(sizeof(T) * size)),
         *this);
   }
 };
@@ -142,15 +149,15 @@
     }                                                      \
   } while (0)
 
-#define CUDA_CALL(func) C10_CUDA_CHECK((func))
+#define CUDA_CALL(func) GRAPHBOLT_C10_CHECK((func))
 
 #define CUDA_KERNEL_CALL(kernel, nblks, nthrs, shmem, ...)          \
   {                                                                 \
     if (!graphbolt::cuda::is_zero((nblks)) &&                       \
         !graphbolt::cuda::is_zero((nthrs))) {                       \
       auto stream = graphbolt::cuda::GetCurrentStream();            \
       (kernel)<<<(nblks), (nthrs), (shmem), stream>>>(__VA_ARGS__); \
-      C10_CUDA_KERNEL_LAUNCH_CHECK();                               \
+      GRAPHBOLT_C10_KERNEL_LAUNCH_CHECK();                          \
     }                                                               \
   }
 

diff --git a/graphbolt/src/cuda/extension/gpu_graph_cache.cu b/graphbolt/src/cuda/extension/gpu_graph_cache.cu
@@ -19,14 +19,16 @@
  */
 #include <graphbolt/cuda_ops.h>
 #include <thrust/gather.h>
+#include <thrust/iterator/discard_iterator.h>
 #include <thrust/transform.h>
 
 #include <cstddef>
 #ifdef GRAPHBOLT_USE_HIP
 #include <hipcub/hipcub.hpp>
-#define C10_CUDA_KERNEL_LAUNCH_CHECK C10_HIP_KERNEL_LAUNCH_CHECK
+#define GRAPHBOLT_DISCARD_OUTPUT_ITERATOR thrust::discard_iterator<>{}
 #else
 #include <cub/cub.cuh>
+#define GRAPHBOLT_DISCARD_OUTPUT_ITERATOR cub::DiscardOutputIterator{}
 #endif
 #include <cuco/static_map.cuh>
 #include <cuda/std/atomic>
@@ -147,7 +149,7 @@
             ::cuda::stream_ref{cuda::GetCurrentStream()}};
         map_ = new map_t<index_t>{std::move(map_temp)};
       }));
-  C10_CUDA_KERNEL_LAUNCH_CHECK();  // Check the map constructor's success.
+  GRAPHBOLT_C10_KERNEL_LAUNCH_CHECK();  // Check the map constructor's success.
   const auto options = torch::TensorOptions().device(c10::DeviceType::CUDA);
   TORCH_CHECK(threshold > 0, "Threshold should be a position integer.");
   threshold_ = threshold;
@@ -235,7 +237,7 @@
             output_indices.data_ptr<index_t>());
         CUB_CALL(
             DevicePartition::If, position_and_index, output_position_and_index,
-            cub::DiscardOutputIterator{}, seeds.size(0),
+            GRAPHBOLT_DISCARD_OUTPUT_ITERATOR, seeds.size(0),
             [] __device__(thrust::tuple<index_t, index_t> & x) {
               return thrust::get<0>(x) >= 0;
             });
@@ -407,7 +409,8 @@
                 CUB_CALL(
                     DeviceSelect::Flagged, iota, is_threshold,
                     output_indices.data_ptr<indices_t>(),
-                    cub::DiscardOutputIterator{}, missing_positions.size(0));
+                    GRAPHBOLT_DISCARD_OUTPUT_ITERATOR,
+                    missing_positions.size(0));
                 auto [in_degree, sliced_indptr] =
                     ops::SliceCSCIndptr(indptr, output_indices);
                 while (num_nodes_ + num_threshold >= indptr_.size(0)) {

diff --git a/graphbolt/src/cuda/extension/unique_and_compact_map.cu b/graphbolt/src/cuda/extension/unique_and_compact_map.cu
@@ -25,7 +25,6 @@
 
 #ifdef GRAPHBOLT_USE_HIP
 #include <hipcub/hipcub.hpp>
-#define C10_CUDA_KERNEL_LAUNCH_CHECK C10_HIP_KERNEL_LAUNCH_CHECK
 #else
 #include <cub/cub.cuh>
 #include <cuda/functional>
@@ -193,8 +192,8 @@
            {},
             cuda::CUDAWorkspaceAllocator<cuco::pair<int64_t, int64_t>>{},
             ::cuda::stream_ref{stream},
         };
-        C10_CUDA_KERNEL_LAUNCH_CHECK();  // Check the map constructor's success.
+        GRAPHBOLT_C10_KERNEL_LAUNCH_CHECK();  // Check the map constructor's success.
         const dim3 block(BLOCK_SIZE);
         const dim3 grid(
             (offsets_ptr[2 * num_batches] + BLOCK_SIZE - 1) / BLOCK_SIZE);

diff --git a/include/dgl/hip/cuda_to_hip.h b/include/dgl/hip/cuda_to_hip.h
@@ -222,10 +222,8 @@ using cudaStream_t = hipStream_t;
 #define cusparseXcsrsort hipsparseXcsrsort
 #define cusparseXcsrsort_bufferSizeExt hipsparseXcsrsort_bufferSizeExt
 
-// Used in graphbolt
-#define CUDACachingAllocator HIPCachingAllocator
-#define getCurrentCUDAStream getCurrentHIPStream
-#define C10_CUDA_CHECK C10_HIP_CHECK
-#define C10_CUDA_KERNEL_LAUNCH_CHECK C10_HIP_KERNEL_LAUNCH_CHECK
+// Torch's hipified c10/ATen headers already declare the CUDA-named stream,
+// event, allocator and error-check APIs inside c10::cuda / at::cuda, so
+// redirecting them here would alias them onto HIP names that do not exist.
 
 #endif
diff --git a/python/dgl/graphbolt/internal/utils.py b/python/dgl/graphbolt/internal/utils.py
@@ -164,7 +164,10 @@ def read_edges(dataset_dir, edge_fmt, edge_path):
             os.path.join(dataset_dir, edge_path),
             names=["src", "dst"],
         )
-        src, dst = edge_data["src"].to_numpy(), edge_data["dst"].to_numpy()
+        src, dst = (
+            edge_data["src"].to_numpy(copy=True),
+            edge_data["dst"].to_numpy(copy=True),
+        )
     return (src, dst)
 
 

diff --git a/script/install_graphbolt_deps.sh b/script/install_graphbolt_deps.sh
@@ -83,3 +83,34 @@ cd ${DEPS_DIR}
 # Right now we need to patch the rocPRIM headers to fix the build because these
 # config headers are missing gfx942 (I've added them manually)
 run cp ${FILE_SOURCE_DIR}/*.hpp ${INSTALL_PREFIX}/include/rocprim/device/detail/config/.
+
+# hipCollections installs cuco headers that still need local fixes on recent ROCm
+# stacks (rocThrust 5.0, libhipcxx SFINAE). Patch files live under script/patches/.
+apply_cuco_patches() {
+  local patch_dir="${FILE_SOURCE_DIR}/patches"
+  if [[ ! -d "${patch_dir}" ]]; then
+    return 0
+  fi
+
+  shopt -s nullglob
+  local patches=("${patch_dir}"/*.patch)
+  shopt -u nullglob
+  if [[ ${#patches[@]} -eq 0 ]]; then
+    return 0
+  fi
+
+  for patch_file in "${patches[@]}"; do
+    echo "Applying ${patch_file} to cuco headers under ${INSTALL_PREFIX}/include"
+    if $DRY_RUN; then
+      echo "[dry-run] patch -p1 -d ${INSTALL_PREFIX} < ${patch_file}"
+      continue
+    fi
+
+    if ! sed -n '/^--- /,$p' "${patch_file}" | patch -p1 -d "${INSTALL_PREFIX}" --forward --batch; then
+      echo "Failed to apply ${patch_file}" >&2
+      exit 1
+    fi
+  done
+}
+
+apply_cuco_patches