Fix builds for TensorFlow Lite and TensorRT

bytedeco · Apr 29, 2024 · becfa06 · becfa06
1 parent f53a152
commit becfa06
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 9 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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 {