iGPU build refactor (#224) (#225)

Author: kthui

CMakeLists.txt

@@ -78,6 +78,10 @@ project(tritononnxruntimebackend LANGUAGES C CXX)
 #     version that is compatible with the specified version of ONNX
 #     Runtime.
 #
+#   - Optionally set TRITON_BUILD_TARGET_PLATFORM to either linux, windows or
+#     igpu. If not set, the current platform will be used. If building on
+#     Jetpack, always set to igpu to avoid misdetection.
+#
 #   - If you want to disable GPU usage, set TRITON_ENABLE_GPU=OFF.
 #    This will make builds with CUDA and TensorRT flags to fail.
 #
@@ -91,6 +95,7 @@ set(TRITON_BUILD_CONTAINER "" CACHE STRING "Triton container to use a base for b
 set(TRITON_BUILD_CONTAINER_VERSION "" CACHE STRING "Triton container version to target")
 set(TRITON_BUILD_ONNXRUNTIME_VERSION "" CACHE STRING "ONNXRuntime version to build")
 set(TRITON_BUILD_ONNXRUNTIME_OPENVINO_VERSION "" CACHE STRING "ONNXRuntime OpenVINO version to build")
+set(TRITON_BUILD_TARGET_PLATFORM "" CACHE STRING "Target platform for ONNXRuntime build")
 set(TRITON_BUILD_CUDA_VERSION "" CACHE STRING "Version of CUDA install")
 set(TRITON_BUILD_CUDA_HOME "" CACHE PATH "Path to CUDA install")
 set(TRITON_BUILD_CUDNN_HOME "" CACHE PATH "Path to CUDNN install")
@@ -313,6 +318,9 @@ endif() # TRITON_ENABLE_ONNXRUNTIME_OPENVINO
 #
 if(TRITON_ONNXRUNTIME_DOCKER_BUILD)
   set(_GEN_FLAGS "")
+  if(NOT ${TRITON_BUILD_TARGET_PLATFORM} STREQUAL "")
+    set(_GEN_FLAGS ${_GEN_FLAGS} "--target-platform=${TRITON_BUILD_TARGET_PLATFORM}")
+  endif() # TRITON_BUILD_TARGET_PLATFORM
   if(NOT ${TRITON_BUILD_CUDA_VERSION} STREQUAL "")
     set(_GEN_FLAGS ${_GEN_FLAGS} "--cuda-version=${TRITON_BUILD_CUDA_VERSION}")
   endif() # TRITON_BUILD_CUDA_VERSION

tools/gen_ort_dockerfile.py

@@ -210,6 +210,8 @@ def dockerfile_for_linux(output_file):
             ep_flags += ' --cuda_home "{}"'.format(FLAGS.cuda_home)
         if FLAGS.cudnn_home is not None:
             ep_flags += ' --cudnn_home "{}"'.format(FLAGS.cudnn_home)
+        elif target_platform() == "igpu":
+            ep_flags += ' --cudnn_home "/usr/lib/aarch64-linux-gnu"'
         if FLAGS.ort_tensorrt:
             ep_flags += " --use_tensorrt"
             if FLAGS.ort_version >= "1.12.1":
@@ -224,7 +226,11 @@ def dockerfile_for_linux(output_file):
     if FLAGS.ort_openvino is not None:
         ep_flags += " --use_openvino CPU_FP32"
 
-    cuda_archs = "60;61;70;75;80;86;90"
+    if target_platform() == "igpu":
+        ep_flags += " --skip_tests --cmake_extra_defines 'onnxruntime_BUILD_UNIT_TESTS=OFF'"
+        cuda_archs = "53;62;72;87"
+    else:
+        cuda_archs = "60;61;70;75;80;86;90"
 
     df += """
 WORKDIR /workspace/onnxruntime
@@ -264,14 +270,21 @@ def dockerfile_for_linux(output_file):
        /opt/onnxruntime/lib && \
     cp /workspace/build/${ONNXRUNTIME_BUILD_CONFIG}/libonnxruntime.so \
        /opt/onnxruntime/lib
-
+"""
+    if target_platform() == "igpu":
+        df += """
+RUN mkdir -p /opt/onnxruntime/bin
+"""
+    else:
+        df += """
 RUN mkdir -p /opt/onnxruntime/bin && \
     cp /workspace/build/${ONNXRUNTIME_BUILD_CONFIG}/onnxruntime_perf_test \
        /opt/onnxruntime/bin && \
     cp /workspace/build/${ONNXRUNTIME_BUILD_CONFIG}/onnx_test_runner \
        /opt/onnxruntime/bin && \
     (cd /opt/onnxruntime/bin && chmod a+x *)
 """
+
     if FLAGS.enable_gpu:
         df += """
 RUN cp /workspace/build/${ONNXRUNTIME_BUILD_CONFIG}/libonnxruntime_providers_cuda.so \
@@ -334,6 +347,13 @@ def dockerfile_for_linux(output_file):
     done
 
 # For testing copy ONNX custom op library and model
+"""
+    if target_platform() == "igpu":
+        df += """
+RUN mkdir -p /opt/onnxruntime/test
+"""
+    else:
+        df += """
 RUN mkdir -p /opt/onnxruntime/test && \
     cp /workspace/build/${ONNXRUNTIME_BUILD_CONFIG}/libcustom_op_library.so \
        /opt/onnxruntime/test && \
@@ -533,7 +553,7 @@ def preprocess_gpu_flags():
         "--target-platform",
         required=False,
         default=None,
-        help='Target for build, can be "ubuntu", "windows" or "jetpack". If not specified, build targets the current platform.',
+        help='Target for build, can be "linux", "windows" or "igpu". If not specified, build targets the current platform.',
     )
 
     parser.add_argument(