From 6142d58c89ccdb83f71961a5465bd3b479b25d77 Mon Sep 17 00:00:00 2001
From: Dongxu Yang <dongxuy@nvidia.com>
Date: Wed, 30 Aug 2023 15:23:51 +0800
Subject: [PATCH] fix docs build and optimize host continuous register

---
 build.sh                                                  | 8 +++++---
 cpp/CMakeLists.txt                                        | 2 +-
 cpp/src/wholememory/memory_handle.cpp                     | 4 +++-
 .../pylibwholegraph/pylibwholegraph/torch/data_loader.py  | 1 -
 4 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/build.sh b/build.sh
index ddf0d46a3..535988f54 100755
--- a/build.sh
+++ b/build.sh
@@ -212,6 +212,8 @@ if hasArg clean; then
     find ${REPODIR}/python/pylibwholegraph -name "*.cpython*.so" -type f -delete
 
     # remove docs build
+    rm -rf ${REPODIR}/cpp/html
+    rm -rf ${REPODIR}/cpp/xml
     cd ${REPODIR}/docs/wholegraph
     make BUILDDIR=${DOCS_BUILD_DIR} clean
     rm -rf ${REPODIR}/docs/wholegraph/_xml
@@ -303,13 +305,13 @@ if hasArg docs; then
               ${CMAKE_GENERATOR_OPTION} \
               ${CMAKE_VERBOSE_OPTION}
     fi
-    cd ${LIBWHOLEGRAPH_BUILD_DIR}
+    cd ${REPODIR}/cpp
     cmake --build "${LIBWHOLEGRAPH_BUILD_DIR}" -j${PARALLEL_LEVEL} --target doxygen ${VERBOSE_FLAG}
     mkdir -p ${REPODIR}/docs/wholegraph/_html/doxygen_docs/libwholegraph/html
-    mv ${LIBWHOLEGRAPH_BUILD_DIR}/html/* ${REPODIR}/docs/wholegraph/_html/doxygen_docs/libwholegraph/html
+    mv ${REPODIR}/cpp/html/* ${REPODIR}/docs/wholegraph/_html/doxygen_docs/libwholegraph/html
     mkdir -p ${REPODIR}/docs/wholegraph/_xml
     # _xml is used for sphinx breathe project
-    mv ${LIBWHOLEGRAPH_BUILD_DIR}/xml/* "${REPODIR}/docs/wholegraph/_xml"
+    mv ${REPODIR}/cpp/xml/* "${REPODIR}/docs/wholegraph/_xml"
     cd ${REPODIR}/docs/wholegraph
     PYTHONPATH=${REPODIR}/python/pylibwholegraph:${PYTHONPATH} make BUILDDIR=${DOCS_BUILD_DIR} html
     mv ${REPODIR}/docs/wholegraph/_html/doxygen_docs ${REPODIR}/docs/wholegraph/${DOCS_BUILD_DIR}/html/
diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index a072f6d0b..2162665c2 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -292,7 +292,7 @@ rapids_export(
 find_package(Doxygen 1.8.11)
 if(Doxygen_FOUND)
         add_custom_command(OUTPUT WHOLEGRAPH_DOXYGEN
-                           WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}
+                           WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
                            COMMAND doxygen Doxyfile
                            VERBATIM)
 
diff --git a/cpp/src/wholememory/memory_handle.cpp b/cpp/src/wholememory/memory_handle.cpp
index 78909c08e..f20668612 100644
--- a/cpp/src/wholememory/memory_handle.cpp
+++ b/cpp/src/wholememory/memory_handle.cpp
@@ -497,7 +497,9 @@ class global_mapped_host_wholememory_impl : public wholememory_impl {
       nullptr, alloc_strategy_.total_alloc_size, PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd, 0);
     WHOLEMEMORY_CHECK(mmap_ptr != (void*)-1);
 #endif
-    memset(mmap_ptr, 0, alloc_strategy_.total_alloc_size);
+    memset(static_cast<char*>(mmap_ptr) + rank_partition_strategy_.local_mem_offset,
+           0,
+           rank_partition_strategy_.local_mem_size);
     WM_CUDA_CHECK_NO_THROW(
       cudaHostRegister(mmap_ptr, alloc_strategy_.total_alloc_size, cudaHostRegisterDefault));
 #ifndef USE_SYSTEMV_SHM
diff --git a/python/pylibwholegraph/pylibwholegraph/torch/data_loader.py b/python/pylibwholegraph/pylibwholegraph/torch/data_loader.py
index 5edc2aade..4499ca216 100644
--- a/python/pylibwholegraph/pylibwholegraph/torch/data_loader.py
+++ b/python/pylibwholegraph/pylibwholegraph/torch/data_loader.py
@@ -81,7 +81,6 @@ def get_train_dataloader(
         batch_size=batch_size,
         num_workers=num_workers,
         pin_memory=True,
-        prefetch_factor=8 if num_workers > 0 else None,
         persistent_workers=True if num_workers > 0 else None,
         sampler=train_sampler,
     )