From fbb4cc7379733163f0d0140d11069f8a21a99ae6 Mon Sep 17 00:00:00 2001
From: William Tambellini <wtambellini@sdl.com>
Date: Thu, 10 Jul 2025 19:59:44 -0700
Subject: [PATCH] =?UTF-8?q?[TRTLLM-4770][feat]=20Enhance=20cpp=20executor?=
 =?UTF-8?q?=20cmake=20to=20listen=20to=20ENABLE=5FMU=E2=80=A6=20(#5104)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

...LTI_DEVICE

Signed-off-by: William Tambellini <wtambellini@sdl.com>
---
 examples/cpp/executor/CMakeLists.txt | 54 ++++++++++++++++++----------
 scripts/build_cpp_examples.py        |  6 +++-
 2 files changed, 40 insertions(+), 20 deletions(-)

diff --git a/examples/cpp/executor/CMakeLists.txt b/examples/cpp/executor/CMakeLists.txt
index 120d756ed6..af029c272a 100644
--- a/examples/cpp/executor/CMakeLists.txt
+++ b/examples/cpp/executor/CMakeLists.txt
@@ -22,11 +22,20 @@ if(NOT TRTLLM_BUILD_DIR)
   set(TRTLLM_BUILD_DIR "${TRTLLM_DIR}/cpp/build")
 endif()
 set(TRTLLM_LIB_PATH "${TRTLLM_BUILD_DIR}/tensorrt_llm/libtensorrt_llm.so")
+if(NOT EXISTS ${TRTLLM_LIB_PATH})
+  message(FATAL_ERROR "Cannot find ${TRTLLM_LIB_PATH}")
+endif()
+
 set(TRTLLM_PLUGIN_PATH
     "${TRTLLM_BUILD_DIR}/tensorrt_llm/plugins/libnvinfer_plugin_tensorrt_llm.so"
 )
 set(TRTLLM_INCLUDE_DIR "${TRTLLM_DIR}/cpp/include")
 
+option(
+  ENABLE_MULTI_DEVICE
+  "Enable multi device/instance examples building (requires MPI headers/libs)"
+  ON)
+
 # Determine CXX11 ABI compatibility
 execute_process(
   COMMAND bash -c "nm -f posix -D ${TRTLLM_LIB_PATH} | grep __cxx11"
@@ -47,8 +56,13 @@ set(CMAKE_VERBOSE_MAKEFILE 1)
 # Define project name
 project(executorExamples)
 
-# Compile options
-set(CMAKE_CXX_FLAGS "-Wall -pthread -lstdc++ -DENABLE_MULTI_DEVICE=1")
+# Compile options $<BOOL:${ENABLE_MULTI_DEVICE}> ?
+if(ENABLE_MULTI_DEVICE)
+  set(EMD 1)
+else()
+  set(EMD 0)
+endif()
+set(CMAKE_CXX_FLAGS "-Wall -pthread -lstdc++ -DENABLE_MULTI_DEVICE=${EMD} ")
 set(CMAKE_CXX_FLAGS_RELEASE "-O3")
 set(CMAKE_BUILD_TYPE release)
 
@@ -120,23 +134,25 @@ target_link_libraries(executorExampleAdvanced nvinfer_plugin_tensorrt_llm
                       cxxopts::cxxopts)
 
 # MultiInstance
+if(ENABLE_MULTI_DEVICE)
+  find_package(MPI REQUIRED)
+  message(STATUS "Using MPI_C_INCLUDE_DIRS: ${MPI_C_INCLUDE_DIRS}")
+  message(STATUS "Using MPI_C_LIBRARIES: ${MPI_C_LIBRARIES}")
+  include_directories(${MPI_C_INCLUDE_DIRS})
 
-find_package(MPI REQUIRED)
-message(STATUS "Using MPI_C_INCLUDE_DIRS: ${MPI_C_INCLUDE_DIRS}")
-message(STATUS "Using MPI_C_LIBRARIES: ${MPI_C_LIBRARIES}")
-include_directories(${MPI_C_INCLUDE_DIRS})
+  add_executable(executorExampleAdvancedMultiInstances
+                 executorExampleAdvancedMultiInstances.cpp)
+  target_link_libraries(
+    executorExampleAdvancedMultiInstances nvinfer_plugin_tensorrt_llm
+    cxxopts::cxxopts ${MPI_C_LIBRARIES})
 
-add_executable(executorExampleAdvancedMultiInstances
-               executorExampleAdvancedMultiInstances.cpp)
-target_link_libraries(
-  executorExampleAdvancedMultiInstances nvinfer_plugin_tensorrt_llm
-  cxxopts::cxxopts ${MPI_C_LIBRARIES})
+  # FastLogits
+  add_executable(executorExampleFastLogits executorExampleFastLogits.cpp)
+  target_link_libraries(executorExampleFastLogits nvinfer_plugin_tensorrt_llm
+                        cxxopts::cxxopts ${MPI_C_LIBRARIES})
 
-# FastLogits
-add_executable(executorExampleFastLogits executorExampleFastLogits.cpp)
-target_link_libraries(executorExampleFastLogits nvinfer_plugin_tensorrt_llm
-                      cxxopts::cxxopts ${MPI_C_LIBRARIES})
-
-add_executable(executorExampleDisaggregated executorExampleDisaggregated.cpp)
-target_link_libraries(executorExampleDisaggregated nvinfer_plugin_tensorrt_llm
-                      cxxopts::cxxopts ${MPI_C_LIBRARIES})
+  add_executable(executorExampleDisaggregated executorExampleDisaggregated.cpp)
+  target_link_libraries(
+    executorExampleDisaggregated nvinfer_plugin_tensorrt_llm cxxopts::cxxopts
+    ${MPI_C_LIBRARIES})
+endif()
diff --git a/scripts/build_cpp_examples.py b/scripts/build_cpp_examples.py
index f30d917bbe..cb2591acfe 100644
--- a/scripts/build_cpp_examples.py
+++ b/scripts/build_cpp_examples.py
@@ -21,7 +21,7 @@ def working_directory(path: PathLike):
 
 
 def build_cpp_examples(build_dir: PathLike, trt_dir: PathLike,
-                       loglevel: int) -> None:
+                       enable_multi_device: str, loglevel: int) -> None:
     logging.basicConfig(level=loglevel,
                         format='%(asctime)s - %(levelname)s - %(message)s')
     # Convert input paths to pathlib.Path objects
@@ -52,6 +52,7 @@ def build_cpp_examples(build_dir: PathLike, trt_dir: PathLike,
             '-B',
             '.',
             f'-DTensorRT_ROOT={cmake_parse(trt_dir)}',
+            f'-DENABLE_MULTI_DEVICE={enable_multi_device}',
         ] + generator
         logging.info(f"Executing {generate_command}")
         subprocess.run(generate_command, check=True)
@@ -70,6 +71,9 @@ if __name__ == '__main__':
     parser.add_argument('--trt-dir',
                         default='/usr/local/tensorrt',
                         help='TensorRT directory path')
+    parser.add_argument('--enable-multi-device',
+                        default='ON',
+                        help='Enable multi device support (requires MPI)')
     parser.add_argument('-v',
                         '--verbose',
                         help="verbose",