Pass computes value to libyolo_layer build (#10210)

docker/tensorrt/Dockerfile.base

@@ -8,6 +8,8 @@ ARG TRT_BASE=nvcr.io/nvidia/tensorrt:23.03-py3
 # Build TensorRT-specific library
 FROM ${TRT_BASE} AS trt-deps
 
+ARG COMPUTE_LEVEL
+
 RUN apt-get update \
     && apt-get install -y git build-essential cuda-nvcc-* cuda-nvtx-* libnvinfer-dev libnvinfer-plugin-dev libnvparsers-dev libnvonnxparsers-dev \
     && rm -rf /var/lib/apt/lists/*

docker/tensorrt/detector/tensorrt_libyolo.sh

@@ -11,7 +11,7 @@ git clone --depth 1 https://github.com/NateMeyer/tensorrt_demos.git -b condition
 if [ ! -e /usr/local/cuda ]; then
     ln -s /usr/local/cuda-* /usr/local/cuda
 fi
-cd ./tensorrt_demos/plugins && make all -j$(nproc)
+cd ./tensorrt_demos/plugins && make all -j$(nproc) computes="${COMPUTE_LEVEL:-}"
 cp libyolo_layer.so /usr/local/lib/libyolo_layer.so
 
 # Store yolo scripts for later conversion

docker/tensorrt/trt.hcl

@@ -10,12 +10,16 @@ variable "SLIM_BASE" {
 variable "TRT_BASE" {
   default = null
 }
+variable "COMPUTE_LEVEL" {
+  default = ""
+}
 
 target "_build_args" {
   args = {
     BASE_IMAGE = BASE_IMAGE,
     SLIM_BASE = SLIM_BASE,
-    TRT_BASE = TRT_BASE
+    TRT_BASE = TRT_BASE,
+    COMPUTE_LEVEL = COMPUTE_LEVEL
   }
   platforms = ["linux/${ARCH}"]
 }

docker/tensorrt/trt.mk

@@ -2,7 +2,7 @@ BOARDS += trt
 
 JETPACK4_BASE ?= timongentzsch/l4t-ubuntu20-opencv:latest	# L4T 32.7.1 JetPack 4.6.1
 JETPACK5_BASE ?= nvcr.io/nvidia/l4t-tensorrt:r8.5.2-runtime	# L4T 35.3.1 JetPack 5.1.1
-X86_DGPU_ARGS := ARCH=amd64
+X86_DGPU_ARGS := ARCH=amd64 COMPUTE_LEVEL="50 60 70 80 90"
 JETPACK4_ARGS := ARCH=arm64 BASE_IMAGE=$(JETPACK4_BASE) SLIM_BASE=$(JETPACK4_BASE) TRT_BASE=$(JETPACK4_BASE)
 JETPACK5_ARGS := ARCH=arm64 BASE_IMAGE=$(JETPACK5_BASE) SLIM_BASE=$(JETPACK5_BASE) TRT_BASE=$(JETPACK5_BASE)