From 79d63db0c888a87a4d7d0a4898527f4fe584bdc4 Mon Sep 17 00:00:00 2001
From: David Gross <dgross@google.com>
Date: Thu, 7 Oct 2021 16:47:25 -0700
Subject: [PATCH] Add PACK operation to NNAPI feature level 6.

Bug: 200281183
Bug: 200280665

Test: m -j NeuralNetworksTest_static
Test: VtsHalNeuralnetworksTargetTest

Change-Id: Ic15d047b70c62437b4f0db6f2ca10127591ae07c
---
 .../neuralnetworks/OperationType.aidl         |  1 +
 .../neuralnetworks/OperationType.aidl         | 49 +++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/neuralnetworks/aidl/aidl_api/android.hardware.neuralnetworks/current/android/hardware/neuralnetworks/OperationType.aidl b/neuralnetworks/aidl/aidl_api/android.hardware.neuralnetworks/current/android/hardware/neuralnetworks/OperationType.aidl
index 425914347b..2eff11b146 100644
--- a/neuralnetworks/aidl/aidl_api/android.hardware.neuralnetworks/current/android/hardware/neuralnetworks/OperationType.aidl
+++ b/neuralnetworks/aidl/aidl_api/android.hardware.neuralnetworks/current/android/hardware/neuralnetworks/OperationType.aidl
@@ -137,4 +137,5 @@ enum OperationType {
   FILL = 100,
   RANK = 101,
   BATCH_MATMUL = 102,
+  PACK = 103,
 }
diff --git a/neuralnetworks/aidl/android/hardware/neuralnetworks/OperationType.aidl b/neuralnetworks/aidl/android/hardware/neuralnetworks/OperationType.aidl
index d9951d58ef..2ec91acf30 100644
--- a/neuralnetworks/aidl/android/hardware/neuralnetworks/OperationType.aidl
+++ b/neuralnetworks/aidl/android/hardware/neuralnetworks/OperationType.aidl
@@ -5273,4 +5273,53 @@ enum OperationType {
      *      c_o = r_y if adj_y else c_y
      */
     BATCH_MATMUL = 102,
+
+    /**
+     * Packs N input tensors (N >= 1) of rank R into one output tensor of rank R+1.
+     * The tensors are packed along a given axis.
+     *
+     * The input tensors must have identical {@link OperandType} and dimensions.
+     *
+     * For example, suppose there are N input tensors of shape (A, B, C).
+     * If axis is 0, the output tensor will have shape (N, A, B, C).
+     * If axis is 1, the output tensor will have shape (A, N, B, C).
+     *
+     * All dimensions through the axis dimension determine the output tile count;
+     * the remaining dimensions determine the tile shape.
+     *
+     * Return to the example of N input tensors of shape (A, B, C).
+     * If axis is 0, there are N tiles in the output, each of shape (A, B, C).
+     * If axis is 1, there are A*N tiles in the output, each of shape (B, C).
+     *
+     * The coordinates of a tile within the output tensor are (t[0],...,t[axis]).
+     * The coordinates of a tile within an input tensor are (t[0],...,t[axis-1]).
+     * (If axis is 0, an input tensor consists of a single tile.)
+     * If we index input tensors starting with 0 (rather than by operand number),
+     * then output_tile[t[0],...,t[axis]] = input_tile[t[axis]][t[0],...,t[axis-1]].
+     * That is, all output tile coordinates except for the axis coordinate select
+     * the corresponding location within some input tensor; and the axis coordinate
+     * selects the input tensor.
+     *
+     * Supported tensor {@link OperandType}:
+     * * {@link OperandType::TENSOR_FLOAT16}
+     * * {@link OperandType::TENSOR_FLOAT32}
+     * * {@link OperandType::TENSOR_QUANT8_ASYMM}
+     * * {@link OperandType::TENSOR_QUANT8_ASYMM_SIGNED}
+     * * {@link OperandType::TENSOR_INT32}
+     *
+     * Supported input tensor rank: from 1
+     *
+     * Inputs:
+     * * 0: A scalar of type {@link OperandType::INT32}, specifying
+     *      the axis along which to pack.  The valid range is [0, R+1).
+     * * 1 ~ N: Input tensors to be packed together.
+     *          For {@link OperandType::TENSOR_QUANT8_ASYMM} and
+     *          {@link OperandType::TENSOR_QUANT8_ASYMM_SIGNED} tensors,
+     *          the scales and zeroPoint must be the same for all input tensors,
+     *          and will be the same for the output tensor.
+     *
+     * Outputs:
+     * * 0: The packed tensor.
+     */
+    PACK = 103,
 }