From 820907ba902e2a4e95b5c2db3a7899632232b963 Mon Sep 17 00:00:00 2001
From: hongtaozhang <hongtaozhang@microsoft.com>
Date: Sat, 16 Nov 2024 00:22:13 -0800
Subject: [PATCH] Add nvbandwidth build.

---
 .gitmodules             |  3 +++
 third_party/Makefile    | 17 +++++++++++++++--
 third_party/nvbandwidth |  1 +
 3 files changed, 19 insertions(+), 2 deletions(-)
 create mode 160000 third_party/nvbandwidth

diff --git a/.gitmodules b/.gitmodules
index 339520d19..9be41b59b 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -30,3 +30,6 @@
 [submodule "third_party/Megatron/Megatron-DeepSpeed"]
 	path = third_party/Megatron/Megatron-DeepSpeed
 	url = https://github.com/microsoft/Megatron-DeepSpeed.git
+[submodule "third_party/nvbandwidth"]
+	path = third_party/nvbandwidth
+	url = https://github.com/NVIDIA/nvbandwidth.git
diff --git a/third_party/Makefile b/third_party/Makefile
index 63ca48f36..39333b1b3 100755
--- a/third_party/Makefile
+++ b/third_party/Makefile
@@ -16,12 +16,12 @@ ROCM_VER ?= $(shell hipconfig -R | grep -oP '\d+\.\d+\.\d+' || echo "0.0.0")
 
 NUM_MAKE_JOBS ?= $(shell nproc --ignore=2)
 
-.PHONY: all cuda_with_msccl cuda rocm common cuda_cutlass cuda_bandwidthTest cuda_nccl_tests cuda_perftest cuda_msccl rocm_perftest fio rocm_rccl_tests rocm_rocblas rocm_bandwidthTest gpcnet cuda_gpuburn cpu_stream cpu_hpl directx_amf_encoding_latency directx_amd rocm_hipblaslt megatron_lm megatron_deepspeed apex_rocm
+.PHONY: all cuda_with_msccl cuda rocm common cuda_cutlass cuda_bandwidthTest cuda_nccl_tests cuda_perftest cuda_msccl rocm_perftest fio rocm_rccl_tests rocm_rocblas rocm_bandwidthTest gpcnet cuda_gpuburn cpu_stream cpu_hpl directx_amf_encoding_latency directx_amd rocm_hipblaslt megatron_lm megatron_deepspeed apex_rocm nvbandwidth
 
 # Build targets.
 all: cuda rocm
 cuda_with_msccl: cuda cuda_msccl
-cuda: common cuda_cutlass cuda_bandwidthTest cuda_nccl_tests cuda_perftest gpcnet cuda_gpuburn megatron_lm megatron_deepspeed
+cuda: common cuda_cutlass cuda_bandwidthTest cuda_nccl_tests cuda_perftest gpcnet cuda_gpuburn megatron_lm megatron_deepspeed nvbandwidth
 rocm: common rocm_perftest rocm_rccl_tests rocm_rocblas rocm_bandwidthTest rocm_hipblaslt megatron_deepspeed apex_rocm
 cpu: common cpu_perftest cpu_stream
 common: fio
@@ -239,3 +239,16 @@ ifneq (,$(wildcard msccl/tests/msccl-tests-nccl/Makefile))
 	mkdir -p $(SB_MICRO_PATH)/bin/msccl-tests-nccl && \
 	cp -r -v ./msccl/tests/msccl-tests-nccl/build/* $(SB_MICRO_PATH)/bin/msccl-tests-nccl/
 endif
+
+# Build nvbandwidth.
+REQUIRED_CMAKE_VERSION := 3.2.0
+CMAKE_VERSION := $(shell cmake --version 2>/dev/null | head -n 1 | sed 's/[^0-9.]*\([0-9.]*\).*/\1/')
+nvbandwidth: sb_micro_path
+	@if [ "$(CMAKE_VERSION)" = "" ]; then \
+        echo "Skip nvbandwidth for CMake is not installed."; \
+    elif [ "$(shell printf '%s\n' "$(REQUIRED_CMAKE_VERSION)" "$(CMAKE_VERSION)" | sort -V | head -n 1)" = "$(REQUIRED_CMAKE_VERSION)" ]; then \
+        cd ./nvbandwidth && cmake . && make && cd ..; \
+        cp -v ./nvbandwidth/nvbandwidth $(SB_MICRO_PATH)/bin; \
+    else \
+        echo "Skip nvbandwidth for CMake version $(CMAKE_VERSION) is too old. Required version is $(REQUIRED_CMAKE_VERSION) or higher."; \
+    fi
diff --git a/third_party/nvbandwidth b/third_party/nvbandwidth
new file mode 160000
index 000000000..445d8aef7
--- /dev/null
+++ b/third_party/nvbandwidth
@@ -0,0 +1 @@
+Subproject commit 445d8aef742e8a48a69779a939996f9e8863df9d