1# SPDX-FileCopyrightText: Copyright 2010-2024 Arm Limited and/or its affiliates <open-source-office@arm.com>
2#
3# SPDX-License-Identifier: Apache-2.0
4#
5# Licensed under the Apache License, Version 2.0 (the License); you may
6# not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an AS IS BASIS, WITHOUT
13# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17from test_settings import TestSettings
18
19import numpy as np
20import tensorflow as tf
21import tf_keras as keras
22
23class PoolingSettings(TestSettings):
24
25    def __init__(self,
26                 dataset,
27                 testtype,
28                 regenerate_weights,
29                 regenerate_input,
30                 regenerate_biases,
31                 schema_file,
32                 channels=8,
33                 x_in=4,
34                 y_in=4,
35                 w_x=4,
36                 w_y=4,
37                 stride_x=1,
38                 stride_y=1,
39                 randmin=TestSettings.INT8_MIN,
40                 randmax=TestSettings.INT8_MAX,
41                 bias_min=TestSettings.INT32_MIN,
42                 bias_max=TestSettings.INT32_MAX,
43                 batches=1,
44                 pad=False,
45                 relu6=False,
46                 out_activation_min=None,
47                 out_activation_max=None,
48                 int16xint8=False,
49                 interpreter="tensorflow"):
50        super().__init__(dataset,
51                         testtype,
52                         regenerate_weights,
53                         regenerate_input,
54                         regenerate_biases,
55                         schema_file,
56                         channels,
57                         channels,
58                         x_in,
59                         y_in,
60                         w_x,
61                         w_y,
62                         stride_x,
63                         stride_y,
64                         pad,
65                         randmin=randmin,
66                         randmax=randmax,
67                         batches=batches,
68                         relu6=relu6,
69                         out_activation_min=out_activation_min,
70                         out_activation_max=out_activation_max,
71                         int16xint8=int16xint8,
72                         interpreter=interpreter)
73
74    def generate_data(self, input_data=None) -> None:
75        if self.is_int16xint8:
76            datatype = "int16_t"
77            inttype = tf.int16
78        else:
79            datatype = "int8_t"
80            inttype = tf.int8
81
82        input_data = self.get_randomized_input_data(input_data)
83        self.generate_c_array(self.input_data_file_prefix, input_data, datatype=datatype)
84
85        input_data = tf.cast(input_data, tf.float32)
86
87        # Create a one-layer Keras model
88        model = keras.models.Sequential()
89        input_shape = (self.batches, self.y_input, self.x_input, self.input_ch)
90        model.add(keras.layers.InputLayer(input_shape=input_shape[1:], batch_size=self.batches))
91        if self.test_type == 'avgpool':
92            model.add(
93                keras.layers.AveragePooling2D(pool_size=(self.filter_y, self.filter_x),
94                                                 strides=(self.stride_y, self.stride_x),
95                                                 padding=self.padding,
96                                                 input_shape=input_shape[1:]))
97        elif self.test_type == 'maxpool':
98            model.add(
99                keras.layers.MaxPooling2D(pool_size=(self.filter_y, self.filter_x),
100                                             strides=(self.stride_y, self.stride_x),
101                                             padding=self.padding,
102                                             input_shape=input_shape[1:]))
103        else:
104            raise RuntimeError("Wrong test type")
105
106        interpreter = self.convert_and_interpret(model, inttype, input_data)
107
108        output_details = interpreter.get_output_details()
109
110        self.x_output = output_details[0]['shape'][2]
111        self.y_output = output_details[0]['shape'][1]
112
113        self.calculate_padding(self.x_output, self.y_output, self.x_input, self.y_input)
114
115        # Generate reference
116        interpreter.invoke()
117        output_data = interpreter.get_tensor(output_details[0]["index"])
118        self.generate_c_array(self.output_data_file_prefix,
119                              np.clip(output_data, self.out_activation_min, self.out_activation_max),
120                              datatype=datatype)
121
122        self.write_c_config_header()
123        self.write_c_header_wrapper()
124
125    def write_c_config_header(self) -> None:
126        super().write_c_config_header()
127
128        filename = self.config_data
129        filepath = self.headers_dir + filename
130        prefix = self.testdataset.upper()
131
132        with open(filepath, "a") as f:
133            self.write_common_config(f, prefix)
134