From dc7c297ef13a016222afd77c5ac400b0a5960db0 Mon Sep 17 00:00:00 2001
From: Naushir Patuck <naush@raspberrypi.com>
Date: Mon, 9 Sep 2024 13:37:07 +0100
Subject: [PATCH] ipa: rpi: Handle the new CNN controls in the IPA

Add code to handle the new CNN vendor controls in the Raspberry Pi IPA.

The value of CnnInputTensorInfo is cached as it is the only stateful
input control.

All other controls are output controls, and the values are copied into
directly from the rpiMetadata object if present. The camera helpers
populate the rpiMetadata object if the sensor supports on-board CNN
processing, such as the IMX500.

Signed-off-by: Naushir Patuck <naush@raspberrypi.com>
---
 src/ipa/rpi/common/ipa_base.cpp | 52 ++++++++++++++++++++++++++++++++-
 src/ipa/rpi/common/ipa_base.h   |  2 ++
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/src/ipa/rpi/common/ipa_base.cpp b/src/ipa/rpi/common/ipa_base.cpp
index 463f6d384..e7ed560b9 100644
--- a/src/ipa/rpi/common/ipa_base.cpp
+++ b/src/ipa/rpi/common/ipa_base.cpp
@@ -75,6 +75,7 @@ const ControlInfoMap::Map ipaControls{
 	{ &controls::FrameDurationLimits, ControlInfo(INT64_C(33333), INT64_C(120000)) },
 	{ &controls::draft::NoiseReductionMode, ControlInfo(controls::draft::NoiseReductionModeValues) },
 	{ &controls::rpi::StatsOutputEnable, ControlInfo(false, true, false) },
+	{ &controls::rpi::CnnEnableInputTensor, ControlInfo(false, true, false) },
 };
 
 /* IPA controls handled conditionally, if the sensor is not mono */
@@ -106,7 +107,7 @@ namespace ipa::RPi {
 IpaBase::IpaBase()
 	: controller_(), frameLengths_(FrameLengthsQueueSize, 0s), statsMetadataOutput_(false),
 	  stitchSwapBuffers_(false), frameCount_(0), mistrustCount_(0), lastRunTimestamp_(0),
-	  firstStart_(true), flickerState_({ 0, 0s })
+	  firstStart_(true), flickerState_({ 0, 0s }), cnnEnableInputTensor_(false)
 {
 }
 
@@ -1249,6 +1250,10 @@ void IpaBase::applyControls(const ControlList &controls)
 			statsMetadataOutput_ = ctrl.second.get<bool>();
 			break;
 
+		case controls::rpi::CNN_ENABLE_INPUT_TENSOR:
+			cnnEnableInputTensor_ = ctrl.second.get<bool>();
+			break;
+
 		default:
 			LOG(IPARPI, Warning)
 				<< "Ctrl " << controls::controls.at(ctrl.first)->name()
@@ -1425,6 +1430,51 @@ void IpaBase::reportMetadata(unsigned int ipaContext)
 			libcameraMetadata_.set(controls::HdrChannel, controls::HdrChannelNone);
 	}
 
+	const std::shared_ptr<uint8_t[]> *inputTensor =
+		rpiMetadata.getLocked<std::shared_ptr<uint8_t[]>>("cnn.input_tensor");
+	if (cnnEnableInputTensor_ && inputTensor) {
+		unsigned int size = *rpiMetadata.getLocked<unsigned int>("cnn.input_tensor_size");
+		Span<const uint8_t> tensor{ inputTensor->get(), size };
+		libcameraMetadata_.set(controls::rpi::CnnInputTensor, tensor);
+		/* No need to keep these big buffers any more. */
+		rpiMetadata.eraseLocked("cnn.input_tensor");
+	}
+
+	const RPiController::CnnInputTensorInfo *inputTensorInfo =
+		rpiMetadata.getLocked<RPiController::CnnInputTensorInfo>("cnn.input_tensor_info");
+	if (inputTensorInfo) {
+		Span<const uint8_t> tensorInfo{ reinterpret_cast<const uint8_t *>(inputTensorInfo),
+						sizeof(*inputTensorInfo) };
+		libcameraMetadata_.set(controls::rpi::CnnInputTensorInfo, tensorInfo);
+	}
+
+	const std::shared_ptr<float[]> *outputTensor =
+		rpiMetadata.getLocked<std::shared_ptr<float[]>>("cnn.output_tensor");
+	if (outputTensor) {
+		unsigned int size = *rpiMetadata.getLocked<unsigned int>("cnn.output_tensor_size");
+		Span<const float> tensor{ reinterpret_cast<const float *>(outputTensor->get()),
+					  size };
+		libcameraMetadata_.set(controls::rpi::CnnOutputTensor, tensor);
+		/* No need to keep these big buffers any more. */
+		rpiMetadata.eraseLocked("cnn.output_tensor");
+	}
+
+	const RPiController::CnnOutputTensorInfo *outputTensorInfo =
+		rpiMetadata.getLocked<RPiController::CnnOutputTensorInfo>("cnn.output_tensor_info");
+	if (outputTensorInfo) {
+		Span<const uint8_t> tensorInfo{ reinterpret_cast<const uint8_t *>(outputTensorInfo),
+						sizeof(*outputTensorInfo) };
+		libcameraMetadata_.set(controls::rpi::CnnOutputTensorInfo, tensorInfo);
+	}
+
+	const RPiController::CnnKpiInfo *kpiInfo =
+		rpiMetadata.getLocked<RPiController::CnnKpiInfo>("cnn.kpi_info");
+	if (kpiInfo) {
+		libcameraMetadata_.set(controls::rpi::CnnKpiInfo,
+				       { static_cast<int32_t>(kpiInfo->dnnRuntime),
+					 static_cast<int32_t>(kpiInfo->dspRuntime) });
+	}
+
 	metadataReady.emit(libcameraMetadata_);
 }
 
diff --git a/src/ipa/rpi/common/ipa_base.h b/src/ipa/rpi/common/ipa_base.h
index 1a811beb3..a55ce7ca9 100644
--- a/src/ipa/rpi/common/ipa_base.h
+++ b/src/ipa/rpi/common/ipa_base.h
@@ -136,6 +136,8 @@ class IpaBase : public IPARPiInterface
 		int32_t mode;
 		utils::Duration manualPeriod;
 	} flickerState_;
+
+	bool cnnEnableInputTensor_;
 };
 
 } /* namespace ipa::RPi */