1 /*
2  * Copyright (C) 2017 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #ifndef ANDROID_FRAMEWORKS_ML_NN_COMMON_OPERATIONS_RNN_H
18 #define ANDROID_FRAMEWORKS_ML_NN_COMMON_OPERATIONS_RNN_H
19 
20 #include <vector>
21 
22 #include "ActivationFunctor.h"
23 #include "HalInterfaces.h"
24 
25 namespace android {
26 namespace nn {
27 
28 struct RunTimeOperandInfo;
29 struct Shape;
30 
31 class RNN {
32    public:
33     RNN(const hal::Operation& operation, RunTimeOperandInfo* operands);
34 
35     static bool Prepare(const hal::Operation& operation, RunTimeOperandInfo* operands,
36                         Shape* hiddenStateShape, Shape* outputShape);
37     bool Eval();
38 
39     static constexpr int kInputTensor = 0;
40     static constexpr int kWeightsTensor = 1;
41     static constexpr int kRecurrentWeightsTensor = 2;
42     static constexpr int kBiasTensor = 3;
43     static constexpr int kHiddenStateInTensor = 4;
44     static constexpr int kActivationParam = 5;
45 
46     static constexpr int kHiddenStateOutTensor = 0;
47     static constexpr int kOutputTensor = 1;
48 
49     template <typename T>
50     static bool RNNStep(const T* inputData, const Shape& inputShape, const T* hiddenStateInputData,
51                         const T* biasData, const T* weightsData, const Shape& weightsShape,
52                         const T* recurrentWeightsData, const Shape& recurrentWeightsShape,
53                         int32_t activation, T* outputData);
54 
55     template <typename T>
56     static bool RNNStep(const T* inputData, const Shape& inputShape, const T* auxInputData,
57                         const Shape& auxInputShape, const T* hiddenStateInputData,
58                         const T* biasData, const T* weightsData, const Shape& weightsShape,
59                         const T* auxWeightsData, const Shape& auxWeightsShape,
60                         const T* recurrentWeightsData, const Shape& recurrentWeightsShape,
61                         int32_t activation, uint32_t outputBatchStride, uint32_t outputBatchStep,
62                         T* outputData, T* hiddenStateOutput = nullptr);
63 
64    private:
65     ActivationFn activation_;
66 
67     const RunTimeOperandInfo* input_;
68     const RunTimeOperandInfo* weights_;
69     const RunTimeOperandInfo* recurrent_weights_;
70     const RunTimeOperandInfo* bias_;
71     const RunTimeOperandInfo* hidden_state_in_;
72 
73     RunTimeOperandInfo* hidden_state_out_;
74     RunTimeOperandInfo* output_;
75 };
76 
77 }  // namespace nn
78 }  // namespace android
79 
80 #endif  // ANDROID_FRAMEWORKS_ML_NN_COMMON_OPERATIONS_RNN_H
81