1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 #include "tensorflow/lite/delegates/gpu/common/tasks/conv_weights_converter_test_util.h"
17 
18 #include <memory>
19 #include <vector>
20 
21 #include "tensorflow/lite/delegates/gpu/common/status.h"
22 #include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
23 #include "tensorflow/lite/delegates/gpu/common/task/weights_conversion.h"
24 #include "tensorflow/lite/delegates/gpu/common/tasks/conv_weights_converter.h"
25 
26 namespace tflite {
27 namespace gpu {
28 namespace {
ConvolutionWeightsConverterTest(const Tensor<OHWI,DataType::FLOAT32> & weights,const WeightsDescription & weight_desc,TestExecutionEnvironment * env,const OperationDef & op_def)29 absl::Status ConvolutionWeightsConverterTest(
30     const Tensor<OHWI, DataType::FLOAT32>& weights,
31     const WeightsDescription& weight_desc, TestExecutionEnvironment* env,
32     const OperationDef& op_def) {
33   // reinterpreting weights as HWIO-BHWC tensor
34   TensorFloat32 src_tensor_as_hwio;
35   src_tensor_as_hwio.shape =
36       BHWC(weights.shape.h, weights.shape.w, weights.shape.i, weights.shape.o);
37   src_tensor_as_hwio.data.resize(src_tensor_as_hwio.shape.DimensionsProduct(),
38                                  2.0);
39   // reinterpreting weights as OHWI-BHWC tensor
40   TensorFloat32 src_tensor_as_ohwi;
41   src_tensor_as_ohwi.shape =
42       BHWC(weights.shape.o, weights.shape.h, weights.shape.w, weights.shape.i);
43   src_tensor_as_ohwi.data.resize(src_tensor_as_ohwi.shape.DimensionsProduct(),
44                                  2.0);
45   for (int o = 0; o < weights.shape.o; ++o) {
46     for (int y = 0; y < weights.shape.h; ++y) {
47       for (int x = 0; x < weights.shape.w; ++x) {
48         for (int i = 0; i < weights.shape.i; ++i) {
49           const int f_index = weights.shape.LinearIndex({o, y, x, i});
50           const int s_index_hwio =
51               src_tensor_as_hwio.shape.LinearIndex({y, x, i, o});
52           src_tensor_as_hwio.data[s_index_hwio] = weights.data[f_index];
53           const int s_index_ohwi =
54               src_tensor_as_ohwi.shape.LinearIndex({o, y, x, i});
55           src_tensor_as_ohwi.data[s_index_ohwi] = weights.data[f_index];
56         }
57       }
58     }
59   }
60 
61   WeightsDescription weight_desc_copy = weight_desc;
62   weight_desc_copy.type = DataType::FLOAT32;
63   const int flt_count =
64       GetTotalElementsCountForLayout(weight_desc_copy, weights.shape);
65   DataType weights_type = DataType::FLOAT32;
66 
67   std::vector<uint8_t> weights_data(flt_count * SizeOf(weights_type));
68   RearrangeWeights(weights, weight_desc_copy, absl::MakeSpan(weights_data));
69 
70   std::vector<TensorFloat32> dst_tensors;
71   if (weight_desc_copy.layout ==
72           WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4 ||
73       weight_desc_copy.layout ==
74           WeightsLayout::k2DX4O4YIsSpatialIAndXIsOOGroupI4) {
75     dst_tensors.resize(4);
76     const int dst_depth = AlignByN(DivideRoundUp(weights.shape.o, 4),
77                                    weight_desc_copy.output_group_size);
78     const int src_depth = DivideRoundUp(weights.shape.i, 4);
79     const int kernel_x = weights.shape.w;
80     const int kernel_y = weights.shape.h;
81     int texture_width = dst_depth;
82     int texture_height = src_depth * kernel_x * kernel_y;
83     int sub_size = SizeOf(weights_type) * 4 * texture_width * texture_height;
84     for (int i = 0; i < 4; ++i) {
85       dst_tensors[i].shape = BHWC(1, texture_height, texture_width, 4);
86       dst_tensors[i].data.resize(4 * texture_width * texture_height);
87       memcpy(dst_tensors[i].data.data(), weights_data.data() + sub_size * i,
88              sub_size);
89     }
90   } else {
91     dst_tensors.resize(1);
92     dst_tensors[0].shape = BHWC(1, 1, 1, flt_count);
93     dst_tensors[0].data.resize(flt_count);
94     memcpy(dst_tensors[0].data.data(), weights_data.data(),
95            flt_count * SizeOf(weights_type));
96   }
97 
98   std::vector<TensorFloat32> dst_tensors_gpu(dst_tensors.size());
99   std::vector<TensorFloat32*> dst_ptrs;
100   std::vector<BHWC> dst_shapes;
101   for (int i = 0; i < dst_tensors.size(); ++i) {
102     dst_shapes.push_back(dst_tensors[i].shape);
103     dst_ptrs.push_back(&dst_tensors_gpu[i]);
104   }
105 
106   auto converter_from_ohwi = ConverterToConvWeights(
107       op_def, weight_desc, /*input layout*/ Layout::OHWI);
108   RETURN_IF_ERROR(env->ExecuteGPUOperation(
109       {src_tensor_as_ohwi},
110       std::make_unique<ConverterToConvWeights>(std::move(converter_from_ohwi)),
111       dst_shapes, dst_ptrs));
112   for (int i = 0; i < dst_tensors.size(); ++i) {
113     RETURN_IF_ERROR(
114         PointWiseNear(dst_tensors[i].data, dst_tensors_gpu[i].data, 0.0f));
115   }
116 
117   auto converter_from_hwio = ConverterToConvWeights(
118       op_def, weight_desc, /*input layout*/ Layout::HWIO);
119   RETURN_IF_ERROR(env->ExecuteGPUOperation(
120       {src_tensor_as_hwio},
121       std::make_unique<ConverterToConvWeights>(std::move(converter_from_hwio)),
122       dst_shapes, dst_ptrs));
123   for (int i = 0; i < dst_tensors.size(); ++i) {
124     RETURN_IF_ERROR(
125         PointWiseNear(dst_tensors[i].data, dst_tensors_gpu[i].data, 0.0f));
126   }
127   return absl::OkStatus();
128 }
129 
130 }  // namespace
131 
ConverterToConvWeights1x1OutX4Test(TestExecutionEnvironment * env)132 absl::Status ConverterToConvWeights1x1OutX4Test(TestExecutionEnvironment* env) {
133   const int kSrcChannels = 8;
134   const int kDstChannels = 32;
135   auto weights_shape = OHWI(kDstChannels, 1, 1, kSrcChannels);
136   WeightsDescription conv_weight_desc;
137   conv_weight_desc.output_group_size = 4;
138 
139   Tensor<OHWI, DataType::FLOAT32> weights;
140   weights.shape = weights_shape;
141   weights.data.resize(weights_shape.DimensionsProduct());
142   for (int i = 0; i < weights.data.size(); ++i) {
143     weights.data[i] = half(static_cast<float>(i));
144   }
145 
146   for (auto precision : env->GetSupportedPrecisions()) {
147     auto data_type = DeduceDataTypeFromPrecision(precision);
148     for (auto storage : env->GetSupportedStorages(data_type)) {
149       for (auto weights_layout : {WeightsLayout::kOSpatialIOGroupI4O4,
150                                   WeightsLayout::kOSpatialIOGroupO4I4}) {
151         conv_weight_desc.layout = weights_layout;
152         OperationDef op_def;
153         op_def.precision = precision;
154         op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
155         op_def.dst_tensors.push_back(
156             {data_type, TensorStorageType::BUFFER, Layout::UNKNOWN});
157         RETURN_IF_ERROR(ConvolutionWeightsConverterTest(
158             weights, conv_weight_desc, env, op_def));
159       }
160     }
161   }
162   return absl::OkStatus();
163 }
164 
ConverterToConvWeights1x1OutX4UnalignedTest(TestExecutionEnvironment * env)165 absl::Status ConverterToConvWeights1x1OutX4UnalignedTest(
166     TestExecutionEnvironment* env) {
167   const int kSrcChannels = 8;
168   const int kDstChannels = 17;
169   auto weights_shape = OHWI(kDstChannels, 1, 1, kSrcChannels);
170   WeightsDescription conv_weight_desc;
171   conv_weight_desc.output_group_size = 4;
172 
173   Tensor<OHWI, DataType::FLOAT32> weights;
174   weights.shape = weights_shape;
175   weights.data.resize(weights_shape.DimensionsProduct());
176   for (int i = 0; i < weights.data.size(); ++i) {
177     weights.data[i] = half(static_cast<float>(i));
178   }
179 
180   for (auto precision : env->GetSupportedPrecisions()) {
181     auto data_type = DeduceDataTypeFromPrecision(precision);
182     for (auto storage : env->GetSupportedStorages(data_type)) {
183       for (auto weights_layout : {WeightsLayout::kOSpatialIOGroupI4O4,
184                                   WeightsLayout::kOSpatialIOGroupO4I4}) {
185         conv_weight_desc.layout = weights_layout;
186         OperationDef op_def;
187         op_def.precision = precision;
188         op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
189         op_def.dst_tensors.push_back(
190             {data_type, TensorStorageType::BUFFER, Layout::UNKNOWN});
191         RETURN_IF_ERROR(ConvolutionWeightsConverterTest(
192             weights, conv_weight_desc, env, op_def));
193       }
194     }
195   }
196   return absl::OkStatus();
197 }
198 
ConverterToConvWeights1x1OutX2Test(TestExecutionEnvironment * env)199 absl::Status ConverterToConvWeights1x1OutX2Test(TestExecutionEnvironment* env) {
200   const int kSrcChannels = 7;
201   const int kDstChannels = 37;
202   auto weights_shape = OHWI(kDstChannels, 1, 1, kSrcChannels);
203   WeightsDescription conv_weight_desc;
204   conv_weight_desc.output_group_size = 2;
205 
206   Tensor<OHWI, DataType::FLOAT32> weights;
207   weights.shape = weights_shape;
208   weights.data.resize(weights_shape.DimensionsProduct());
209   for (int i = 0; i < weights.data.size(); ++i) {
210     weights.data[i] = half(static_cast<float>(i));
211   }
212 
213   for (auto precision : env->GetSupportedPrecisions()) {
214     auto data_type = DeduceDataTypeFromPrecision(precision);
215     for (auto storage : env->GetSupportedStorages(data_type)) {
216       for (auto weights_layout : {WeightsLayout::kOSpatialIOGroupI4O4,
217                                   WeightsLayout::kOSpatialIOGroupO4I4}) {
218         conv_weight_desc.layout = weights_layout;
219         OperationDef op_def;
220         op_def.precision = precision;
221         op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
222         op_def.dst_tensors.push_back(
223             {data_type, TensorStorageType::BUFFER, Layout::UNKNOWN});
224         RETURN_IF_ERROR(ConvolutionWeightsConverterTest(
225             weights, conv_weight_desc, env, op_def));
226       }
227     }
228   }
229   return absl::OkStatus();
230 }
231 
ConverterToConvWeightsOutX2Test(TestExecutionEnvironment * env)232 absl::Status ConverterToConvWeightsOutX2Test(TestExecutionEnvironment* env) {
233   const int kSrcChannels = 8;
234   const int kDstChannels = 38;
235   auto weights_shape = OHWI(kDstChannels, 3, 4, kSrcChannels);
236   WeightsDescription conv_weight_desc;
237   conv_weight_desc.output_group_size = 2;
238 
239   Tensor<OHWI, DataType::FLOAT32> weights;
240   weights.shape = weights_shape;
241   weights.data.resize(weights_shape.DimensionsProduct());
242   for (int i = 0; i < weights.data.size(); ++i) {
243     weights.data[i] = half(static_cast<float>(i));
244   }
245 
246   for (auto precision : env->GetSupportedPrecisions()) {
247     auto data_type = DeduceDataTypeFromPrecision(precision);
248     for (auto storage : env->GetSupportedStorages(data_type)) {
249       for (auto weights_layout : {WeightsLayout::kOSpatialIOGroupI4O4,
250                                   WeightsLayout::kOSpatialIOGroupO4I4}) {
251         conv_weight_desc.layout = weights_layout;
252         OperationDef op_def;
253         op_def.precision = precision;
254         op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
255         op_def.dst_tensors.push_back(
256             {data_type, TensorStorageType::BUFFER, Layout::UNKNOWN});
257         RETURN_IF_ERROR(ConvolutionWeightsConverterTest(
258             weights, conv_weight_desc, env, op_def));
259       }
260     }
261   }
262   return absl::OkStatus();
263 }
264 
ConverterToConvTransposedWeights4x4Test(TestExecutionEnvironment * env)265 absl::Status ConverterToConvTransposedWeights4x4Test(
266     TestExecutionEnvironment* env) {
267   const int kSrcChannels = 7;
268   const int kDstChannels = 11;
269   auto weights_shape = OHWI(kDstChannels, 4, 4, kSrcChannels);
270   WeightsDescription weight_desc;
271   weight_desc.spatial_remap = {10, 11, 14, 15, 8, 9, 12, 13,
272                                2,  3,  6,  7,  0, 1, 4,  5};
273 
274   Tensor<OHWI, DataType::FLOAT32> weights;
275   weights.shape = weights_shape;
276   weights.data.resize(weights_shape.DimensionsProduct());
277   for (int i = 0; i < weights.data.size(); ++i) {
278     weights.data[i] = half(static_cast<float>(i));
279   }
280 
281   for (auto precision : env->GetSupportedPrecisions()) {
282     auto data_type = DeduceDataTypeFromPrecision(precision);
283     for (auto storage : env->GetSupportedStorages(data_type)) {
284       for (auto weights_layout : {WeightsLayout::kOICustomSpatialI4O4,
285                                   WeightsLayout::kOICustomSpatialO4I4}) {
286         weight_desc.layout = weights_layout;
287         OperationDef op_def;
288         op_def.precision = precision;
289         op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
290         op_def.dst_tensors.push_back(
291             {data_type, TensorStorageType::BUFFER, Layout::UNKNOWN});
292         RETURN_IF_ERROR(
293             ConvolutionWeightsConverterTest(weights, weight_desc, env, op_def));
294       }
295     }
296   }
297   return absl::OkStatus();
298 }
299 
ConverterToConvWeights4xTexturesTest(TestExecutionEnvironment * env)300 absl::Status ConverterToConvWeights4xTexturesTest(
301     TestExecutionEnvironment* env) {
302   const int src_channels = 9;
303   const int dst_channels = 17;
304   auto weights_shape = OHWI(dst_channels, 1, 1, src_channels);
305   WeightsDescription conv_weight_desc;
306   conv_weight_desc.output_group_size = 4;
307 
308   Tensor<OHWI, DataType::FLOAT32> weights;
309   weights.shape = weights_shape;
310   weights.data.resize(weights_shape.DimensionsProduct());
311   for (int i = 0; i < weights.data.size(); ++i) {
312     weights.data[i] = half(static_cast<float>(i));
313   }
314 
315   for (auto precision : env->GetSupportedPrecisions()) {
316     auto data_type = DeduceDataTypeFromPrecision(precision);
317     for (auto storage : env->GetSupportedStorages(data_type)) {
318       for (auto weights_layout :
319            {WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4,
320             WeightsLayout::k2DX4O4YIsSpatialIAndXIsOOGroupI4}) {
321         conv_weight_desc.layout = weights_layout;
322         OperationDef op_def;
323         op_def.precision = precision;
324         op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
325         op_def.dst_tensors.push_back(
326             {data_type, TensorStorageType::TEXTURE_2D, Layout::HWC});
327         op_def.dst_tensors.push_back(
328             {data_type, TensorStorageType::TEXTURE_2D, Layout::HWC});
329         op_def.dst_tensors.push_back(
330             {data_type, TensorStorageType::TEXTURE_2D, Layout::HWC});
331         op_def.dst_tensors.push_back(
332             {data_type, TensorStorageType::TEXTURE_2D, Layout::HWC});
333         RETURN_IF_ERROR(ConvolutionWeightsConverterTest(
334             weights, conv_weight_desc, env, op_def));
335       }
336     }
337   }
338   return absl::OkStatus();
339 }
340 
341 }  // namespace gpu
342 }  // namespace tflite
343