From becfa06a42e6e58b9371d77abf71a68b4d43c97c Mon Sep 17 00:00:00 2001
From: Samuel Audet <samuel.audet@gmail.com>
Date: Tue, 30 Apr 2024 00:03:53 +0900
Subject: [PATCH] Fix builds for TensorFlow Lite and TensorRT

---
 tensorflow-lite/cppbuild.sh                      |  4 ++--
 .../org/bytedeco/tensorrt/presets/nvinfer.java   |  2 ++
 .../tensorrt/presets/nvinfer_plugin.java         | 13 ++++++++++---
 .../bytedeco/tensorrt/presets/nvonnxparser.java  | 16 ++++++++++++----
 4 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/tensorflow-lite/cppbuild.sh b/tensorflow-lite/cppbuild.sh
index b03054b526f..f1800d7ebff 100755
--- a/tensorflow-lite/cppbuild.sh
+++ b/tensorflow-lite/cppbuild.sh
@@ -46,7 +46,7 @@ case $PLATFORM in
         export CMAKE_FLAGS="-DCMAKE_TOOLCHAIN_FILE=${PLATFORM_ROOT}/build/cmake/android.toolchain.cmake -DANDROID_ABI=x86 -DANDROID_NATIVE_API_LEVEL=24 $CMAKE_FLAGS"
         ;;
     android-x86_64)
-        export CMAKE_FLAGS="-DCMAKE_TOOLCHAIN_FILE=${PLATFORM_ROOT}/build/cmake/android.toolchain.cmake -DANDROID_ABI=x86_64 -DANDROID_NATIVE_API_LEVEL=24 $CMAKE_FLAGS"
+        export CMAKE_FLAGS="-DCMAKE_TOOLCHAIN_FILE=${PLATFORM_ROOT}/build/cmake/android.toolchain.cmake -DANDROID_ABI=x86_64 -DANDROID_NATIVE_API_LEVEL=24 -DXNNPACK_ENABLE_AVXVNNI=OFF $CMAKE_FLAGS"
         ;;
     linux-armhf)
         export CC="arm-linux-gnueabihf-gcc -funsafe-math-optimizations"
@@ -56,7 +56,7 @@ case $PLATFORM in
     linux-arm64)
         export CC="aarch64-linux-gnu-gcc -funsafe-math-optimizations"
         export CXX="aarch64-linux-gnu-g++ -funsafe-math-optimizations"
-        export CMAKE_FLAGS="-DCMAKE_SYSTEM_NAME=Linux -DCMAKE_SYSTEM_PROCESSOR=aarch64 -DTFLITE_ENABLE_XNNPACK=OFF $CMAKE_FLAGS"
+        export CMAKE_FLAGS="-DCMAKE_SYSTEM_NAME=Linux -DCMAKE_SYSTEM_PROCESSOR=aarch64 -DXNNPACK_ENABLE_ARM_I8MM=OFF $CMAKE_FLAGS"
         ;;
     linux-x86)
         export CC="gcc -m32"
diff --git a/tensorrt/src/main/java/org/bytedeco/tensorrt/presets/nvinfer.java b/tensorrt/src/main/java/org/bytedeco/tensorrt/presets/nvinfer.java
index 9043f0d247d..7d144a9facb 100644
--- a/tensorrt/src/main/java/org/bytedeco/tensorrt/presets/nvinfer.java
+++ b/tensorrt/src/main/java/org/bytedeco/tensorrt/presets/nvinfer.java
@@ -70,6 +70,8 @@
         ),
         @Platform(
             value = "windows-x86_64",
+            link = "nvinfer_10",
+            preload = "nvinfer_builder_resource_10",
             includepath = "C:/Program Files/NVIDIA GPU Computing Toolkit/TensorRT/include",
             linkpath = "C:/Program Files/NVIDIA GPU Computing Toolkit/TensorRT/lib/"
         )
diff --git a/tensorrt/src/main/java/org/bytedeco/tensorrt/presets/nvinfer_plugin.java b/tensorrt/src/main/java/org/bytedeco/tensorrt/presets/nvinfer_plugin.java
index f6dc6f08530..80284cead7c 100644
--- a/tensorrt/src/main/java/org/bytedeco/tensorrt/presets/nvinfer_plugin.java
+++ b/tensorrt/src/main/java/org/bytedeco/tensorrt/presets/nvinfer_plugin.java
@@ -34,9 +34,16 @@
  */
 @Properties(
     inherit = nvinfer.class,
-    value = @Platform(
-        include = {"NvInferPlugin.h", "NvInferPluginUtils.h"},
-        link = "nvinfer_plugin@.10"),
+    value = {
+        @Platform(
+            include = {"NvInferPlugin.h", "NvInferPluginUtils.h"},
+            link = "nvinfer_plugin@.10"
+        ),
+        @Platform(
+            value = "windows-x86_64",
+            link = "nvinfer_plugin_10"
+        )
+    },
     target = "org.bytedeco.tensorrt.nvinfer_plugin",
     global = "org.bytedeco.tensorrt.global.nvinfer_plugin")
 public class nvinfer_plugin implements InfoMapper {
diff --git a/tensorrt/src/main/java/org/bytedeco/tensorrt/presets/nvonnxparser.java b/tensorrt/src/main/java/org/bytedeco/tensorrt/presets/nvonnxparser.java
index aede025be26..cba6d45fa81 100644
--- a/tensorrt/src/main/java/org/bytedeco/tensorrt/presets/nvonnxparser.java
+++ b/tensorrt/src/main/java/org/bytedeco/tensorrt/presets/nvonnxparser.java
@@ -34,10 +34,18 @@
  */
 @Properties(
     inherit = nvinfer_plugin.class,
-    value = @Platform(
-        include = "NvOnnxParser.h",
-        link = "nvonnxparser@.10",
-        preload = "nvinfer_vc_plugin@.10"),
+    value = {
+        @Platform(
+            include = "NvOnnxParser.h",
+            link = "nvonnxparser@.10",
+            preload = "nvinfer_vc_plugin@.10"
+        ),
+        @Platform(
+            value = "windows-x86_64",
+            link = "nvonnxparser_10",
+            preload = "nvinfer_vc_plugin_10"
+        )
+    },
     target = "org.bytedeco.tensorrt.nvonnxparser",
     global = "org.bytedeco.tensorrt.global.nvonnxparser")
 public class nvonnxparser implements InfoMapper {