From 390403d95788b3b8f5a68a665e1686a9ff0cf305 Mon Sep 17 00:00:00 2001
From: Nate Meyer <Nate.Devel@gmail.com>
Date: Tue, 5 Mar 2024 07:04:13 -0500
Subject: [PATCH] Pass computes value to libyolo_layer build (#10210)

---
 docker/tensorrt/Dockerfile.base              | 2 ++
 docker/tensorrt/detector/tensorrt_libyolo.sh | 2 +-
 docker/tensorrt/trt.hcl                      | 6 +++++-
 docker/tensorrt/trt.mk                       | 2 +-
 4 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/docker/tensorrt/Dockerfile.base b/docker/tensorrt/Dockerfile.base
index b0015016d..59ead46f5 100644
--- a/docker/tensorrt/Dockerfile.base
+++ b/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/*
diff --git a/docker/tensorrt/detector/tensorrt_libyolo.sh b/docker/tensorrt/detector/tensorrt_libyolo.sh
index 91b9340a9..46e4077fa 100755
--- a/docker/tensorrt/detector/tensorrt_libyolo.sh
+++ b/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
diff --git a/docker/tensorrt/trt.hcl b/docker/tensorrt/trt.hcl
index 56e294100..3195fb5bf 100644
--- a/docker/tensorrt/trt.hcl
+++ b/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}"]
 }
diff --git a/docker/tensorrt/trt.mk b/docker/tensorrt/trt.mk
index 0e01c1402..ad2425c81 100644
--- a/docker/tensorrt/trt.mk
+++ b/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)