chore: import upstream snapshot with attribution

tensorflow/lite/delegates/gpu/common/tasks/add.cc

@@ -0,0 +1,94 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/add.h"
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/strings/str_cat.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+
+namespace tflite {
+namespace gpu {
+namespace {
+GPUOperation CreateUnequalAdd(const OperationDef& op_def) {
+  GPUOperation op(op_def);
+  op.AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
+  for (int i = 0; i < op_def.src_tensors.size(); ++i) {
+    const std::string tensor_name = absl::StrCat("src_tensor_", i);
+    op.AddSrcTensor(tensor_name, op_def.src_tensors[i]);
+  }
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+  std::string c;
+  c += "MAIN_FUNCTION($0) {\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+    for (int i = 0; i < op_def.src_tensors.size(); ++i) {
+      const std::string tensor_name = absl::StrCat("src_tensor_", i);
+      c += "  args." + tensor_name + ".SetBatchRef(B);\n";
+    }
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  c += "  int Y = GLOBAL_ID_1;\n";
+  c += "  int S = GLOBAL_ID_2;\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "S >= args.dst_tensor.Slices()) return; \n";
+  c += "  args.src_tensor_0::type src = args.src_tensor_0::zero_value;\n";
+  for (int i = 0; i < op_def.src_tensors.size(); ++i) {
+    const std::string tensor_name = absl::StrCat("src_tensor_", i);
+    c += "  if (S < args." + tensor_name + ".Slices()) {\n";
+    c += "    src += args." + tensor_name + ".Read(X, Y, S);\n";
+    c += "  }\n";
+  }
+  c += "  args.dst_tensor.Write(src, X, Y, S);\n";
+  c += "} \n";
+  op.code_ = std::move(c);
+  return op;
+}
+}  // namespace
+
+GPUOperation CreateAdd(const OperationDef& definition,
+                       const std::vector<int>& channels, int dst_channels) {
+  if (dst_channels != channels[0]) {
+    return CreateUnequalAdd(definition);
+  }
+  ElementwiseDescriptor op_desc;
+  op_desc.code = "  out_value = in_value;\n";
+  for (int i = 1; i < definition.src_tensors.size(); ++i) {
+    const std::string tensor_name = absl::StrCat("src_tensor_", i);
+    std::string coords = "X_COORD, Y_COORD";
+    if (definition.src_tensors[i].HasAxis(Axis::DEPTH)) {
+      coords += ", Z_COORD";
+    }
+    coords += ", S_COORD";
+    if (definition.src_tensors[i].HasAxis(Axis::BATCH)) {
+      coords += ", B_COORD";
+    }
+    op_desc.code += "if (S_COORD < args." + tensor_name + ".Slices()) {\n";
+    op_desc.code +=
+        "  out_value += args." + tensor_name + ".Read(" + coords + ");\n";
+    op_desc.code += "}\n";
+  }
+  return CreateGpuOperation(definition, std::move(op_desc));
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/add.h

@@ -0,0 +1,37 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ADD_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ADD_H_
+
+#include <string>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+
+namespace tflite {
+namespace gpu {
+
+// Add operation supports not equal tensors on input (for possibility to
+// remove Padding operation with zeroes in channels dimension)
+GPUOperation CreateAdd(const OperationDef& definition,
+                       const std::vector<int>& channels, int dst_channels);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ADD_H_

tensorflow/lite/delegates/gpu/common/tasks/add_test_util.cc

@@ -0,0 +1,215 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/add_test_util.h"
+
+#include <memory>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/add.h"
+
+namespace tflite {
+namespace gpu {
+
+template <DataType T>
+absl::Status AddTwoEqualIntTensorsTest(TestExecutionEnvironment* env) {
+  tflite::gpu::Tensor<BHWC, T> src0, src1;
+  src0.shape = BHWC(1, 2, 1, 2);
+  src0.data = {3, 4, 5, 6};
+  src1.shape = BHWC(1, 2, 1, 2);
+  src1.data = {-4, 12, 17, -2};
+  std::vector<int> channels = {2, 2};
+  tflite::gpu::Tensor<BHWC, T> ref_tensor;
+  ref_tensor.shape = BHWC(1, 2, 1, 2);
+  ref_tensor.data = {-1, 16, 22, 4};
+
+  for (auto storage : env->GetSupportedStorages(T)) {
+    OperationDef op_def;
+    op_def.precision = CalculationsPrecision::F32;
+    op_def.src_tensors.push_back({T, storage, Layout::HWC});
+    op_def.src_tensors.push_back({T, storage, Layout::HWC});
+    op_def.dst_tensors.push_back({T, storage, Layout::HWC});
+    TensorDescriptor src_0, src_1, dst;
+    src_0 = op_def.src_tensors[0];
+    src_1 = op_def.src_tensors[1];
+    src_0.UploadData(src0);
+    src_1.UploadData(src1);
+    dst.SetBHWCShape(BHWC(1, 2, 1, 2));
+    GPUOperation operation = CreateAdd(op_def, channels, channels[0]);
+    RETURN_IF_ERROR(env->ExecuteGPUOperation(
+        {&src_0, &src_1}, {&dst},
+        std::make_unique<GPUOperation>(std::move(operation))));
+    tflite::gpu::Tensor<BHWC, T> dst_tensor;
+    dst.DownloadData(&dst_tensor);
+    if (dst_tensor.data != ref_tensor.data) {
+      return absl::InternalError("not equal");
+    }
+  }
+  return absl::OkStatus();
+}
+
+template absl::Status AddTwoEqualIntTensorsTest<DataType::INT32>(
+    TestExecutionEnvironment* env);
+template absl::Status AddTwoEqualIntTensorsTest<DataType::INT16>(
+    TestExecutionEnvironment* env);
+template absl::Status AddTwoEqualIntTensorsTest<DataType::INT8>(
+    TestExecutionEnvironment* env);
+
+template <DataType T>
+absl::Status AddTwoEqualUintTensorsTest(TestExecutionEnvironment* env) {
+  tflite::gpu::Tensor<BHWC, T> src0, src1;
+  src0.shape = BHWC(1, 2, 1, 2);
+  src0.data = {3, 4, 5, 6};
+  src1.shape = BHWC(1, 2, 1, 2);
+  src1.data = {4, 12, 17, 2};
+  std::vector<int> channels = {2, 2};
+  tflite::gpu::Tensor<BHWC, T> ref_tensor;
+  ref_tensor.shape = BHWC(1, 2, 1, 2);
+  ref_tensor.data = {7, 16, 22, 8};
+
+  for (auto storage : env->GetSupportedStorages(T)) {
+    OperationDef op_def;
+    op_def.precision = CalculationsPrecision::F32;
+    op_def.src_tensors.push_back({T, storage, Layout::HWC});
+    op_def.src_tensors.push_back({T, storage, Layout::HWC});
+    op_def.dst_tensors.push_back({T, storage, Layout::HWC});
+    TensorDescriptor src_0, src_1, dst;
+    src_0 = op_def.src_tensors[0];
+    src_1 = op_def.src_tensors[1];
+    src_0.UploadData(src0);
+    src_1.UploadData(src1);
+    dst.SetBHWCShape(BHWC(1, 2, 1, 2));
+    GPUOperation operation = CreateAdd(op_def, channels, channels[0]);
+    RETURN_IF_ERROR(env->ExecuteGPUOperation(
+        {&src_0, &src_1}, {&dst},
+        std::make_unique<GPUOperation>(std::move(operation))));
+    tflite::gpu::Tensor<BHWC, T> dst_tensor;
+    dst.DownloadData(&dst_tensor);
+    if (dst_tensor.data != ref_tensor.data) {
+      return absl::InternalError("not equal");
+    }
+  }
+  return absl::OkStatus();
+}
+
+template absl::Status AddTwoEqualUintTensorsTest<DataType::UINT32>(
+    TestExecutionEnvironment* env);
+template absl::Status AddTwoEqualUintTensorsTest<DataType::UINT16>(
+    TestExecutionEnvironment* env);
+template absl::Status AddTwoEqualUintTensorsTest<DataType::UINT8>(
+    TestExecutionEnvironment* env);
+
+absl::Status AddTwoEqualTensorsTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src0, src1;
+  src0.shape = BHWC(1, 2, 1, 2);
+  src0.data = {0.0f, -1.0f, -0.05f, 0.045f};
+  src1.shape = BHWC(1, 2, 1, 2);
+  src1.data = {0.0f, 1.0f, -0.05f, -0.045f};
+  std::vector<int> channels = {2, 2};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreateAdd(op_def, channels, channels[0]);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {src0, src1}, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 2, 1, 2), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({0.0f, 0.0f, -0.1f, 0.0f}, dst_tensor.data, eps));
+    }
+  }
+
+  RETURN_IF_ERROR(AddTwoEqualIntTensorsTest<DataType::INT32>(env));
+  RETURN_IF_ERROR(AddTwoEqualIntTensorsTest<DataType::INT16>(env));
+  RETURN_IF_ERROR(AddTwoEqualIntTensorsTest<DataType::INT8>(env));
+  RETURN_IF_ERROR(AddTwoEqualUintTensorsTest<DataType::UINT32>(env));
+  RETURN_IF_ERROR(AddTwoEqualUintTensorsTest<DataType::UINT16>(env));
+  RETURN_IF_ERROR(AddTwoEqualUintTensorsTest<DataType::UINT8>(env));
+  return absl::OkStatus();
+}
+
+absl::Status AddFirstTensorHasMoreChannelsThanSecondTest(
+    TestExecutionEnvironment* env) {
+  TensorFloat32 src0, src1;
+  src0.shape = BHWC(1, 2, 1, 6);
+  src0.data = {0.0f,   -1.0f,  -0.05f, 0.045f, 1.0f,   -2.0f,
+               -1.05f, 1.045f, 2.0f,   -3.0f,  -2.05f, 2.045f};
+  src1.shape = BHWC(1, 2, 1, 2);
+  src1.data = {0.0f, 1.0f, -0.05f, -0.045f};
+  std::vector<int> channels = {6, 2};
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreateAdd(op_def, channels, channels[0]);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {src0, src1}, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 2, 1, 6), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear({0.0f, 0.0f, -0.05f, 0.045f, 1.0f, -2.0f,
+                                     -1.1f, 1.0f, 2.0f, -3.0f, -2.05f, 2.045f},
+                                    dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status AddFirstTensorHasLessChannelsThanSecond(
+    TestExecutionEnvironment* env) {
+  TensorFloat32 src0, src1;
+  src1.shape = BHWC(1, 2, 1, 6);
+  src1.data = {0.0f,   -1.0f,  -0.05f, 0.045f, 1.0f,   -2.0f,
+               -1.05f, 1.045f, 2.0f,   -3.0f,  -2.05f, 2.045f};
+  src0.shape = BHWC(1, 2, 1, 2);
+  src0.data = {0.0f, 1.0f, -0.05f, -0.045f};
+  std::vector<int> channels = {2, 6};
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreateAdd(op_def, channels, 6);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {src0, src1}, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 2, 1, 6), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear({0.0f, 0.0f, -0.05f, 0.045f, 1.0f, -2.0f,
+                                     -1.1f, 1.0f, 2.0f, -3.0f, -2.05f, 2.045f},
+                                    dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/add_test_util.h

@@ -0,0 +1,36 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ADD_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ADD_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status AddTwoEqualTensorsTest(TestExecutionEnvironment* env);
+
+absl::Status AddFirstTensorHasMoreChannelsThanSecondTest(
+    TestExecutionEnvironment* env);
+
+absl::Status AddFirstTensorHasLessChannelsThanSecond(
+    TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ADD_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/cast.cc

@@ -0,0 +1,40 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/cast.h"
+
+#include <map>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/strings/substitute.h"
+#include "tensorflow/lite/delegates/gpu/common/task/util.h"
+
+namespace tflite {
+namespace gpu {
+GPUOperation CreateCast(const OperationDef& definition,
+                        const GpuInfo& gpu_info) {
+  ElementwiseDescriptor op_desc;
+  const std::string conversion =
+      GetTypeConversion(gpu_info, definition.src_tensors[0].GetDataType(),
+                        definition.dst_tensors[0].GetDataType(), 4);
+  op_desc.code =
+      "out_value = " + absl::Substitute(conversion, "in_value") + ";\n";
+  return CreateGpuOperation(definition, std::move(op_desc));
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/cast.h

@@ -0,0 +1,30 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CAST_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CAST_H_
+
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+
+namespace tflite {
+namespace gpu {
+
+GPUOperation CreateCast(const OperationDef& definition,
+                        const GpuInfo& gpu_info);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CAST_H_

tensorflow/lite/delegates/gpu/common/tasks/cast_test_util.cc

@@ -0,0 +1,133 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/cast_test_util.h"
+
+#include <memory>
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/cast.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status CastTests(TestExecutionEnvironment* env) {
+  tflite::gpu::Tensor<BHWC, DataType::FLOAT32> src;
+  src.shape = BHWC(1, 2, 1, 2);
+  src.data = {0.0f, -1.3f, -7.4f, 12.45f};
+
+  tflite::gpu::Tensor<BHWC, DataType::INT32> ref_tensor;
+  ref_tensor.shape = BHWC(1, 2, 1, 2);
+  ref_tensor.data = {0, -1, -7, 12};
+
+  for (auto storage_src : env->GetSupportedStorages(DataType::FLOAT16)) {
+    for (auto storage_dst : env->GetSupportedStorages(DataType::INT32)) {
+      OperationDef op_def;
+      op_def.precision = CalculationsPrecision::F16;
+      op_def.src_tensors.push_back(
+          {DataType::FLOAT16, storage_src, Layout::HWC});
+      op_def.dst_tensors.push_back({DataType::INT32, storage_dst, Layout::HWC});
+      TensorDescriptor src_desc, dst_desc;
+      src_desc = op_def.src_tensors[0];
+      src_desc.UploadData(src);
+      dst_desc.SetBHWCShape(BHWC(1, 2, 1, 2));
+      GPUOperation operation = CreateCast(op_def, env->GetGpuInfo());
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {&src_desc}, {&dst_desc},
+          std::make_unique<GPUOperation>(std::move(operation))));
+
+      tflite::gpu::Tensor<BHWC, DataType::INT32> dst_tensor;
+      dst_desc.DownloadData(&dst_tensor);
+      if (dst_tensor.data != ref_tensor.data) {
+        return absl::InternalError("not equal");
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status CastToBoolTests(TestExecutionEnvironment* env) {
+  tflite::gpu::Tensor<BHWC, DataType::FLOAT32> src;
+  src.shape = BHWC(1, 2, 1, 2);
+  src.data = {0.0f, -1.3f, -7.4f, 12.45f};
+
+  tflite::gpu::Tensor<BHWC, DataType::BOOL> ref_tensor;
+  ref_tensor.shape = BHWC(1, 2, 1, 2);
+  ref_tensor.data = {false, true, true, true};
+
+  for (auto storage_src : env->GetSupportedStorages(DataType::FLOAT32)) {
+    for (auto storage_dst : env->GetSupportedStorages(DataType::BOOL)) {
+      OperationDef op_def;
+      op_def.precision = CalculationsPrecision::F32;
+      op_def.src_tensors.push_back(
+          {DataType::FLOAT32, storage_src, Layout::HWC});
+      op_def.dst_tensors.push_back({DataType::BOOL, storage_dst, Layout::HWC});
+      TensorDescriptor src_desc, dst_desc;
+      src_desc = op_def.src_tensors[0];
+      src_desc.UploadData(src);
+      dst_desc.SetBHWCShape(BHWC(1, 2, 1, 2));
+      GPUOperation operation = CreateCast(op_def, env->GetGpuInfo());
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {&src_desc}, {&dst_desc},
+          std::make_unique<GPUOperation>(std::move(operation))));
+
+      tflite::gpu::Tensor<BHWC, DataType::BOOL> dst_tensor;
+      dst_desc.DownloadData(&dst_tensor);
+      if (dst_tensor.data != ref_tensor.data) {
+        return absl::InternalError("not equal");
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status CastFromBoolTests(TestExecutionEnvironment* env) {
+  tflite::gpu::Tensor<BHWC, DataType::BOOL> src;
+  src.shape = BHWC(1, 2, 1, 2);
+  src.data = {false, true, true, true};
+
+  tflite::gpu::Tensor<BHWC, DataType::FLOAT32> ref_tensor;
+  ref_tensor.shape = BHWC(1, 2, 1, 2);
+  ref_tensor.data = {0.0, 1.0, 1.0, 1.0};
+
+  for (auto storage_src : env->GetSupportedStorages(DataType::BOOL)) {
+    for (auto storage_dst : env->GetSupportedStorages(DataType::FLOAT32)) {
+      OperationDef op_def;
+      op_def.precision = CalculationsPrecision::F32;
+      op_def.src_tensors.push_back({DataType::BOOL, storage_src, Layout::HWC});
+      op_def.dst_tensors.push_back(
+          {DataType::FLOAT32, storage_dst, Layout::HWC});
+      TensorDescriptor src_desc, dst_desc;
+      src_desc = op_def.src_tensors[0];
+      src_desc.UploadData(src);
+      dst_desc.SetBHWCShape(BHWC(1, 2, 1, 2));
+      GPUOperation operation = CreateCast(op_def, env->GetGpuInfo());
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {&src_desc}, {&dst_desc},
+          std::make_unique<GPUOperation>(std::move(operation))));
+
+      tflite::gpu::Tensor<BHWC, DataType::FLOAT32> dst_tensor;
+      dst_desc.DownloadData(&dst_tensor);
+      if (dst_tensor.data != ref_tensor.data) {
+        return absl::InternalError("not equal");
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/cast_test_util.h

@@ -0,0 +1,34 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CAST_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CAST_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status CastTests(TestExecutionEnvironment* env);
+
+absl::Status CastToBoolTests(TestExecutionEnvironment* env);
+
+absl::Status CastFromBoolTests(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CAST_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/concat_test_util.cc

@@ -0,0 +1,171 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/concat_test_util.h"
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/concat_xy.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/concat_z.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConcatWidthTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src0, src1;
+  src0.shape = BHWC(1, 2, 1, 2);
+  src0.data = {half(0.0f), half(-1.0f), half(-0.05f), half(0.045f)};
+  src1.shape = BHWC(1, 2, 2, 2);
+  src1.data = {half(1.0f), half(-1.2f), half(-0.45f), half(1.045f),
+               half(1.1f), half(-1.3f), half(-0.55f), half(2.045f)};
+
+  ConcatAttributes attr;
+  attr.axis = Axis::WIDTH;
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreateConcatXY(op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {src0, src1}, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 2, 3, 2), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({half(0.0f), half(-1.0f), half(1.0f), half(-1.2f),
+                         half(-0.45f), half(1.045f), half(-0.05f), half(0.045f),
+                         half(1.1f), half(-1.3f), half(-0.55f), half(2.045f)},
+                        dst_tensor.data, 0.0f));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConcatHeightTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src0, src1;
+  src0.shape = BHWC(1, 2, 1, 2);
+  src0.data = {half(0.0f), half(-1.0f), half(-0.05f), half(0.045f)};
+  src1.shape = BHWC(1, 1, 1, 2);
+  src1.data = {half(1.0f), half(-1.2f)};
+
+  ConcatAttributes attr;
+  attr.axis = Axis::HEIGHT;
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreateConcatXY(op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {src0, src1}, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 3, 1, 2), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear({half(0.0f), half(-1.0f), half(-0.05f),
+                                     half(0.045f), half(1.0f), half(-1.2f)},
+                                    dst_tensor.data, 0.0f));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConcatChannelsTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src0, src1, src2;
+  src0.shape = BHWC(1, 2, 1, 1);
+  src0.data = {half(0.0f), half(-1.0f)};
+  src1.shape = BHWC(1, 2, 1, 2);
+  src1.data = {half(1.0f), half(2.0f), half(3.0f), half(4.0f)};
+  src2.shape = BHWC(1, 2, 1, 3);
+  src2.data = {half(5.0f), half(6.0f), half(7.0f),
+               half(8.0f), half(9.0),  half(10.0f)};
+
+  ConcatAttributes attr;
+  attr.axis = Axis::CHANNELS;
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation =
+          CreateConcatZ(op_def, {1, 2, 3}, env->GetGpuInfo());
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {src0, src1, src2},
+          std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 2, 1, 6), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({half(0.0f), half(1.0f), half(2.0f), half(5.0f),
+                         half(6.0f), half(7.0f), half(-1.0f), half(3.0f),
+                         half(4.0f), half(8.0f), half(9.0), half(10.0f)},
+                        dst_tensor.data, 0.0f));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConcatChannelsAlignedx4Test(TestExecutionEnvironment* env) {
+  TensorFloat32 src0, src1;
+  src0.shape = BHWC(1, 2, 1, 4);
+  src0.data = {half(-1.0f), half(-2.0f), half(-3.0f), half(-4.0f),
+               half(1.0f),  half(2.0f),  half(3.0f),  half(4.0f)};
+  src1.shape = BHWC(1, 2, 1, 4);
+  src1.data = {half(5.0f),  half(6.0f),  half(7.0f),  half(8.0f),
+               half(-5.0f), half(-6.0f), half(-7.0f), half(-8.0f)};
+
+  ConcatAttributes attr;
+  attr.axis = Axis::CHANNELS;
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreateConcatZ(op_def, {4, 4}, env->GetGpuInfo());
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {src0, src1}, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 2, 1, 8), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({half(-1.0f), half(-2.0f), half(-3.0f), half(-4.0f),
+                         half(5.0f), half(6.0f), half(7.0f), half(8.0f),
+                         half(1.0f), half(2.0f), half(3.0f), half(4.0f),
+                         half(-5.0f), half(-6.0f), half(-7.0f), half(-8.0f)},
+                        dst_tensor.data, 0.0f));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/concat_test_util.h

@@ -0,0 +1,33 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONCAT_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONCAT_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConcatWidthTest(TestExecutionEnvironment* env);
+absl::Status ConcatHeightTest(TestExecutionEnvironment* env);
+absl::Status ConcatChannelsTest(TestExecutionEnvironment* env);
+absl::Status ConcatChannelsAlignedx4Test(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONCAT_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/concat_xy.cc

@@ -0,0 +1,128 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/concat_xy.h"
+
+#include <map>
+#include <string>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+namespace {
+std::string GetConcatKernelCode(const OperationDef& op_def,
+                                const ConcatAttributes& attr) {
+  std::vector<std::string> tensor_names(op_def.src_tensors.size());
+  for (int i = 0; i < op_def.src_tensors.size(); ++i) {
+    tensor_names[i] = "src_tensor_" + std::to_string(i);
+  }
+
+  std::map<Axis, std::string> axis_to_selector = {
+      {Axis::WIDTH, "Width"}, {Axis::HEIGHT, "Height"},
+      {Axis::DEPTH, "Depth"}, {Axis::CHANNELS, "Channels"},
+      {Axis::BATCH, "Batch"},
+  };
+  std::map<Axis, std::string> axis_to_coord = {
+      {Axis::WIDTH, "X"},    {Axis::HEIGHT, "Y"}, {Axis::DEPTH, "D"},
+      {Axis::CHANNELS, "S"}, {Axis::BATCH, "B"},
+  };
+
+  std::vector<std::string> src_coords;
+  std::vector<std::string> dst_coords;
+  for (auto axis :
+       {Axis::WIDTH, Axis::HEIGHT, Axis::DEPTH, Axis::CHANNELS, Axis::BATCH}) {
+    if (op_def.src_tensors[0].HasAxis(axis) && axis != Axis::BATCH) {
+      if (axis == attr.axis) {
+        src_coords.push_back("coord");
+      } else {
+        src_coords.push_back(axis_to_coord[axis]);
+      }
+    }
+    if (op_def.dst_tensors[0].HasAxis(axis)) {
+      dst_coords.push_back(axis_to_coord[axis]);
+    }
+  }
+  std::string src_coord = src_coords[0];
+  for (int i = 1; i < src_coords.size(); ++i) {
+    src_coord += ", " + src_coords[i];
+  }
+  std::string dst_coord = dst_coords[0];
+  for (int i = 1; i < dst_coords.size(); ++i) {
+    dst_coord += ", " + dst_coords[i];
+  }
+
+  std::string c;
+  c += "MAIN_FUNCTION($0) {\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int linear_id_0 = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id_0 / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id_0 % args.dst_tensor.Batch();\n";
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  int linear_id_1 = GLOBAL_ID_1;\n";
+    c += "  int Y = linear_id_1 / args.dst_tensor.Depth();\n";
+    c += "  int D = linear_id_1 % args.dst_tensor.Depth();\n";
+  } else {
+    c += "  int Y = GLOBAL_ID_1;\n";
+  }
+  c += "  int S = GLOBAL_ID_2;\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "S >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
+  c += "  args.src_tensor_0::type result = args.src_tensor_0::zero_value;\n";
+  c += "  int coord = " + axis_to_coord[attr.axis] + ";\n";
+  for (int i = 0; i < op_def.src_tensors.size(); ++i) {
+    const std::string field =
+        "args." + tensor_names[i] + "." + axis_to_selector[attr.axis] + "()";
+    c += "  if (coord >= 0 && coord < " + field + ") { \n";
+    if (op_def.src_tensors[i].HasAxis(Axis::BATCH)) {
+      if (attr.axis == Axis::BATCH) {
+        c += "  args." + tensor_names[i] + ".SetBatchRef(coord);\n";
+      } else {
+        c += "  args." + tensor_names[i] + ".SetBatchRef(B);\n";
+      }
+    }
+    c += "    result = args." + tensor_names[i] + ".Read(" + src_coord + ");\n";
+    c += "  } \n";
+    c += "  coord -= " + field + ";\n";
+  }
+  c += "  args.dst_tensor.Write(result, " + dst_coord + ");\n";
+  c += "}\n";
+  return c;
+}
+}  // namespace
+
+GPUOperation CreateConcatXY(const OperationDef& definition,
+                            const ConcatAttributes& attr) {
+  GPUOperation op(definition);
+  for (int i = 0; i < definition.src_tensors.size(); ++i) {
+    const std::string name = "src_tensor_" + std::to_string(i);
+    op.AddSrcTensor(name, definition.src_tensors[i]);
+  }
+  op.AddDstTensor("dst_tensor", definition.dst_tensors[0]);
+  op.code_ = GetConcatKernelCode(definition, attr);
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+  return op;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/concat_xy.h

@@ -0,0 +1,33 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONCAT_XY_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONCAT_XY_H_
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+GPUOperation CreateConcatXY(const OperationDef& definition,
+                            const ConcatAttributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONCAT_XY_H_

tensorflow/lite/delegates/gpu/common/tasks/concat_z.cc

@@ -0,0 +1,157 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/concat_z.h"
+
+#include <algorithm>
+#include <string>
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+namespace {
+
+bool IsAllChannelsX4(const std::vector<int>& channels) {
+  for (int channel : channels) {
+    if (channel % 4 != 0) {
+      return false;
+    }
+  }
+  return true;
+}
+
+std::string GetConcatKernelCode(const OperationDef& op_def,
+                                const std::vector<int>& channels) {
+  std::vector<std::string> tensor_names(op_def.src_tensors.size());
+  for (int i = 0; i < op_def.src_tensors.size(); ++i) {
+    tensor_names[i] = "src_tensor_" + std::to_string(i);
+  }
+
+  std::string c;
+  c += "MAIN_FUNCTION($0) {\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+    for (int i = 0; i < op_def.src_tensors.size(); ++i) {
+      c += "  args." + tensor_names[i] + ".SetBatchRef(B);\n";
+    }
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  c += "  int Y = GLOBAL_ID_1;\n";
+  std::string coords = "X, Y";
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  int Z = GLOBAL_ID_2;\n";
+    c += "  if (Z >= args.dst_tensor.Depth()) return;\n";
+    coords = "X, Y, Z";
+  }
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height()) "
+       "return; \n";
+
+  if (IsAllChannelsX4(channels)) {
+    // When all channels % 4 == 0 we can read/assign/write VEC4 elements easily.
+    // Also it is easy to write a loop in this case, to prevent long kernel
+    // generation.
+    c += "  int S = 0;\n";
+    for (int i = 0; i < channels.size(); ++i) {
+      std::string t_name = "args." + tensor_names[i];
+      const int src_depth = DivideRoundUp(channels[i], 4);
+      if (src_depth % 2 == 0) {
+        // We can read more at once inside of loop in case src_depth % 2 == 0
+        // it should be better for reading latency hiding
+        c += "  for (int i = 0; i < " + t_name + ".Slices(); i += 2) {\n";
+        c += "    " + t_name + "::type result0 = " + t_name + ".Read(" +
+             coords + ", i);\n";
+        c += "    " + t_name + "::type result1 = " + t_name + ".Read(" +
+             coords + ", i + 1);\n";
+        c += "    args.dst_tensor.Write(result0, " + coords + ", S);\n";
+        c += "    args.dst_tensor.Write(result1, " + coords + ", S + 1);\n";
+        c += "    S += 2;\n";
+        c += "  }\n";
+      } else {
+        c += "  for (int i = 0; i < " + t_name + ".Slices(); ++i) {\n";
+        c += "    " + t_name + "::type result = " + t_name + ".Read(" + coords +
+             ", i);\n";
+        c += "    args.dst_tensor.Write(result, " + coords + ", S);\n";
+        c += "    S++;\n";
+        c += "  }\n";
+      }
+    }
+  } else {
+    c += "  args.src_tensor_0::type result = args.src_tensor_0::zero_value;\n";
+    int out_channel = 0;
+    int read_index = 0;
+    int z = 0;
+    const std::string postfix[] = {".x", ".y", ".z", ".w"};
+    for (int i = 0; i < channels.size(); ++i) {
+      std::string tensor_name = "args." + tensor_names[i];
+      const int depth = DivideRoundUp(channels[i], 4);
+      for (int d = 0; d < depth; ++d) {
+        const int channels_in_group = std::min(4, channels[i] - d * 4);
+        const std::string temp_name = "t" + std::to_string(read_index);
+        c += "  " + tensor_name + "::type " + temp_name + " = " + tensor_name +
+             ".Read(" + coords + ", " + std::to_string(d) + ");\n";
+        for (int ch = 0; ch < channels_in_group; ++ch) {
+          c += "  result" + postfix[out_channel] + " = ";
+          c += temp_name + postfix[ch] + ";\n";
+          out_channel++;
+          if (out_channel == 4) {
+            out_channel = 0;
+            c += "  args.dst_tensor.Write(result, " + coords + ", " +
+                 std::to_string(z) + ");\n";
+            z++;
+          }
+        }
+        read_index++;
+      }
+    }
+    if (out_channel != 0) {
+      c += "  args.dst_tensor.Write(result, " + coords + ", " +
+           std::to_string(z) + ");\n";
+    }
+  }
+  c += "}\n";
+  return c;
+}
+
+}  // namespace
+
+GPUOperation CreateConcatZ(const OperationDef& definition,
+                           const std::vector<int>& channels,
+                           const GpuInfo& gpu_info) {
+  GPUOperation op(definition);
+  for (int i = 0; i < definition.src_tensors.size(); ++i) {
+    const std::string name = "src_tensor_" + std::to_string(i);
+    op.AddSrcTensor(name, definition.src_tensors[i]);
+  }
+  op.AddDstTensor("dst_tensor", definition.dst_tensors[0]);
+  op.code_ = GetConcatKernelCode(definition, channels);
+  if (gpu_info.IsPowerVR() &&
+      definition.precision == CalculationsPrecision::F32 &&
+      !IsAllChannelsX4(channels)) {
+    // BUG, some PowerVRs (GE8320) produce incorrect result without it
+    op.compiler_options_.push_back(CompilerOptions::kClDisableOptimizations);
+  }
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HToY_DToZ;
+  return op;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/concat_z.h

@@ -0,0 +1,35 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONCAT_Z_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONCAT_Z_H_
+
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+GPUOperation CreateConcatZ(const OperationDef& definition,
+                           const std::vector<int>& channels,
+                           const GpuInfo& gpu_info);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONCAT_Z_H_

tensorflow/lite/delegates/gpu/common/tasks/conv_constants.cc

@@ -0,0 +1,296 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/conv_constants.h"
+
+#include <algorithm>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+
+namespace tflite {
+namespace gpu {
+
+namespace {
+// Adreno can provide up to ~3-4KB of constant memory, but in some cases even
+// 3KB can have very bad performance.
+int GetAdrenoOptimalMaxConstantSize(const AdrenoInfo& adreno_info) {
+  if (adreno_info.IsAdreno3xx() || adreno_info.IsAdreno4xx() ||
+      adreno_info.IsAdreno5xx()) {
+    return 256 * 10;  // 2.5KB
+  } else {
+    return 256 * 14;  // 3.5KB
+  }
+}
+
+int GetOptimalMaxConstantSize(const GpuInfo& gpu_info) {
+  if (gpu_info.IsAdreno()) {
+    return GetAdrenoOptimalMaxConstantSize(gpu_info.adreno_info);
+  } else if (gpu_info.IsAMD()) {
+    return 4096;
+  } else {
+    return 1024;  // 1KB
+  }
+}
+
+void AppendConditionally(const std::string& value, const std::string& delimeter,
+                         std::string* result) {
+  if (!result->empty()) {
+    *result += delimeter;
+  }
+  *result += value;
+}
+
+// src_size and dst_size must be <= 4;
+std::string GenerateConv(int src_size, int dst_size, bool use_dot_conv,
+                         int const_mem_offset, CalculationsPrecision precision,
+                         const std::string& dst, const std::string& src) {
+  std::string result;
+  const std::string postfixes[] = {".x", ".y", ".z", ".w"};
+  if (use_dot_conv) {
+    const std::string src_postfixes[] = {".x", ".xy", ".xyz", ""};
+    const std::string src_postfix = src_postfixes[src_size - 1];
+    for (int i = 0; i < dst_size; ++i) {
+      result += "      " + dst + postfixes[i] + " += dot(" + src +
+                ", args.weights.Read(" + std::to_string(const_mem_offset + i) +
+                ")" + src_postfix + ");\n";
+    }
+  } else {
+    const std::string dst_postfixes[] = {".x", ".xy", ".xyz", ""};
+    const std::string dst_postfix = dst_postfixes[dst_size - 1];
+    if (precision == CalculationsPrecision::F32_F16) {
+      for (int i = 0; i < src_size; ++i) {
+        if (i != 0) {
+          result += " + ";
+        }
+        std::string src_name = src;
+        if (src_size != 1) {
+          src_name += postfixes[i];
+        }
+        result += src_name + " * args.weights.Read(" +
+                  std::to_string(const_mem_offset + i) + ")" + dst_postfix;
+      }
+      std::string size = dst_size == 1 ? "" : std::to_string(dst_size);
+      result = "      " + dst + dst_postfix + " += TO_ACCUM_FLT" + size + "(" +
+               result + ");\n";
+    } else {
+      for (int i = 0; i < src_size; ++i) {
+        std::string src_name = src;
+        if (src_size != 1) {
+          src_name += postfixes[i];
+        }
+        result += "      " + dst + dst_postfix + " += " + src_name +
+                  " * args.weights.Read(" +
+                  std::to_string(const_mem_offset + i) + ")" + dst_postfix +
+                  ";\n";
+      }
+    }
+  }
+  return result;
+}
+
+std::string GenerateConvolutionConstantCode(const GpuInfo& gpu_info,
+                                            const OperationDef& op_def,
+                                            const OHWI& weights_shape,
+                                            bool x_oob_reads, bool y_oob_reads,
+                                            bool use_dot_conv,
+                                            GPUOperation* op) {
+  auto src_desc = op_def.src_tensors[0];
+  op->AddSrcTensor("src_tensor", src_desc);
+  op->AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
+
+  const int out_z = DivideRoundUp(weights_shape.o, 4);
+  const std::string kOutZ = std::to_string(out_z);
+  const int src_depth = DivideRoundUp(weights_shape.i, 4);
+
+  const std::string postfixes[] = {".x", ".xy", ".xyz", ""};
+
+  std::string c;
+  c += "MAIN_FUNCTION($0) {\n";
+  if (src_desc.HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  c += "  int Y = GLOBAL_ID_1;\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height()) "
+       "return;\n";
+  c += "  int start_x = X * args.stride_x + args.padding_x;\n";
+  c += "  int start_y = Y * args.stride_y + args.padding_y;\n";
+  for (int i = 0; i < out_z; ++i) {
+    c += "  ACCUM_FLT4 r" + std::to_string(i) + " = INIT_ACCUM_FLT4(0.0f);\n";
+  }
+  std::string check;
+  if (y_oob_reads && !src_desc.SupportsZeroClamp(Axis::HEIGHT, gpu_info)) {
+    AppendConditionally("inside_y", " && ", &check);
+  }
+  if (x_oob_reads && !src_desc.SupportsZeroClamp(Axis::WIDTH, gpu_info)) {
+    AppendConditionally("inside_x", " && ", &check);
+  }
+  int filters_counter = 0;
+  for (int s = 0; s < src_depth; ++s) {
+    const int src_ch_count = std::min(4, weights_shape.i - s * 4);
+    const std::string s_count =
+        src_ch_count == 1 ? "" : std::to_string(src_ch_count);
+    const std::string s_type = absl::StrCat("FLT", s_count);
+    const std::string s_postfix = postfixes[src_ch_count - 1];
+    for (int ky = 0; ky < weights_shape.h; ++ky) {
+      std::string s_y = absl::StrCat("(start_y + ", ky, " * args.dilation_y)");
+      c += "  {\n";
+      c += "    int y_c = start_y + " + std::to_string(ky) +
+           " * args.dilation_y;\n";
+      if (y_oob_reads && !src_desc.SupportsZeroClamp(Axis::HEIGHT, gpu_info)) {
+        c +=
+            "    bool inside_y = y_c >= 0 && y_c < args.src_tensor.Height();\n";
+        c += "    y_c = clamp(y_c, 0, args.src_tensor.Height() - 1);\n";
+      }
+      for (int kx = 0; kx < weights_shape.w; ++kx) {
+        c += "    {\n";
+        c += "      int x_c = start_x + " + std::to_string(kx) +
+             " * args.dilation_x;\n";
+        if (x_oob_reads && !src_desc.SupportsZeroClamp(Axis::WIDTH, gpu_info)) {
+          c += "      bool inside_x = x_c >= 0 && x_c < "
+               "args.src_tensor.Width();\n";
+          c += "      x_c = clamp(x_c, 0, args.src_tensor.Width() - 1);\n";
+        }
+        c += "      " + s_type + " src = args.src_tensor.Read(x_c, y_c, " +
+             std::to_string(s) + ")" + s_postfix + ";\n";
+        if (!check.empty()) {
+          c += "      src *= INIT_FLT(" + check + ");\n";
+        }
+        for (int d = 0; d < out_z; ++d) {
+          const int dst_ch_count = std::min(4, weights_shape.o - d * 4);
+          c += GenerateConv(src_ch_count, dst_ch_count, use_dot_conv,
+                            filters_counter, op_def.precision,
+                            "r" + std::to_string(d), "src");
+          filters_counter += use_dot_conv ? dst_ch_count : src_ch_count;
+        }
+        c += "    }\n";
+      }
+      c += "  }\n";
+    }
+  }
+  for (int i = 0; i < out_z; ++i) {
+    std::string s_i = std::to_string(i);
+    c += "  {\n";
+    c += "    FLT4 res = TO_FLT4(r" + s_i + ") + args.biases.Read(" + s_i +
+         ");\n";
+    c += "    args.dst_tensor.Write(res, X, Y, " + s_i + ");\n";
+    c += "  }\n";
+  }
+  c += "}\n";
+  return c;
+}
+
+bool IsDotConvBetter(int src_channels, int dst_channels) {
+  if (dst_channels % 4 == 0) {
+    return false;
+  }
+
+  // dst_channels % 4 != 0
+  if (src_channels % 4 == 0) {
+    return true;
+  }
+
+  // dst_channels % 4 != 0 && src_channels % 4 != 0
+  const int src_depth = DivideRoundUp(src_channels, 4);
+  const int dst_depth = DivideRoundUp(dst_channels, 4);
+  return dst_channels * src_depth < src_channels * dst_depth;
+}
+
+}  // namespace
+
+bool IsConvConstantsSupported(const GpuInfo& gpu_info,
+                              const OperationDef& definition,
+                              const Convolution2DAttributes& attr) {
+  if (gpu_info.IsApiOpenCl() && gpu_info.IsAdreno()) {
+    const std::string kBadDriver =
+        "OpenCL 2.0 QUALCOMM build: commit #7ff4f54 changeid #I4460aa6217 "
+        "Date: 12/30/18";
+    if (absl::StrContains(gpu_info.opencl_info.platform_version, kBadDriver)) {
+      return false;
+    }
+  }
+
+  if (attr.groups != 1) {
+    return false;
+  }
+
+  const bool use_dot_conv =
+      IsDotConvBetter(attr.weights.shape.i, attr.weights.shape.o);
+  const auto& w_shape = attr.weights.shape;
+  const int src_depth = DivideRoundUp(w_shape.i, 4);
+  const int dst_depth = DivideRoundUp(w_shape.o, 4);
+  const int aligned_ch_count =
+      use_dot_conv ? w_shape.o * src_depth * 4 : w_shape.i * dst_depth * 4;
+  const int filters_count = aligned_ch_count * w_shape.h * w_shape.w;
+  const int float_size = definition.precision == CalculationsPrecision::F32
+                             ? sizeof(float)
+                             : sizeof(half);
+  const int filters_buffer_size = filters_count * float_size;
+  const int kConstantMaxSize = GetOptimalMaxConstantSize(gpu_info);
+  const int flt4_registers = DivideRoundUp(w_shape.o, 4);
+  return filters_buffer_size <= kConstantMaxSize && flt4_registers <= 8;
+}
+
+GPUOperation CreateConvConstants(const GpuInfo& gpu_info,
+                                 const OperationDef& definition,
+                                 const Convolution2DAttributes& attr) {
+  const bool use_dot_conv =
+      IsDotConvBetter(attr.weights.shape.i, attr.weights.shape.o);
+  GPUOperation op(definition);
+  UploadWeightsForConvConstants(attr.weights, gpu_info, definition.precision,
+                                use_dot_conv, &op);
+  op.args_.AddInt("stride_x", attr.strides.w);
+  op.args_.AddInt("stride_y", attr.strides.h);
+  op.args_.AddInt("padding_x", -attr.padding.prepended.w);
+  op.args_.AddInt("padding_y", -attr.padding.prepended.h);
+  op.args_.AddInt("dilation_x", attr.dilations.w);
+  op.args_.AddInt("dilation_y", attr.dilations.h);
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_ZIs1;
+
+  bool x_oob_reads =
+      attr.padding.appended.w != 0 || attr.padding.prepended.w != 0;
+  bool y_oob_reads =
+      attr.padding.appended.h != 0 || attr.padding.prepended.h != 0;
+  op.code_ = GenerateConvolutionConstantCode(gpu_info, definition,
+                                             attr.weights.shape, x_oob_reads,
+                                             y_oob_reads, use_dot_conv, &op);
+  if (definition.precision == CalculationsPrecision::F16 &&
+      gpu_info.IsAdreno() && gpu_info.adreno_info.IsAdreno3xx()) {
+    op.compiler_options_.push_back(CompilerOptions::kAdrenoFullSimd);
+  }
+  if (definition.precision != CalculationsPrecision::F32 &&
+      gpu_info.IsPowerVR()) {
+    // BUG, some PowerVRs (GE8320) produce incorrect result without it
+    op.compiler_options_.push_back(CompilerOptions::kClDisableOptimizations);
+  }
+
+  TensorDescriptor bias_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), attr.bias);
+  op.args_.AddObject("biases", std::make_unique<TensorDescriptor>(
+                                   std::move(bias_tensor_desc)));
+  return op;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/conv_constants.h

@@ -0,0 +1,170 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_CONSTANTS_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_CONSTANTS_H_
+
+#include <memory>
+#include <utility>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/buffer_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/tensor.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+template <DataType S, typename T>
+void RearrangeWeightsForConvConstants(
+    const tflite::gpu::Tensor<OHWI, S>& weights, absl::Span<T> dst) {
+  const int dst_depth = DivideRoundUp(weights.shape.o, 4);
+  const int src_depth = DivideRoundUp(weights.shape.i, 4);
+  const int kernel_x = weights.shape.w;
+  const int kernel_y = weights.shape.h;
+
+  int counter = 0;
+  for (int s = 0; s < src_depth; ++s) {
+    for (int y = 0; y < kernel_y; ++y) {
+      for (int x = 0; x < kernel_x; ++x) {
+        for (int d = 0; d < dst_depth; ++d) {
+          const int channels_count = std::min(4, weights.shape.i - s * 4);
+          T filters[4];
+          for (int i = 0; i < 4; ++i) {
+            for (int j = 0; j < channels_count; ++j) {
+              const int s_ch = s * 4 + j;
+              const int d_ch = d * 4 + i;
+              if (s_ch < weights.shape.i && d_ch < weights.shape.o) {
+                const int f_index =
+                    weights.shape.LinearIndex({d_ch, y, x, s_ch});
+                filters[j][i] = weights.data[f_index];
+              } else {
+                filters[j][i] = 0.0f;
+              }
+            }
+          }
+          for (int i = 0; i < channels_count; ++i) {
+            dst[counter++] = filters[i];
+          }
+        }
+      }
+    }
+  }
+}
+
+template <DataType S, typename T>
+void RearrangeWeightsForConvConstantsDot(
+    const tflite::gpu::Tensor<OHWI, S>& weights, absl::Span<T> dst) {
+  const int dst_depth = DivideRoundUp(weights.shape.o, 4);
+  const int src_depth = DivideRoundUp(weights.shape.i, 4);
+  const int kernel_x = weights.shape.w;
+  const int kernel_y = weights.shape.h;
+
+  int counter = 0;
+  for (int s = 0; s < src_depth; ++s) {
+    for (int y = 0; y < kernel_y; ++y) {
+      for (int x = 0; x < kernel_x; ++x) {
+        for (int d = 0; d < dst_depth; ++d) {
+          const int channels_count = std::min(4, weights.shape.o - d * 4);
+          T filters[4];
+          for (int j = 0; j < channels_count; ++j) {
+            for (int i = 0; i < 4; ++i) {
+              const int s_ch = s * 4 + i;
+              const int d_ch = d * 4 + j;
+              if (s_ch < weights.shape.i && d_ch < weights.shape.o) {
+                const int f_index =
+                    weights.shape.LinearIndex({d_ch, y, x, s_ch});
+                filters[j][i] = weights.data[f_index];
+              } else {
+                filters[j][i] = 0.0f;
+              }
+            }
+          }
+          for (int i = 0; i < channels_count; ++i) {
+            dst[counter++] = filters[i];
+          }
+        }
+      }
+    }
+  }
+}
+
+template <DataType T>
+void UploadWeightsForConvConstants(const tflite::gpu::Tensor<OHWI, T>& weights,
+                                   const GpuInfo& gpu_info,
+                                   CalculationsPrecision precision,
+                                   bool use_dot_conv, GPUOperation* op) {
+  const int src_depth = DivideRoundUp(weights.shape.i, 4);
+  const int dst_depth = DivideRoundUp(weights.shape.o, 4);
+  const int kernel_x = weights.shape.w;
+  const int kernel_y = weights.shape.h;
+
+  const bool f32_weights = precision == CalculationsPrecision::F32;
+  const int float_size = f32_weights ? 4 : 2;
+  const int aligned_ch_count = use_dot_conv ? weights.shape.o * src_depth * 4
+                                            : weights.shape.i * dst_depth * 4;
+  const int float_count = aligned_ch_count * kernel_x * kernel_y;
+
+  BufferDescriptor desc;
+  desc.element_type = f32_weights ? DataType::FLOAT32 : DataType::FLOAT16;
+  desc.element_size = 4;
+  if (gpu_info.IsApiOpenCl() || gpu_info.IsApiMetal()) {
+    desc.memory_type = MemoryType::CONSTANT;
+  } else {
+    desc.memory_type = MemoryType::GLOBAL;
+  }
+  desc.size = float_size * float_count;
+  desc.data.resize(desc.size);
+
+  if (f32_weights) {
+    float4* ptr = reinterpret_cast<float4*>(desc.data.data());
+    if (use_dot_conv) {
+      RearrangeWeightsForConvConstantsDot(weights,
+                                          absl::MakeSpan(ptr, float_count / 4));
+    } else {
+      RearrangeWeightsForConvConstants(weights,
+                                       absl::MakeSpan(ptr, float_count / 4));
+    }
+  } else {
+    half4* ptr = reinterpret_cast<half4*>(desc.data.data());
+    if (use_dot_conv) {
+      RearrangeWeightsForConvConstantsDot(weights,
+                                          absl::MakeSpan(ptr, float_count / 4));
+    } else {
+      RearrangeWeightsForConvConstants(weights,
+                                       absl::MakeSpan(ptr, float_count / 4));
+    }
+  }
+
+  op->args_.AddObject("weights",
+                      std::make_unique<BufferDescriptor>(std::move(desc)));
+}
+
+bool IsConvConstantsSupported(const GpuInfo& gpu_info,
+                              const OperationDef& definition,
+                              const Convolution2DAttributes& attr);
+
+GPUOperation CreateConvConstants(const GpuInfo& gpu_info,
+                                 const OperationDef& definition,
+                                 const Convolution2DAttributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_CONSTANTS_H_

tensorflow/lite/delegates/gpu/common/tasks/conv_constants_test_util.cc

@@ -0,0 +1,104 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/conv_constants_test_util.h"
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/conv_constants.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConvConstantsSimpleWeightsTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  Convolution2DAttributes attr;
+  attr.padding.prepended = HW(0, 0);
+  attr.padding.appended = HW(1, 1);
+  attr.strides = HW(1, 1);
+  attr.dilations = HW(1, 1);
+  attr.weights.shape = OHWI(1, 2, 2, 2);
+  attr.weights.data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f};
+  attr.bias.shape = Linear(1);
+  attr.bias.data = {0.0f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation =
+          CreateConvConstants(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 2, 2, 1), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({28.0f, 18.0f, 22.0f, 13.0f}, dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConvConstantsTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  Convolution2DAttributes attr;
+  attr.padding.prepended = HW(0, 0);
+  attr.padding.appended = HW(1, 1);
+  attr.strides = HW(1, 1);
+  attr.dilations = HW(1, 1);
+  attr.weights.shape = OHWI(2, 2, 2, 2);
+  attr.weights.data = {1.0f, 2.0f,  3.0f,  4.0f,  5.0f,  6.0f,  7.0f,  8.0f,
+                       9.0f, 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f, 16.0f};
+  attr.bias.shape = Linear(2);
+  attr.bias.data = {0.5f, -0.5f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation =
+          CreateConvConstants(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 2, 2, 2), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear(
+          {168.5f, 391.5f, 80.5f, 223.5f, 60.5f, 235.5f, 20.5f, 123.5f},
+          dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/conv_constants_test_util.h

@@ -0,0 +1,31 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_CONSTANTS_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_CONSTANTS_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConvConstantsSimpleWeightsTest(TestExecutionEnvironment* env);
+absl::Status ConvConstantsTest(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_CONSTANTS_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/conv_generic.h

@@ -0,0 +1,359 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_GENERIC_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_GENERIC_H_
+
+#include <cstring>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/buffer_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tensor_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/weights_conversion.h"
+#include "tensorflow/lite/delegates/gpu/common/task/weights_layout.h"
+#include "tensorflow/lite/delegates/gpu/common/tensor.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+#include "tensorflow/lite/delegates/gpu/common/winograd_util.h"
+
+namespace tflite {
+namespace gpu {
+
+class ConvGeneric : public GPUOperation {
+ public:
+  enum class WeightsUploadType {
+    LOCAL_MEM_ASYNC_SUBGROUP,  // we use it for PowerVR with workgroup size = 32
+    LOCAL_MEM_BY_THREADS,
+    GLOBAL_MEM,
+    CONSTANT_MEM,
+    PRIVATE_MEM_SIMD_BROADCAST,
+    TEXTURES_MEM_X4,  // 4 textures for weights
+  };
+  struct ConvParams {
+    DataType weights_data_type;  // used for weights and biases
+    int4 block_size;             // WHDS
+    bool fixed_work_group_size;
+    int3 work_group_size;
+    int3 work_group_launch_order;
+    bool linear_spatial;  // spatial dimensions are Width/Height/Depth
+    bool linear_all;  // linear_spatial & linear_all can not be used together,
+                      // linear_all can not be used with WeightsUploadTypes
+                      // that use workgroups(subgroups) for
+                      // uploading(LOCAL_MEM_BY_THREADS for example).
+    bool different_weights_for_height;
+    bool groups_support = false;  // convolution groups
+    int src_depth_loop_size;
+    bool need_src_loop = true;
+    bool need_dst_loop = true;
+    WeightsUploadType weights_upload_type;
+    bool x_kernel_is_1 = false;
+    bool y_kernel_is_1 = false;
+    bool z_kernel_is_1 = false;
+    WeightsLayout weights_layout;
+
+    // used only with PRIVATE_MEM_SIMD_BROADCAST
+    int simd_size = 1;
+
+    bool AreWeightsBuffer() const {
+      return weights_upload_type != WeightsUploadType::TEXTURES_MEM_X4;
+    }
+
+    bool IsPrivateMemBroadcast() const {
+      return weights_upload_type ==
+             WeightsUploadType::PRIVATE_MEM_SIMD_BROADCAST;
+    }
+  };
+  ConvGeneric() = default;
+  void GetPossibleKernelWorkGroups(
+      TuningType tuning_type, const GpuInfo& gpu_info,
+      const KernelInfo& kernel_info,
+      std::vector<int3>* work_groups) const override;
+  absl::Status BindArguments(ArgumentsBinder* args) override;
+  int3 GetGridSize() const override;
+
+  WeightsDescription GetWeightsDescription() const {
+    WeightsDescription desc;
+    desc.type = conv_params_.weights_data_type;
+    desc.layout = conv_params_.weights_layout;
+    desc.output_group_size = conv_params_.block_size.w;
+    return desc;
+  }
+
+  // Move only
+  ConvGeneric(ConvGeneric&& operation);
+  ConvGeneric& operator=(ConvGeneric&& operation);
+  ConvGeneric(const ConvGeneric&) = delete;
+  ConvGeneric& operator=(const ConvGeneric&) = delete;
+
+ private:
+  ConvGeneric(const OperationDef& definition,
+              const Convolution2DAttributes& attr, const GpuInfo& gpu_info,
+              const BHWC* dst_shape = nullptr);
+  ConvGeneric(const OperationDef& definition,
+              const Convolution2DAttributes& attr, const BHWC& weights_shape,
+              const GpuInfo& gpu_info, const BHWC* dst_shape = nullptr);
+  ConvGeneric(const OperationDef& definition,
+              const FullyConnectedAttributes& attr, const GpuInfo& gpu_info,
+              const BHWC* dst_shape = nullptr);
+  explicit ConvGeneric(const OperationDef& definition);
+  ConvGeneric(const OperationDef& definition,
+              const Convolution3DAttributes& attr, const GpuInfo& gpu_info,
+              const BHWDC* dst_shape = nullptr);
+
+  void GenerateCode(const GpuInfo& gpu_info);
+
+  template <DataType T>
+  void UploadData(const tflite::gpu::Tensor<OHWI, T>& weights,
+                  const tflite::gpu::Tensor<Linear, T>& biases);
+  template <DataType T>
+  void UploadDataForWinograd4x4To6x6(
+      const tflite::gpu::Tensor<OHWI, T>& weights);
+
+  template <DataType T>
+  void UploadWeights(const tflite::gpu::Tensor<OHWI, T>& weights);
+
+  template <DataType T>
+  void UploadWeights(const tflite::gpu::Tensor<OHWDI, T>& weights);
+
+  template <DataType T>
+  void UploadBias(const tflite::gpu::Tensor<Linear, T>& bias);
+
+  friend ConvGeneric CreateConvGeneric(const GpuInfo& gpu_info,
+                                       const OperationDef& definition,
+                                       const Convolution2DAttributes& attr,
+                                       const BHWC* dst_shape);
+
+  friend ConvGeneric CreateConvGeneric(const GpuInfo& gpu_info,
+                                       const OperationDef& definition,
+                                       const FullyConnectedAttributes& attr,
+                                       const BHWC* dst_shape);
+
+  friend ConvGeneric CreateConvGenericBatchedMatMul(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const OHWI& weights_shape, const BHWC* dst_shape);
+
+  friend ConvGeneric CreateConvGenericDynamicWeights(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const Convolution2DAttributes& attr, const BHWC& weights_shape,
+      const BHWC* dst_shape);
+
+  friend ConvGeneric CreateConvGenericWino4x4To6x6(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const Convolution2DAttributes& attr, const BHWC* dst_shape);
+
+  friend ConvGeneric CreateConvGeneric3D(const GpuInfo& gpu_info,
+                                         const OperationDef& definition,
+                                         const Convolution3DAttributes& attr,
+                                         const BHWDC* dst_shape);
+
+  ConvParams GuessBestParams(const GpuInfo& gpu_info,
+                             const OperationDef& definition,
+                             const Convolution2DAttributes& attr,
+                             const BHWC* dst_shape = nullptr);
+  ConvParams GuessBestParams(const GpuInfo& gpu_info,
+                             const OperationDef& definition,
+                             const Convolution2DAttributes& attr,
+                             const BHWC& weights_shape,
+                             const BHWC* dst_shape = nullptr);
+  ConvParams GuessBestParams(const GpuInfo& gpu_info,
+                             const OperationDef& definition,
+                             const FullyConnectedAttributes& attr,
+                             const BHWC* dst_shape = nullptr);
+  ConvParams GuessBestParamsPointwise(const GpuInfo& gpu_info,
+                                      const OperationDef& definition,
+                                      const OHWI& weights_shape,
+                                      const BHWC* dst_shape = nullptr);
+  ConvParams GuessBestParams(const GpuInfo& gpu_info,
+                             const OperationDef& definition,
+                             const Convolution3DAttributes& attr,
+                             const BHWDC* dst_shape = nullptr);
+  ConvParams GuessBestParams(const GpuInfo& gpu_info,
+                             const OperationDef& definition, int src_depth,
+                             int dst_depth, bool x_kernel_is_1,
+                             bool y_kernel_is_1,
+                             bool different_weights_for_height,
+                             const BHWC* dst_shape = nullptr);
+  ConvParams GuessBestParamsApple(const GpuInfo& gpu_info,
+                                  const OperationDef& definition, int src_depth,
+                                  int dst_depth, bool x_kernel_is_1,
+                                  bool y_kernel_is_1,
+                                  bool different_weights_for_height,
+                                  const BHWC& dst_shape);
+
+  std::string GenerateConv(const GpuInfo& gpu_info, const OperationDef& op_def,
+                           const ConvParams& conv_params);
+
+  int4 stride_;
+  int4 padding_;
+  int4 kernel_size_;
+  int4 dilation_;
+  ConvParams conv_params_;
+};
+
+template <DataType T>
+void ConvGeneric::UploadData(const tflite::gpu::Tensor<OHWI, T>& weights,
+                             const tflite::gpu::Tensor<Linear, T>& biases) {
+  UploadWeights(weights);
+  UploadBias(biases);
+}
+
+template <DataType T>
+void ConvGeneric::UploadDataForWinograd4x4To6x6(
+    const tflite::gpu::Tensor<OHWI, T>& weights) {
+  tflite::gpu::Tensor<OHWI, T> wino_weights;
+  RearrangeWeightsToWinograd4x4To6x6Weights(weights, &wino_weights);
+  UploadWeights(wino_weights);
+  tflite::gpu::Tensor<Linear, DataType::FLOAT32> biases;
+  biases.shape = Linear(weights.shape.o);
+  biases.data.resize(weights.shape.o, 0.0f);
+  UploadBias(biases);
+}
+
+template <DataType T>
+void ConvGeneric::UploadBias(const tflite::gpu::Tensor<Linear, T>& bias) {
+  BufferDescriptor desc;
+  desc.element_type = conv_params_.weights_data_type;
+  desc.element_size = 4;
+  desc.memory_type = conv_params_.weights_upload_type ==
+                             ConvGeneric::WeightsUploadType::CONSTANT_MEM
+                         ? MemoryType::CONSTANT
+                         : MemoryType::GLOBAL;
+  const int float_size = conv_params_.weights_data_type == DataType::FLOAT32
+                             ? sizeof(float)
+                             : sizeof(half);
+  int aligned_channels = AlignByN(bias.shape.v, 4 * conv_params_.block_size.w);
+  desc.size = float_size * aligned_channels;
+  desc.data.resize(desc.size);
+  if (conv_params_.weights_data_type == DataType::FLOAT32) {
+    float* gpu_data = reinterpret_cast<float*>(desc.data.data());
+    for (int i = 0; i < aligned_channels; ++i) {
+      gpu_data[i] = i < bias.shape.v ? bias.data[i] : 0.0f;
+    }
+  } else {
+    half* gpu_data = reinterpret_cast<half*>(desc.data.data());
+    for (int i = 0; i < aligned_channels; ++i) {
+      gpu_data[i] = i < bias.shape.v ? bias.data[i] : 0.0f;
+    }
+  }
+  args_.AddObject("biases",
+                  std::make_unique<BufferDescriptor>(std::move(desc)));
+}
+
+template <DataType T>
+void ConvGeneric::UploadWeights(const tflite::gpu::Tensor<OHWI, T>& weights) {
+  const auto weights_desc = GetWeightsDescription();
+  const int flt_count =
+      GetTotalElementsCountForLayout(weights_desc, weights.shape);
+
+  std::vector<uint8_t> weights_data(flt_count * SizeOf(weights_desc.type));
+  RearrangeWeights(weights, weights_desc, absl::MakeSpan(weights_data));
+
+  if (conv_params_.AreWeightsBuffer()) {
+    BufferDescriptor desc;
+    desc.element_type = weights_desc.type;
+    desc.element_size = 4;
+    desc.memory_type = conv_params_.weights_upload_type ==
+                               ConvGeneric::WeightsUploadType::CONSTANT_MEM
+                           ? MemoryType::CONSTANT
+                           : MemoryType::GLOBAL;
+    desc.size = weights_data.size();
+    desc.data = std::move(weights_data);
+    args_.AddObject("weights",
+                    std::make_unique<BufferDescriptor>(std::move(desc)));
+  } else {
+    uint2 tex_size = Get2dResourceSize(weights_desc, weights.shape);
+    int sub_size = SizeOf(weights_desc.type) * 4 * tex_size.x * tex_size.y;
+    for (int i = 0; i < 4; ++i) {
+      TensorDescriptor desc = CreateConstantHWVec4TensorDescriptor(
+          weights_desc.type, TensorStorageType::TEXTURE_2D, tex_size.x,
+          tex_size.y, weights_data.data() + sub_size * i);
+      args_.AddObject("weights" + std::to_string(i),
+                      std::make_unique<TensorDescriptor>(std::move(desc)));
+    }
+  }
+}
+
+template <DataType T>
+void ConvGeneric::UploadWeights(const tflite::gpu::Tensor<OHWDI, T>& weights) {
+  const auto weights_desc = GetWeightsDescription();
+  const int flt_count =
+      GetTotalElementsCountForLayout(weights_desc, weights.shape);
+
+  std::vector<uint8_t> weights_data(flt_count * SizeOf(weights_desc.type));
+  RearrangeWeights(weights, weights_desc, absl::MakeSpan(weights_data));
+
+  if (conv_params_.AreWeightsBuffer()) {
+    BufferDescriptor desc;
+    desc.element_type = weights_desc.type;
+    desc.element_size = 4;
+    desc.size = weights_data.size();
+    desc.data = std::move(weights_data);
+    args_.AddObject("weights",
+                    std::make_unique<BufferDescriptor>(std::move(desc)));
+  } else {
+    uint2 tex_size = Get2dResourceSize(weights_desc, weights.shape);
+    int sub_size = SizeOf(weights_desc.type) * 4 * tex_size.x * tex_size.y;
+    for (int i = 0; i < 4; ++i) {
+      TensorDescriptor desc = CreateConstantHWVec4TensorDescriptor(
+          weights_desc.type, TensorStorageType::TEXTURE_2D, tex_size.x,
+          tex_size.y, weights_data.data() + sub_size * i);
+      args_.AddObject("weights" + std::to_string(i),
+                      std::make_unique<TensorDescriptor>(std::move(desc)));
+    }
+  }
+}
+
+ConvGeneric CreateConvGeneric(const GpuInfo& gpu_info,
+                              const OperationDef& definition,
+                              const Convolution2DAttributes& attr,
+                              const BHWC* dst_shape = nullptr);
+
+ConvGeneric CreateConvGeneric(const GpuInfo& gpu_info,
+                              const OperationDef& definition,
+                              const FullyConnectedAttributes& attr,
+                              const BHWC* dst_shape = nullptr);
+
+ConvGeneric CreateConvGenericDynamicWeights(const GpuInfo& gpu_info,
+                                            const OperationDef& definition,
+                                            const Convolution2DAttributes& attr,
+                                            const BHWC& weights_shape,
+                                            const BHWC* dst_shape = nullptr);
+
+ConvGeneric CreateConvGenericBatchedMatMul(const GpuInfo& gpu_info,
+                                           const OperationDef& definition,
+                                           const OHWI& weights_shape,
+                                           const BHWC* dst_shape = nullptr);
+
+ConvGeneric CreateConvGenericWino4x4To6x6(const GpuInfo& gpu_info,
+                                          const OperationDef& definition,
+                                          const Convolution2DAttributes& attr,
+                                          const BHWC* dst_shape = nullptr);
+
+ConvGeneric CreateConvGeneric3D(const GpuInfo& gpu_info,
+                                const OperationDef& definition,
+                                const Convolution3DAttributes& attr,
+                                const BHWDC* dst_shape = nullptr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_GENERIC_H_

tensorflow/lite/delegates/gpu/common/tasks/conv_generic_test_util.cc

@@ -0,0 +1,219 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/conv_generic_test_util.h"
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/conv_generic.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConvGeneric1x1SimpleWeightsTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  Convolution2DAttributes attr;
+  attr.padding.prepended = HW(0, 0);
+  attr.padding.appended = HW(0, 0);
+  attr.strides = HW(1, 1);
+  attr.dilations = HW(1, 1);
+  attr.weights.shape = OHWI(2, 1, 1, 2);
+  attr.weights.data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f};
+  attr.bias.shape = Linear(1);
+  attr.bias.data = {0.0f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      ConvGeneric operation =
+          CreateConvGeneric(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<ConvGeneric>(std::move(operation)),
+          BHWC(1, 2, 2, 2), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({1.0f, 1.0f, 5.0f, 5.0f, 9.0f, 9.0f, 13.0f, 13.0f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConvGeneric1x1Test(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  Convolution2DAttributes attr;
+  attr.padding.prepended = HW(0, 0);
+  attr.padding.appended = HW(0, 0);
+  attr.strides = HW(1, 1);
+  attr.dilations = HW(1, 1);
+  attr.weights.shape = OHWI(2, 1, 1, 2);
+  attr.weights.data = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f};
+  attr.bias.shape = Linear(2);
+  attr.bias.data = {0.5f, -0.5f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      ConvGeneric operation =
+          CreateConvGeneric(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<ConvGeneric>(std::move(operation)),
+          BHWC(1, 2, 2, 2), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({2.5f, 3.5f, 8.5f, 17.5f, 14.5f, 31.5f, 20.5f, 45.5f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConvGenericSimpleWeightsTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  Convolution2DAttributes attr;
+  attr.padding.prepended = HW(0, 0);
+  attr.padding.appended = HW(1, 1);
+  attr.strides = HW(1, 1);
+  attr.dilations = HW(1, 1);
+  attr.weights.shape = OHWI(1, 2, 2, 2);
+  attr.weights.data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f};
+  attr.bias.shape = Linear(1);
+  attr.bias.data = {0.0f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      ConvGeneric operation =
+          CreateConvGeneric(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<ConvGeneric>(std::move(operation)),
+          BHWC(1, 2, 2, 1), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({28.0f, 18.0f, 22.0f, 13.0f}, dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConvGenericTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  Convolution2DAttributes attr;
+  attr.padding.prepended = HW(0, 0);
+  attr.padding.appended = HW(1, 1);
+  attr.strides = HW(1, 1);
+  attr.dilations = HW(1, 1);
+  attr.weights.shape = OHWI(2, 2, 2, 2);
+  attr.weights.data = {1.0f, 2.0f,  3.0f,  4.0f,  5.0f,  6.0f,  7.0f,  8.0f,
+                       9.0f, 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f, 16.0f};
+  attr.bias.shape = Linear(2);
+  attr.bias.data = {0.5f, -0.5f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      ConvGeneric operation =
+          CreateConvGeneric(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<ConvGeneric>(std::move(operation)),
+          BHWC(1, 2, 2, 2), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear(
+          {168.5f, 391.5f, 80.5f, 223.5f, 60.5f, 235.5f, 20.5f, 123.5f},
+          dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConvGenericGroupedTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 1, 1, 8);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  Convolution2DAttributes attr;
+  attr.groups = 2;
+  attr.padding.prepended = HW(0, 0);
+  attr.padding.appended = HW(0, 0);
+  attr.strides = HW(1, 1);
+  attr.dilations = HW(1, 1);
+  attr.weights.shape = OHWI(8, 1, 1, 4);
+  attr.weights.data = {1.0f,  2.0f,  3.0f,  4.0f,  5.0f,  6.0f,  7.0f,  8.0f,
+                       9.0f,  10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f, 16.0f,
+                       17.0f, 18.0f, 19.0f, 20.0f, 21.0f, 22.0f, 23.0f, 24.0f,
+                       25.0f, 26.0f, 27.0f, 28.0f, 29.0f, 30.0f, 31.0f, 32.0f};
+  attr.bias.shape = Linear(8);
+  attr.bias.data = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      ConvGeneric operation =
+          CreateConvGeneric(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<ConvGeneric>(std::move(operation)),
+          BHWC(1, 1, 1, 8), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear(
+          {20.0f, 44.0f, 68.0f, 92.0f, 412.0f, 500.0f, 588.0f, 676.0f},
+          dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/conv_generic_test_util.h

@@ -0,0 +1,34 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_GENERIC_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_GENERIC_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConvGeneric1x1SimpleWeightsTest(TestExecutionEnvironment* env);
+absl::Status ConvGeneric1x1Test(TestExecutionEnvironment* env);
+absl::Status ConvGenericSimpleWeightsTest(TestExecutionEnvironment* env);
+absl::Status ConvGenericTest(TestExecutionEnvironment* env);
+absl::Status ConvGenericGroupedTest(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_GENERIC_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/conv_metal_simd.cc

@@ -0,0 +1,560 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/conv_metal_simd.h"
+
+#include <cmath>
+#include <cstdint>
+#include <functional>
+#include <map>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/gpu_info.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+#include "tensorflow/lite/delegates/gpu/common/util.h"
+
+namespace tflite {
+namespace gpu {
+namespace {
+std::string GenerateDstCoords(const int3& work_group_launch_order,
+                              bool linear_spatial, bool need_depth,
+                              bool need_batch) {
+  std::string c;
+  int3 launch_remap;
+  launch_remap[work_group_launch_order.x] = 0;
+  launch_remap[work_group_launch_order.y] = 1;
+  launch_remap[work_group_launch_order.z] = 2;
+  if (linear_spatial) {
+    if (work_group_launch_order[0] == 0) {
+      c += "  int linear_spatial = GLOBAL_ID_0;\n";
+    } else {
+      c += "  int linear_spatial = GROUP_ID_" +
+           std::to_string(launch_remap[0]) + " * GROUP_SIZE_0 + LOCAL_ID_0;\n";
+    }
+    if (need_batch) {
+      c += "  int B = linear_spatial % args.dst_tensor.Batch();\n";
+      c += "  linear_spatial = linear_spatial / args.dst_tensor.Batch();\n";
+    }
+    if (need_depth) {
+      c += "  int DST_X = linear_spatial % args.dst_tensor.Width();\n";
+      c += "  linear_spatial = linear_spatial / args.dst_tensor.Width();\n";
+      c += "  int DST_Y = linear_spatial % args.dst_tensor.Height();\n";
+      c += "  int DST_Z = linear_spatial / args.dst_tensor.Height();\n";
+    } else {
+      c += "  int DST_Y = linear_spatial / args.dst_tensor.Width();\n";
+      c += "  int DST_X = linear_spatial % args.dst_tensor.Width();\n";
+    }
+    if (work_group_launch_order[1] == 1) {
+      c += "  int DST_S = GLOBAL_ID_1;\n";
+    } else {
+      c += "  int DST_S = GROUP_ID_" + std::to_string(launch_remap[1]) +
+           " * GROUP_SIZE_1 + LOCAL_ID_1;\n";
+    }
+  } else {
+    if (work_group_launch_order[0] == 0) {
+      c += "  int DST_X = GLOBAL_ID_0;\n";
+    } else {
+      c += "  int DST_X = GROUP_ID_" + std::to_string(launch_remap[0]) +
+           " * GROUP_SIZE_0 + LOCAL_ID_0;\n";
+    }
+    if (need_batch) {
+      c += "  int B = DST_X % args.dst_tensor.Batch();\n";
+      c += "  DST_X = DST_X / args.dst_tensor.Batch();\n";
+    }
+    std::string global_id_1;
+    if (work_group_launch_order[1] == 1) {
+      global_id_1 = "GLOBAL_ID_1";
+    } else {
+      global_id_1 = "GROUP_ID_" + std::to_string(launch_remap[1]) +
+                    " * GROUP_SIZE_1 + LOCAL_ID_1";
+    }
+    if (need_depth) {
+      c += "  int linear_id_1 = " + global_id_1 + ";\n";
+      c += "  int DST_Z = linear_id_1 / dst_tensor.Height();\n";
+      c += "  int DST_Y = linear_id_1 % dst_tensor.Height();\n";
+    } else {
+      c += "  int DST_Y = " + global_id_1 + ";\n";
+    }
+    if (work_group_launch_order[2] == 2) {
+      c += "  int DST_S = GLOBAL_ID_2;\n";
+    } else {
+      c += "  int DST_S = GROUP_ID_" + std::to_string(launch_remap[2]) +
+           " * GROUP_SIZE_2 + LOCAL_ID_2;\n";
+    }
+  }
+
+  return c;
+}
+
+std::string GenerateConvolution(
+    const OperationDef& definition,
+    const ConvolutionMetalSimd::ConvParams& conv_params) {
+  std::string c;
+  c += "#define MMA simdgroup_multiply_accumulate\n";
+  const int spatial_threads = conv_params.GetSpatialThreadsCount();
+  c += "#define SPATIAL_THREADS " + std::to_string(spatial_threads) + "\n";
+  c += "MAIN_FUNCTION($0) {\n";
+  c += GenerateDstCoords(conv_params.work_group_launch_order,
+                         conv_params.linear_spatial,
+                         definition.src_tensors[0].HasAxis(Axis::DEPTH),
+                         definition.src_tensors[0].HasAxis(Axis::BATCH));
+  if (definition.src_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+  }
+  if (conv_params.slices_per_thread != 1) {
+    c += "  DST_S *= " + std::to_string(conv_params.slices_per_thread) + ";\n";
+  }
+  c += "  device FLT4* f_offseted = args.weights.GetPtr() + DST_S * 4 * "
+       "args.src_tensor.Slices();\n";
+  const bool cache_weight = spatial_threads > 32;
+  const int src_x4_slices = conv_params.GetX4SlicesCount();
+  const int src_x8_slices = src_x4_slices / 2;
+  const int dst_x4_slices = conv_params.slices_per_thread;
+  const int dst_x8_slices = dst_x4_slices / 2;
+  const int weights_tiles8x8_per_spatial = src_x8_slices * dst_x8_slices;
+  const int weights_flt4_per_spatial = weights_tiles8x8_per_spatial * 8 * 8 / 4;
+  if (conv_params.linear_spatial) {
+    c += "  int spatial_id = LOCAL_ID_0;\n";
+    c += "  int slice_id = LOCAL_ID_1;\n";
+  } else {
+    c += "  int spatial_id = LOCAL_ID_1 * GROUP_SIZE_0 + LOCAL_ID_0;\n";
+    c += "  int slice_id = LOCAL_ID_2;\n";
+  }
+  c += "  int tid = slice_id * SPATIAL_THREADS + spatial_id;\n";
+  if (cache_weight) {
+    c += "  threadgroup FLT4 tmp_w[" +
+         std::to_string(weights_flt4_per_spatial *
+                        conv_params.GetX4SlicesCount()) +
+         "];\n";
+    c += "  threadgroup FLT* tmp_w_x1 = (threadgroup FLT*)tmp_w;\n";
+    c += "  tmp_w_x1 += " + std::to_string(weights_flt4_per_spatial * 4) +
+         " * slice_id;\n";
+    c += "  threadgroup FLT4* tmp_w_x4 = (threadgroup FLT4*)tmp_w_x1;\n\n";
+  } else {
+    c += "  device FLT* f_offseted_x1 = (device FLT*)f_offseted;\n\n";
+  }
+  c += "  threadgroup FLT4 tmp_src[SPATIAL_THREADS * " +
+       std::to_string(src_x4_slices) + "];\n";
+  c += "  threadgroup FLT* tmp_src_x1 = (threadgroup FLT*)tmp_src;\n\n";
+  c += "  // sp - spatial dimensions, ch - channels dimension\n";
+  c += "  // indexing relative to simdgroup\n";
+  for (int sp = 0; sp < 32; sp += 8) {
+    const std::string sp_start = std::to_string(sp);
+    const std::string sp_end = std::to_string(sp + 8);
+    const std::string dst_name = "dst_sp" + sp_start + "_" + sp_end;
+    for (int slice = 0; slice < dst_x8_slices; slice += 1) {
+      const std::string sl_start = std::to_string(slice * 8);
+      const std::string sl_end = std::to_string(slice * 8 + 8);
+      c += "  simdgroup_matrix<FLT, 8, 8> " + dst_name + "_ch" + sl_start +
+           "_" + sl_end + "(0.0f);\n";
+    }
+  }
+  if (spatial_threads > 32) {
+    c += "  int spatial_group = spatial_id / 32;\n";
+    c += "  tmp_src_x1 += 8 * 8 * 4 * spatial_group;\n";
+  }
+  c += R"(
+  int c_x = min(DST_X, args.src_tensor.Width() - 1);
+  int c_y = min(DST_Y, args.src_tensor.Height() - 1);
+)";
+  if (definition.src_tensors[0].IsLinear()) {
+    c +=
+        "  int src_address = args.src_tensor.GetAddress(c_x, c_y, slice_id);\n";
+  }
+  c += R"(
+  int tid2 = 0;
+  if (tid < SPATIAL_THREADS) {
+    tid2 = tid * 2 + 0;
+  } else if (tid < SPATIAL_THREADS * 2) {
+    tid2 = (tid - SPATIAL_THREADS) * 2 + 1;
+  })";
+  for (int src_s = 1; src_s < src_x8_slices; ++src_s) {
+    c += " else if (tid < SPATIAL_THREADS * " + std::to_string(src_s * 2 + 1) +
+         ") {\n";
+    c += "    tid2 = (tid - SPATIAL_THREADS * " + std::to_string(src_s * 2) +
+         ") * 2 + 0 + SPATIAL_THREADS * " + std::to_string(src_s * 2) + ";\n";
+    c += "  } else if (tid < SPATIAL_THREADS * " +
+         std::to_string(src_s * 2 + 2) + ") {\n";
+    c += "    tid2 = (tid - SPATIAL_THREADS * " +
+         std::to_string(src_s * 2 + 1) + ") * 2 + 1 + SPATIAL_THREADS * " +
+         std::to_string(src_s * 2) + ";\n";
+    c += "  }";
+  }
+  c += "\n\n";
+  c += "  for (int s = 0; s < args.src_tensor.Slices(); s += " +
+       std::to_string(src_x4_slices) + ") {\n";
+  for (int src_s = 0; src_s < src_x8_slices; ++src_s) {
+    const std::string src_range =
+        "i" + std::to_string(src_s * 8) + "_" + std::to_string(src_s * 8 + 8);
+    for (int dst_s = 0; dst_s < dst_x8_slices; ++dst_s) {
+      const std::string dst_range =
+          "o" + std::to_string(dst_s * 8) + "_" + std::to_string(dst_s * 8 + 8);
+      const std::string w_name = "w_" + dst_range + "_" + src_range;
+      c += "    simdgroup_matrix<FLT, 8, 8> " + w_name + ";\n";
+    }
+  }
+  c += "    threadgroup_barrier(mem_flags::mem_threadgroup);\n";
+  if (cache_weight) {
+    const int groups = weights_flt4_per_spatial / spatial_threads;
+    const int reminder = weights_flt4_per_spatial % spatial_threads;
+    for (int i = 0; i < groups; ++i) {
+      c += "    tmp_w_x4[spatial_id + " + std::to_string(spatial_threads * i) +
+           "] = f_offseted[spatial_id + " +
+           std::to_string(spatial_threads * i) + "];\n";
+    }
+    if (reminder != 0) {
+      c += "    if (spatial_id < " + std::to_string(reminder) + ") {\n";
+      c += "      tmp_w_x4[spatial_id + " +
+           std::to_string(spatial_threads * groups) +
+           "] = f_offseted[spatial_id + " +
+           std::to_string(spatial_threads * groups) + "];\n";
+      c += "    }\n";
+    }
+  } else {
+    for (int src_s = 0; src_s < src_x8_slices; ++src_s) {
+      const std::string src_range =
+          "i" + std::to_string(src_s * 8) + "_" + std::to_string(src_s * 8 + 8);
+      for (int dst_s = 0; dst_s < dst_x8_slices; ++dst_s) {
+        const std::string dst_range = "o" + std::to_string(dst_s * 8) + "_" +
+                                      std::to_string(dst_s * 8 + 8);
+        const std::string w_name = "w_" + dst_range + "_" + src_range;
+        c += "    simdgroup_load(" + w_name + ", f_offseted_x1 + " +
+             std::to_string((src_s * dst_x8_slices + dst_s) * 64) + ", 8);\n";
+      }
+    }
+  }
+  if (definition.src_tensors[0].IsLinear()) {
+    c += "    tmp_src[tid2] = args.src_tensor.Read(src_address);\n";
+  } else {
+    c += "    tmp_src[tid2] = args.src_tensor.Read(c_x, c_y, s + slice_id);\n";
+  }
+  if (cache_weight) {
+    c += "    f_offseted += 16 * " +
+         std::to_string(src_x8_slices * dst_x8_slices) + ";\n";
+  } else {
+    c += "    f_offseted_x1 += 64 * " +
+         std::to_string(src_x8_slices * dst_x8_slices) + ";\n";
+  }
+  if (definition.src_tensors[0].IsLinear()) {
+    c += "    src_address += args.src_tensor.SliceStride() * " +
+         std::to_string(src_x4_slices) + ";\n";
+  }
+  c += "    threadgroup_barrier(mem_flags::mem_threadgroup);\n";
+  if (cache_weight) {
+    for (int src_s = 0; src_s < src_x8_slices; ++src_s) {
+      const std::string src_range =
+          "i" + std::to_string(src_s * 8) + "_" + std::to_string(src_s * 8 + 8);
+      for (int dst_s = 0; dst_s < dst_x8_slices; ++dst_s) {
+        const std::string dst_range = "o" + std::to_string(dst_s * 8) + "_" +
+                                      std::to_string(dst_s * 8 + 8);
+        const std::string w_name = "w_" + dst_range + "_" + src_range;
+        c += "    simdgroup_load(" + w_name + ", tmp_w_x1 + " +
+             std::to_string((src_s * dst_x8_slices + dst_s) * 64) + ", 8);\n";
+      }
+    }
+  }
+  c += "    simdgroup_matrix<FLT, 8, 8> mat_src;\n";
+  const int spatial_x8_count = spatial_threads / 8;
+  for (int src_s = 0; src_s < src_x8_slices; ++src_s) {
+    const std::string src_s_range =
+        std::to_string(src_s * 8) + "_" + std::to_string(src_s * 8 + 8);
+    for (int sp = 0; sp < 32; sp += 8) {
+      const std::string sp_range =
+          std::to_string(sp) + "_" + std::to_string(sp + 8);
+      const int src_tile_offset = src_s * spatial_x8_count + (sp / 8);
+      c += "    simdgroup_load(mat_src, tmp_src_x1 + " +
+           std::to_string(src_tile_offset * 64) + ", 8);  // loading sp[" +
+           sp_range + "] src_ch[" + src_s_range + "]\n";
+      for (int dst_s = 0; dst_s < dst_x8_slices; ++dst_s) {
+        const std::string dst_s_range =
+            std::to_string(dst_s * 8) + "_" + std::to_string(dst_s * 8 + 8);
+        const std::string dst_name = "dst_sp" + sp_range + "_ch" + dst_s_range;
+        const std::string w_name = "w_o" + dst_s_range + "_i" + src_s_range;
+        c += "    MMA(" + dst_name + ", mat_src, " + w_name + ", " + dst_name +
+             ");\n";
+      }
+    }
+  }
+  c += "  }\n";
+  for (int slice = 0; slice < dst_x8_slices * 2; slice += 1) {
+    c += "  FLT4 r" + std::to_string(slice) + " = INIT_FLT4(0.0);\n";
+  }
+  c += "  // transferring from simdgroup memory to private registers.\n";
+  c += "  const int kSpatialGroupsCount = " + std::to_string(src_x4_slices) +
+       ";\n";
+  c += "  for (int i = 0; i < kSpatialGroupsCount; ++i) {\n";
+  c += "    int spatial_id = tid - i * SPATIAL_THREADS;\n";
+  c += "    bool current_spatial_group = spatial_id >= 0 && spatial_id < "
+       "SPATIAL_THREADS;\n";
+  for (int dst_s = 0; dst_s < dst_x8_slices; ++dst_s) {
+    const std::string dst_range =
+        "ch" + std::to_string(dst_s * 8) + "_" + std::to_string(dst_s * 8 + 8);
+    c += "    threadgroup_barrier(mem_flags::mem_threadgroup);\n";
+    c += "    if (current_spatial_group) {\n";
+    c += "      simdgroup_store(dst_sp0_8_" + dst_range + ", tmp_src_x1, 8);\n";
+    c += "      simdgroup_store(dst_sp8_16_" + dst_range +
+         ", tmp_src_x1 + 64, 8);\n";
+    c += "      simdgroup_store(dst_sp16_24_" + dst_range +
+         ", tmp_src_x1 + 64 * 2, 8);\n";
+    c += "      simdgroup_store(dst_sp24_32_" + dst_range +
+         ", tmp_src_x1 + 64 * 3, 8);\n";
+    c += "    }\n";
+    c += "    threadgroup_barrier(mem_flags::mem_threadgroup);\n";
+    c += "    if (current_spatial_group) {\n";
+    c += "      r" + std::to_string(dst_s * 2 + 0) +
+         " += tmp_src[spatial_id * 2 + 0];\n";
+    c += "      r" + std::to_string(dst_s * 2 + 1) +
+         " += tmp_src[spatial_id * 2 + 1];\n";
+    c += "    }\n";
+  }
+  c += "  }\n";
+  c += "  if (DST_X >= args.dst_tensor.Width() || DST_Y >= "
+       "args.dst_tensor.Height()) {\n";
+  c += "    return;\n";
+  c += "  }\n";
+  for (int slice = 0; slice < dst_x8_slices * 2; slice += 1) {
+    const std::string dst_s = "DST_S + " + std::to_string(slice);
+    const std::string r_name = "r" + std::to_string(slice);
+    c += "  if (" + dst_s + " < args.dst_tensor.Slices()) {\n";
+    c += "    " + r_name + " += args.biases.Read(" + dst_s + ");\n";
+    c += "    args.dst_tensor.Write(" + r_name + ", DST_X, DST_Y, " + dst_s +
+         ");\n";
+    c += "  }\n";
+  }
+  c += "}\n";
+  return c;
+}
+
+void OIToVecOIOGroupIO(const std::vector<float>& src, int o_size, int i_size,
+                       int vec_size, int o_group_size,
+                       std::vector<float>* dst) {
+  int o_slices = DivideRoundUp(o_size, vec_size);
+  int i_slices = DivideRoundUp(i_size, vec_size);
+  int o_groups = DivideRoundUp(o_slices, o_group_size);
+  dst->resize(o_slices * vec_size * i_slices * vec_size);
+  for (int os = 0; os < o_groups; ++os) {
+    for (int is = 0; is < i_slices; ++is) {
+      for (int o_group = 0; o_group < o_group_size; ++o_group) {
+        for (int sub_o = 0; sub_o < vec_size; ++sub_o) {
+          for (int sub_i = 0; sub_i < vec_size; ++sub_i) {
+            float value = 0.0f;
+            int i_ch = is * vec_size + sub_i;
+            int o_ch = (os * o_group_size + o_group) * vec_size + sub_o;
+            if (i_ch < i_size && o_ch < o_size) {
+              value = src[o_ch * i_size + i_ch];
+            }
+            (*dst)[(((os * i_slices + is) * o_group_size + o_group) * vec_size +
+                    sub_i) *
+                       vec_size +
+                   sub_o] = value;
+          }
+        }
+      }
+    }
+  }
+}
+
+std::vector<uint8_t> ReorderWeightsForConv(
+    const tflite::gpu::Tensor<OHWI, DataType::FLOAT32>& weights,
+    const DataType& weights_type, int dst_x8_slices) {
+  std::vector<float> weights_gpu;
+  OIToVecOIOGroupIO(weights.data, weights.shape.o, weights.shape.i, 8,
+                    dst_x8_slices, &weights_gpu);
+  std::vector<uint8_t> result(weights_gpu.size() * SizeOf(weights_type));
+  if (weights_type == DataType::FLOAT32) {
+    float* gpu_data = reinterpret_cast<float*>(result.data());
+    for (int i = 0; i < weights_gpu.size(); ++i) {
+      gpu_data[i] = weights_gpu[i];
+    }
+  } else {
+    half* gpu_data = reinterpret_cast<half*>(result.data());
+    for (int i = 0; i < weights_gpu.size(); ++i) {
+      gpu_data[i] = weights_gpu[i];
+    }
+  }
+  return result;
+}
+
+std::vector<uint8_t> ReorderBiasesForConv(
+    const tflite::gpu::Tensor<Linear, DataType::FLOAT32>& biases,
+    const DataType& biases_type, int output_size) {
+  std::vector<uint8_t> result(output_size * SizeOf(biases_type));
+  if (biases_type == DataType::FLOAT32) {
+    float* gpu_data = reinterpret_cast<float*>(result.data());
+    for (int i = 0; i < output_size; ++i) {
+      gpu_data[i] = i < biases.shape.v ? biases.data[i] : 0.0f;
+    }
+  } else {
+    half* gpu_data = reinterpret_cast<half*>(result.data());
+    for (int i = 0; i < output_size; ++i) {
+      gpu_data[i] = i < biases.shape.v ? biases.data[i] : 0.0f;
+    }
+  }
+  return result;
+}
+
+std::vector<int2> Get2DWorkgroupsEqualTo32() {
+  return {{8, 4}, {16, 2}, {4, 8}, {32, 1}, {2, 16}, {1, 32}};
+}
+
+int Get2dGroupsCount(const BHWC& dst_shape, const int2 group_size) {
+  int x_groups = DivideRoundUp(dst_shape.w * dst_shape.b, group_size.x);
+  int y_groups = DivideRoundUp(dst_shape.h, group_size.y);
+  return x_groups * y_groups;
+}
+
+int2 GetOptimalGroupSize(const BHWC& dst_shape) {
+  const auto base_work_groups = Get2DWorkgroupsEqualTo32();
+  int min_2d_work_groups = Get2dGroupsCount(dst_shape, base_work_groups[0]);
+  int min_index = 0;
+  for (int i = 1; i < base_work_groups.size(); ++i) {
+    int groups_count = Get2dGroupsCount(dst_shape, base_work_groups[i]);
+    if (groups_count < min_2d_work_groups) {
+      min_2d_work_groups = groups_count;
+      min_index = i;
+    }
+  }
+  return base_work_groups[min_index];
+}
+
+}  // namespace
+
+int3 ConvolutionMetalSimd::GetGridSize() const {
+  const int task_size_x = dst_[0]->Width() * dst_[0]->Batch();
+  const int task_size_y = dst_[0]->Height();
+  const int task_size_z = dst_[0]->Depth();
+  const int task_size_s =
+      DivideRoundUp(dst_[0]->Slices(), params_.slices_per_thread);
+  if (params_.linear_spatial) {
+    return int3(task_size_x * task_size_y * task_size_z, task_size_s, 1);
+  } else {
+    return int3(task_size_x, task_size_y * task_size_z, task_size_s);
+  }
+}
+
+ConvolutionMetalSimd CreateConvolutionMetalSimd(
+    const OperationDef& definition, const BHWC& dst_shape,
+    const Convolution2DAttributes& attr, const GpuInfo& gpu_info) {
+  ConvolutionMetalSimd desc(definition);
+  const int2 optimal_2d_group_size = GetOptimalGroupSize(dst_shape);
+  const int groups2d_count = Get2dGroupsCount(dst_shape, optimal_2d_group_size);
+  const int groups1d_count =
+      DivideRoundUp(dst_shape.w * dst_shape.b * dst_shape.h, 32);
+  if (groups1d_count < groups2d_count) {
+    desc.params_.work_group_size = int3(32, 4, 1);
+    desc.params_.work_group_launch_order = int3(0, 1, 2);
+    desc.params_.linear_spatial = true;
+  } else {
+    desc.params_.work_group_size =
+        int3(optimal_2d_group_size.x, optimal_2d_group_size.y, 4);
+    desc.params_.work_group_launch_order = int3(0, 1, 2);
+    desc.params_.linear_spatial = false;
+  }
+  desc.params_.slices_per_thread = 4;
+  desc.params_.x_kernel_is_1 = true;
+  desc.params_.y_kernel_is_1 = true;
+  desc.params_.z_kernel_is_1 = true;
+  desc.code_ = GenerateConvolution(definition, desc.params_);
+
+  desc.AddSrcTensor("src_tensor", definition.src_tensors[0]);
+  desc.AddDstTensor("dst_tensor", definition.dst_tensors[0]);
+
+  auto weights_type = DeduceDataTypeFromPrecision(definition.precision);
+
+  MemoryType mem_type = MemoryType::GLOBAL;
+
+  if (definition.src_tensors.size() == 2) {
+    // dynamic weights
+    BufferDescriptor weights_desc;
+    weights_desc.element_type = definition.src_tensors[1].GetDataType();
+    weights_desc.element_size = 4;
+    weights_desc.memory_type = mem_type;
+    desc.AddSrcBuffer("weights", weights_desc);
+  } else {
+    BufferDescriptor weights_desc;
+    weights_desc.element_type = weights_type;
+    weights_desc.element_size = 4;
+    weights_desc.memory_type = mem_type;
+    weights_desc.data = ReorderWeightsForConv(
+        attr.weights, weights_type, desc.params_.slices_per_thread / 2);
+    weights_desc.size = weights_desc.data.size();
+    desc.args_.AddObject(
+        "weights", std::make_unique<BufferDescriptor>(std::move(weights_desc)));
+  }
+
+  BufferDescriptor bias_desc;
+  bias_desc.element_type = weights_type;
+  bias_desc.element_size = 4;
+  bias_desc.memory_type = mem_type;
+  bias_desc.data = ReorderBiasesForConv(attr.bias, weights_type,
+                                        AlignByN(attr.weights.shape.o, 4 * 4));
+  bias_desc.size = bias_desc.data.size();
+  desc.args_.AddObject(
+      "biases", std::make_unique<BufferDescriptor>(std::move(bias_desc)));
+
+  desc.work_group_size_ = desc.params_.work_group_size;
+  desc.work_group_launch_order_ = desc.params_.work_group_launch_order;
+  if (desc.params_.linear_spatial) {
+    desc.grid_dimension_ = 2;
+  } else {
+    desc.grid_dimension_ = 3;
+  }
+
+  return desc;
+}
+
+bool IsConvolutionMetalSimdSupported(const GpuInfo& gpu_info,
+                                     const OperationDef& definition,
+                                     const Convolution2DAttributes& attr) {
+  if (!gpu_info.IsApple() || !gpu_info.metal_info.IsSIMDMatMulSupported() ||
+      !gpu_info.apple_info.IsSIMDMatMulSupported()) {
+    return false;
+  }
+  const bool genuine_1x1 =
+      attr.weights.shape.w == 1 && attr.weights.shape.h == 1 &&
+      attr.dilations.w == 1 && attr.dilations.h == 1 && attr.strides.w == 1 &&
+      attr.strides.h == 1 && attr.padding.prepended.w == 0 &&
+      attr.padding.prepended.h == 0 && attr.padding.appended.w == 0 &&
+      attr.padding.appended.h == 0 && attr.groups == 1;
+  const int src_slices = DivideRoundUp(attr.weights.shape.i, 4);
+  const int dst_slices = DivideRoundUp(attr.weights.shape.o, 4);
+  return genuine_1x1 && src_slices % 4 == 0 && dst_slices % 16 == 0;
+}
+
+bool IsGoodTaskSizeForAppleConvSimd(const BHWC& dst_shape,
+                                    const GpuInfo& gpu_info) {
+  const uint64_t task_size_spatial = dst_shape.b * dst_shape.h * dst_shape.w;
+  const uint64_t wave_size = 32;
+  const double useful_part = static_cast<double>(task_size_spatial) /
+                             AlignByN(task_size_spatial, wave_size);
+  if (useful_part < 0.625) {
+    return false;
+  }
+  const double task_size_slices = DivideRoundUp(dst_shape.c, 16);
+  const double task_size = task_size_spatial * task_size_slices;
+  const double task_size_per_cu = task_size / gpu_info.GetComputeUnitsCount();
+  const double waves_per_cu = task_size_per_cu / wave_size;
+  return waves_per_cu >= 8.0;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/conv_metal_simd.h

@@ -0,0 +1,105 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_METAL_SIMD_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_METAL_SIMD_H_
+
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/gpu_info.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/weights_layout.h"
+
+namespace tflite {
+namespace gpu {
+
+class ConvolutionMetalSimd : public GPUOperation {
+ public:
+  ConvolutionMetalSimd() = default;
+  void GetPossibleKernelWorkGroups(
+      TuningType tuning_type, const GpuInfo& gpu_info,
+      const KernelInfo& kernel_info,
+      std::vector<int3>* work_groups) const override {
+    work_groups->push_back(work_group_size_);
+  }
+  int3 GetGridSize() const override;
+
+  // Move only
+  ConvolutionMetalSimd(ConvolutionMetalSimd&& kernel) = default;
+  ConvolutionMetalSimd& operator=(ConvolutionMetalSimd&& kernel) = default;
+  ConvolutionMetalSimd(const ConvolutionMetalSimd&) = delete;
+  ConvolutionMetalSimd& operator=(const ConvolutionMetalSimd&) = delete;
+
+  WeightsDescription GetWeightsDescription() const {
+    WeightsDescription desc;
+    desc.type = DeduceDataTypeFromPrecision(definition_.precision);
+    desc.layout = WeightsLayout::kOSpatialIOGroupO4I4;
+    desc.output_group_size = 4;
+    return desc;
+  }
+
+  struct ConvParams {
+    int3 work_group_size;
+    int3 work_group_launch_order;
+    bool linear_spatial;  // spatial dimensions are Width/Height/Depth
+    int slices_per_thread;
+    bool x_kernel_is_1 = true;
+    bool y_kernel_is_1 = true;
+    bool z_kernel_is_1 = true;
+
+    // must be 32 * k
+    int GetSpatialThreadsCount() const {
+      if (linear_spatial) {
+        return work_group_size.x;
+      } else {
+        return work_group_size.x * work_group_size.y;
+      }
+    }
+
+    int GetX4SlicesCount() const {
+      if (linear_spatial) {
+        return work_group_size.y;
+      } else {
+        return work_group_size.z;
+      }
+    }
+  };
+
+  ConvParams params_;
+
+ private:
+  explicit ConvolutionMetalSimd(const OperationDef& definition)
+      : GPUOperation(definition) {}
+  friend ConvolutionMetalSimd CreateConvolutionMetalSimd(
+      const OperationDef& definition, const BHWC& dst_shape,
+      const Convolution2DAttributes& attr, const GpuInfo& gpu_info);
+};
+
+ConvolutionMetalSimd CreateConvolutionMetalSimd(
+    const OperationDef& definition, const BHWC& dst_shape,
+    const Convolution2DAttributes& attr, const GpuInfo& gpu_info);
+
+bool IsConvolutionMetalSimdSupported(const GpuInfo& gpu_info,
+                                     const OperationDef& definition,
+                                     const Convolution2DAttributes& attr);
+
+bool IsGoodTaskSizeForAppleConvSimd(const BHWC& dst_shape,
+                                    const GpuInfo& gpu_info);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_METAL_SIMD_H_

tensorflow/lite/delegates/gpu/common/tasks/conv_weights_converter.cc

@@ -0,0 +1,247 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/conv_weights_converter.h"
+
+#include <cstring>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "tensorflow/lite/delegates/gpu/common/task/util.h"
+
+namespace tflite {
+namespace gpu {
+
+ConverterToConvWeights::ConverterToConvWeights(
+    const OperationDef& definition, const WeightsDescription& weights_desc,
+    Layout input_layout)
+    : GPUOperation(definition),
+      weights_desc_(weights_desc),
+      input_layout_(input_layout) {
+  code_ = GetConverterToConvWeightsCode();
+}
+
+std::string ConverterToConvWeights::GetConverterToConvWeightsCode() {
+  AddSrcTensor("src_tensor", definition_.src_tensors[0]);
+  args_.AddFloat("mask_x");
+  args_.AddFloat("mask_y");
+  args_.AddFloat("mask_z");
+  args_.AddFloat("mask_w");
+  args_.AddInt("out_ch");
+  args_.AddInt("out_ch_x4_groups");
+  args_.AddInt("in_ch");
+  args_.AddInt("in_ch_x4_groups");
+  args_.AddInt("kernel_width");
+  args_.AddInt("kernel_height");
+  args_.AddInt("kernel_spatial_size");
+
+  if (weights_desc_.layout == WeightsLayout::kOICustomSpatialI4O4 ||
+      weights_desc_.layout == WeightsLayout::kOICustomSpatialO4I4) {
+    std::vector<int32_t> remap(weights_desc_.spatial_remap.size());
+    for (int i = 0; i < remap.size(); ++i) {
+      remap[i] = weights_desc_.spatial_remap[i];
+    }
+    BufferDescriptor desc;
+    desc.element_type = DataType::INT32;
+    desc.element_size = 1;
+    desc.memory_type = MemoryType::GLOBAL;
+    desc.size = remap.size() * sizeof(int32_t);
+    desc.data.resize(desc.size);
+    std::memcpy(desc.data.data(), remap.data(), desc.size);
+    args_.AddObject("spatial_remap",
+                    std::make_unique<BufferDescriptor>(std::move(desc)));
+  }
+
+  std::string c;
+  c += "MAIN_FUNCTION($0) {\n";
+  c += "  int O = GLOBAL_ID_0;\n";
+  c += "  int I = GLOBAL_ID_1;\n";
+  c += "  int spatial_linear = GLOBAL_ID_2;\n";
+  c += "  if (O >= args.out_ch_x4_groups) return;\n";
+  c += "  if (I >= args.in_ch_x4_groups) return;\n";
+  c += "  if (spatial_linear >= args.kernel_spatial_size) return;\n";
+  if (weights_desc_.layout == WeightsLayout::kOICustomSpatialI4O4 ||
+      weights_desc_.layout == WeightsLayout::kOICustomSpatialO4I4) {
+    c += "  int linear_remap = args.spatial_remap.Read(spatial_linear);\n";
+    c += "  int W = linear_remap % args.kernel_width;\n";
+    c += "  int H = linear_remap / args.kernel_width;\n";
+  } else {
+    c += "  int W = spatial_linear % args.kernel_width;\n";
+    c += "  int H = spatial_linear / args.kernel_width;\n";
+  }
+  // W and H is src coordinates, spatial_linear is dst coordinate
+  c += "  FLT4 v0 = INIT_FLT4(0.0f);\n";
+  c += "  FLT4 v1 = INIT_FLT4(0.0f);\n";
+  c += "  FLT4 v2 = INIT_FLT4(0.0f);\n";
+  c += "  FLT4 v3 = INIT_FLT4(0.0f);\n";
+  if (input_layout_ == Layout::OHWI) {
+    c += "  if (O * 4 < args.out_ch) {\n";
+    c += "    v0 = args.src_tensor.Read(W, H, I, O * 4);\n";
+    c += "  }\n";
+    c += "  if (O * 4 + 1 < args.out_ch) {\n";
+    c += "    v1 = args.src_tensor.Read(W, H, I, O * 4 + 1);\n";
+    c += "  }\n";
+    c += "  if (O * 4 + 2 < args.out_ch) {\n";
+    c += "    v2 = args.src_tensor.Read(W, H, I, O * 4 + 2);\n";
+    c += "  }\n";
+    c += "  if (O * 4 + 3 < args.out_ch) {\n";
+    c += "    v3 = args.src_tensor.Read(W, H, I, O * 4 + 3);\n";
+    c += "  }\n";
+    c += "  if (I == args.src_tensor.Slices() - 1) {\n";
+    c += "    FLT4 mask = INIT_FLT4v4(args.mask_x, args.mask_y, args.mask_z, "
+         "args.mask_w);\n";
+    c += "    v0 *= mask;\n";
+    c += "    v1 *= mask;\n";
+    c += "    v2 *= mask;\n";
+    c += "    v3 *= mask;\n";
+    c += "  }\n";
+  } else if (input_layout_ == Layout::HWIO) {
+    c += "  if (I * 4 < args.in_ch && O < args.src_tensor.Slices()) {\n";
+    c += "    v0 = args.src_tensor.Read(I * 4, W, O, H);\n";
+    c += "  }\n";
+    c += "  if (I * 4 + 1 < args.in_ch && O < args.src_tensor.Slices()) {\n";
+    c += "    v1 = args.src_tensor.Read(I * 4 + 1, W, O, H);\n";
+    c += "  }\n";
+    c += "  if (I * 4 + 2 < args.in_ch && O < args.src_tensor.Slices()) {\n";
+    c += "    v2 = args.src_tensor.Read(I * 4 + 2, W, O, H);\n";
+    c += "  }\n";
+    c += "  if (I * 4 + 3 < args.in_ch && O < args.src_tensor.Slices()) {\n";
+    c += "    v3 = args.src_tensor.Read(I * 4 + 3, W, O, H);\n";
+    c += "  }\n";
+    c += "  if (O == args.src_tensor.Slices() - 1) {\n";
+    c += "    FLT4 mask = INIT_FLT4v4(args.mask_x, args.mask_y, args.mask_z, "
+         "args.mask_w);\n";
+    c += "    v0 *= mask;\n";
+    c += "    v1 *= mask;\n";
+    c += "    v2 *= mask;\n";
+    c += "    v3 *= mask;\n";
+    c += "  }\n";
+  }
+  const bool need_transpose =
+      (input_layout_ == Layout::HWIO && weights_desc_.IsO4I4()) ||
+      (input_layout_ == Layout::OHWI && weights_desc_.IsI4O4());
+  if (need_transpose) {
+    c += "  FLT4 r0 = INIT_FLT4v4(v0.x, v1.x, v2.x, v3.x);\n";
+    c += "  FLT4 r1 = INIT_FLT4v4(v0.y, v1.y, v2.y, v3.y);\n";
+    c += "  FLT4 r2 = INIT_FLT4v4(v0.z, v1.z, v2.z, v3.z);\n";
+    c += "  FLT4 r3 = INIT_FLT4v4(v0.w, v1.w, v2.w, v3.w);\n";
+  } else {
+    c += "  FLT4 r0 = v0;\n";
+    c += "  FLT4 r1 = v1;\n";
+    c += "  FLT4 r2 = v2;\n";
+    c += "  FLT4 r3 = v3;\n";
+  }
+  if (weights_desc_.layout ==
+          WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4 ||
+      weights_desc_.layout ==
+          WeightsLayout::k2DX4O4YIsSpatialIAndXIsOOGroupI4) {
+    // Writing to 4X Textures 2D
+    AddDstTensor("dst_tensor0", definition_.dst_tensors[0]);
+    AddDstTensor("dst_tensor1", definition_.dst_tensors[1]);
+    AddDstTensor("dst_tensor2", definition_.dst_tensors[2]);
+    AddDstTensor("dst_tensor3", definition_.dst_tensors[3]);
+    c += "  int yc = spatial_linear *  args.in_ch_x4_groups + I;\n";
+    c += "  args.dst_tensor0.Write2D(r0, O, yc);\n";
+    c += "  args.dst_tensor1.Write2D(r1, O, yc);\n";
+    c += "  args.dst_tensor2.Write2D(r2, O, yc);\n";
+    c += "  args.dst_tensor3.Write2D(r3, O, yc);\n";
+    c += "}\n";
+  } else {
+    // Writing to linear buffer
+    AddDstTensor("dst_tensor", definition_.dst_tensors[0]);
+    c += "  int OUTPUT_GROUP_SIZE = " +
+         std::to_string(weights_desc_.GetOutputGroupSize()) + ";\n";
+    c += "  int d_index = (O * 4) / (OUTPUT_GROUP_SIZE * 4);\n";
+    c += "  int k_index = ((O * 4) % (OUTPUT_GROUP_SIZE * 4)) / 4;\n";
+    std::string index;
+    if (weights_desc_.layout == WeightsLayout::kOICustomSpatialI4O4 ||
+        weights_desc_.layout == WeightsLayout::kOICustomSpatialO4I4) {
+      index =
+          "(d_index * args.in_ch_x4_groups + I) * args.kernel_spatial_size + "
+          "spatial_linear";
+    } else if (weights_desc_.layout == WeightsLayout::kOSpatialIOGroupI4O4 ||
+               weights_desc_.layout == WeightsLayout::kOSpatialIOGroupO4I4) {
+      index =
+          "(d_index * args.kernel_spatial_size + spatial_linear) * "
+          "args.in_ch_x4_groups + I";
+    }
+    c += "  int dst_offset = (" + index + ") * OUTPUT_GROUP_SIZE + k_index;\n";
+    c += "  args.dst_tensor.WriteLinear(r0, dst_offset * 4 + 0);\n";
+    c += "  args.dst_tensor.WriteLinear(r1, dst_offset * 4 + 1);\n";
+    c += "  args.dst_tensor.WriteLinear(r2, dst_offset * 4 + 2);\n";
+    c += "  args.dst_tensor.WriteLinear(r3, dst_offset * 4 + 3);\n";
+    c += "}\n";
+  }
+  return c;
+}
+
+OHWI ConverterToConvWeights::GetWeightsSize() const {
+  int output_channels = 0;
+  int input_channels = 0;
+  int kernel_width = 0;
+  int kernel_height = 0;
+  if (input_layout_ == Layout::HWIO) {
+    output_channels = src_[0]->Channels();
+    input_channels = src_[0]->Width();
+    kernel_width = src_[0]->Height();
+    kernel_height = src_[0]->Batch();
+  } else if (input_layout_ == Layout::OHWI) {
+    output_channels = src_[0]->Batch();
+    input_channels = src_[0]->Channels();
+    kernel_width = src_[0]->Width();
+    kernel_height = src_[0]->Height();
+  }
+  return OHWI(output_channels, kernel_height, kernel_width, input_channels);
+}
+
+absl::Status ConverterToConvWeights::BindArguments(ArgumentsBinder* args) {
+  const auto& weights_shape = GetWeightsSize();
+  const int output_channels_x4_groups = DivideRoundUp(
+      AlignByN(weights_shape.o, 4 * weights_desc_.GetOutputGroupSize()), 4);
+  RETURN_IF_ERROR(args->SetInt("out_ch", weights_shape.o));
+  RETURN_IF_ERROR(args->SetInt("out_ch_x4_groups", output_channels_x4_groups));
+  RETURN_IF_ERROR(args->SetInt("in_ch", weights_shape.i));
+  RETURN_IF_ERROR(
+      args->SetInt("in_ch_x4_groups", DivideRoundUp(weights_shape.i, 4)));
+  RETURN_IF_ERROR(args->SetInt("kernel_width", weights_shape.w));
+  RETURN_IF_ERROR(args->SetInt("kernel_height", weights_shape.h));
+  RETURN_IF_ERROR(
+      args->SetInt("kernel_spatial_size", weights_shape.w * weights_shape.h));
+  float4 mask = GetMaskForLastPlane(src_[0]->Channels());
+  RETURN_IF_ERROR(args->SetFloat("mask_x", mask.x));
+  RETURN_IF_ERROR(args->SetFloat("mask_y", mask.y));
+  RETURN_IF_ERROR(args->SetFloat("mask_z", mask.z));
+  return args->SetFloat("mask_w", mask.w);
+}
+
+int3 ConverterToConvWeights::GetGridSize() const {
+  const auto& weights_shape = GetWeightsSize();
+  const int out_group_size = weights_desc_.GetOutputGroupSize();
+  const int grid_x =
+      DivideRoundUp(AlignByN(weights_shape.o, 4 * out_group_size), 4);
+  const int grid_y = DivideRoundUp(weights_shape.i, 4);
+  const int grid_z = weights_shape.w * weights_shape.h;
+  return int3(grid_x, grid_y, grid_z);
+}
+
+ConverterToConvWeights CreateConverterToConvWeights(
+    const OperationDef& definition, const WeightsDescription& weights_desc,
+    Layout input_layout) {
+  return ConverterToConvWeights(definition, weights_desc, input_layout);
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/conv_weights_converter.h

@@ -0,0 +1,63 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_WEIGHTS_CONVERTER_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_WEIGHTS_CONVERTER_H_
+
+#include <string>
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/weights_layout.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+class ConverterToConvWeights : public GPUOperation {
+ public:
+  ConverterToConvWeights(const OperationDef& definition,
+                         const WeightsDescription& weights_desc,
+                         Layout input_layout);
+  absl::Status BindArguments(ArgumentsBinder* args) override;
+  int3 GetGridSize() const override;
+
+  // Move only
+  ConverterToConvWeights(ConverterToConvWeights&& operation) = default;
+  ConverterToConvWeights& operator=(ConverterToConvWeights&& operation) =
+      default;
+  ConverterToConvWeights(const ConverterToConvWeights&) = delete;
+  ConverterToConvWeights& operator=(const ConverterToConvWeights&) = delete;
+
+ private:
+  std::string GetConverterToConvWeightsCode();
+
+  OHWI GetWeightsSize() const;
+
+  WeightsDescription weights_desc_;
+
+  Layout input_layout_;  // Can be only OHWI or HWIO
+  // if input_layout_ is OHWI: reinterpreting weights as OHWI-BHWC tensor
+  // if input_layout_ is HWIO: reinterpreting weights as HWIO-BHWC tensor
+};
+
+ConverterToConvWeights CreateConverterToConvWeights(
+    const OperationDef& definition, const WeightsDescription& weights_desc,
+    Layout input_layout);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_WEIGHTS_CONVERTER_H_

tensorflow/lite/delegates/gpu/common/tasks/conv_weights_converter_test_util.cc

@@ -0,0 +1,342 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/conv_weights_converter_test_util.h"
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/task/weights_conversion.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/conv_weights_converter.h"
+
+namespace tflite {
+namespace gpu {
+namespace {
+absl::Status ConvolutionWeightsConverterTest(
+    const Tensor<OHWI, DataType::FLOAT32>& weights,
+    const WeightsDescription& weight_desc, TestExecutionEnvironment* env,
+    const OperationDef& op_def) {
+  // reinterpreting weights as HWIO-BHWC tensor
+  TensorFloat32 src_tensor_as_hwio;
+  src_tensor_as_hwio.shape =
+      BHWC(weights.shape.h, weights.shape.w, weights.shape.i, weights.shape.o);
+  src_tensor_as_hwio.data.resize(src_tensor_as_hwio.shape.DimensionsProduct(),
+                                 2.0);
+  // reinterpreting weights as OHWI-BHWC tensor
+  TensorFloat32 src_tensor_as_ohwi;
+  src_tensor_as_ohwi.shape =
+      BHWC(weights.shape.o, weights.shape.h, weights.shape.w, weights.shape.i);
+  src_tensor_as_ohwi.data.resize(src_tensor_as_ohwi.shape.DimensionsProduct(),
+                                 2.0);
+  for (int o = 0; o < weights.shape.o; ++o) {
+    for (int y = 0; y < weights.shape.h; ++y) {
+      for (int x = 0; x < weights.shape.w; ++x) {
+        for (int i = 0; i < weights.shape.i; ++i) {
+          const int f_index = weights.shape.LinearIndex({o, y, x, i});
+          const int s_index_hwio =
+              src_tensor_as_hwio.shape.LinearIndex({y, x, i, o});
+          src_tensor_as_hwio.data[s_index_hwio] = weights.data[f_index];
+          const int s_index_ohwi =
+              src_tensor_as_ohwi.shape.LinearIndex({o, y, x, i});
+          src_tensor_as_ohwi.data[s_index_ohwi] = weights.data[f_index];
+        }
+      }
+    }
+  }
+
+  WeightsDescription weight_desc_copy = weight_desc;
+  weight_desc_copy.type = DataType::FLOAT32;
+  const int flt_count =
+      GetTotalElementsCountForLayout(weight_desc_copy, weights.shape);
+  DataType weights_type = DataType::FLOAT32;
+
+  std::vector<uint8_t> weights_data(flt_count * SizeOf(weights_type));
+  RearrangeWeights(weights, weight_desc_copy, absl::MakeSpan(weights_data));
+
+  std::vector<TensorFloat32> dst_tensors;
+  if (weight_desc_copy.layout ==
+          WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4 ||
+      weight_desc_copy.layout ==
+          WeightsLayout::k2DX4O4YIsSpatialIAndXIsOOGroupI4) {
+    dst_tensors.resize(4);
+    const int dst_depth = AlignByN(DivideRoundUp(weights.shape.o, 4),
+                                   weight_desc_copy.output_group_size);
+    const int src_depth = DivideRoundUp(weights.shape.i, 4);
+    const int kernel_x = weights.shape.w;
+    const int kernel_y = weights.shape.h;
+    int texture_width = dst_depth;
+    int texture_height = src_depth * kernel_x * kernel_y;
+    int sub_size = SizeOf(weights_type) * 4 * texture_width * texture_height;
+    for (int i = 0; i < 4; ++i) {
+      dst_tensors[i].shape = BHWC(1, texture_height, texture_width, 4);
+      dst_tensors[i].data.resize(4 * texture_width * texture_height);
+      memcpy(dst_tensors[i].data.data(), weights_data.data() + sub_size * i,
+             sub_size);
+    }
+  } else {
+    dst_tensors.resize(1);
+    dst_tensors[0].shape = BHWC(1, 1, 1, flt_count);
+    dst_tensors[0].data.resize(flt_count);
+    memcpy(dst_tensors[0].data.data(), weights_data.data(),
+           flt_count * SizeOf(weights_type));
+  }
+
+  std::vector<TensorFloat32> dst_tensors_gpu(dst_tensors.size());
+  std::vector<TensorFloat32*> dst_ptrs;
+  std::vector<BHWC> dst_shapes;
+  for (int i = 0; i < dst_tensors.size(); ++i) {
+    dst_shapes.push_back(dst_tensors[i].shape);
+    dst_ptrs.push_back(&dst_tensors_gpu[i]);
+  }
+
+  auto converter_from_ohwi = ConverterToConvWeights(
+      op_def, weight_desc, /*input layout*/ Layout::OHWI);
+  RETURN_IF_ERROR(env->ExecuteGPUOperation(
+      {src_tensor_as_ohwi},
+      std::make_unique<ConverterToConvWeights>(std::move(converter_from_ohwi)),
+      dst_shapes, dst_ptrs));
+  for (int i = 0; i < dst_tensors.size(); ++i) {
+    RETURN_IF_ERROR(
+        PointWiseNear(dst_tensors[i].data, dst_tensors_gpu[i].data, 0.0f));
+  }
+
+  auto converter_from_hwio = ConverterToConvWeights(
+      op_def, weight_desc, /*input layout*/ Layout::HWIO);
+  RETURN_IF_ERROR(env->ExecuteGPUOperation(
+      {src_tensor_as_hwio},
+      std::make_unique<ConverterToConvWeights>(std::move(converter_from_hwio)),
+      dst_shapes, dst_ptrs));
+  for (int i = 0; i < dst_tensors.size(); ++i) {
+    RETURN_IF_ERROR(
+        PointWiseNear(dst_tensors[i].data, dst_tensors_gpu[i].data, 0.0f));
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status ConverterToConvWeights1x1OutX4Test(TestExecutionEnvironment* env) {
+  const int kSrcChannels = 8;
+  const int kDstChannels = 32;
+  auto weights_shape = OHWI(kDstChannels, 1, 1, kSrcChannels);
+  WeightsDescription conv_weight_desc;
+  conv_weight_desc.output_group_size = 4;
+
+  Tensor<OHWI, DataType::FLOAT32> weights;
+  weights.shape = weights_shape;
+  weights.data.resize(weights_shape.DimensionsProduct());
+  for (int i = 0; i < weights.data.size(); ++i) {
+    weights.data[i] = half(static_cast<float>(i));
+  }
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      for (auto weights_layout : {WeightsLayout::kOSpatialIOGroupI4O4,
+                                  WeightsLayout::kOSpatialIOGroupO4I4}) {
+        conv_weight_desc.layout = weights_layout;
+        OperationDef op_def;
+        op_def.precision = precision;
+        op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
+        op_def.dst_tensors.push_back(
+            {data_type, TensorStorageType::BUFFER, Layout::UNKNOWN});
+        RETURN_IF_ERROR(ConvolutionWeightsConverterTest(
+            weights, conv_weight_desc, env, op_def));
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConverterToConvWeights1x1OutX4UnalignedTest(
+    TestExecutionEnvironment* env) {
+  const int kSrcChannels = 8;
+  const int kDstChannels = 17;
+  auto weights_shape = OHWI(kDstChannels, 1, 1, kSrcChannels);
+  WeightsDescription conv_weight_desc;
+  conv_weight_desc.output_group_size = 4;
+
+  Tensor<OHWI, DataType::FLOAT32> weights;
+  weights.shape = weights_shape;
+  weights.data.resize(weights_shape.DimensionsProduct());
+  for (int i = 0; i < weights.data.size(); ++i) {
+    weights.data[i] = half(static_cast<float>(i));
+  }
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      for (auto weights_layout : {WeightsLayout::kOSpatialIOGroupI4O4,
+                                  WeightsLayout::kOSpatialIOGroupO4I4}) {
+        conv_weight_desc.layout = weights_layout;
+        OperationDef op_def;
+        op_def.precision = precision;
+        op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
+        op_def.dst_tensors.push_back(
+            {data_type, TensorStorageType::BUFFER, Layout::UNKNOWN});
+        RETURN_IF_ERROR(ConvolutionWeightsConverterTest(
+            weights, conv_weight_desc, env, op_def));
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConverterToConvWeights1x1OutX2Test(TestExecutionEnvironment* env) {
+  const int kSrcChannels = 7;
+  const int kDstChannels = 37;
+  auto weights_shape = OHWI(kDstChannels, 1, 1, kSrcChannels);
+  WeightsDescription conv_weight_desc;
+  conv_weight_desc.output_group_size = 2;
+
+  Tensor<OHWI, DataType::FLOAT32> weights;
+  weights.shape = weights_shape;
+  weights.data.resize(weights_shape.DimensionsProduct());
+  for (int i = 0; i < weights.data.size(); ++i) {
+    weights.data[i] = half(static_cast<float>(i));
+  }
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      for (auto weights_layout : {WeightsLayout::kOSpatialIOGroupI4O4,
+                                  WeightsLayout::kOSpatialIOGroupO4I4}) {
+        conv_weight_desc.layout = weights_layout;
+        OperationDef op_def;
+        op_def.precision = precision;
+        op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
+        op_def.dst_tensors.push_back(
+            {data_type, TensorStorageType::BUFFER, Layout::UNKNOWN});
+        RETURN_IF_ERROR(ConvolutionWeightsConverterTest(
+            weights, conv_weight_desc, env, op_def));
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConverterToConvWeightsOutX2Test(TestExecutionEnvironment* env) {
+  const int kSrcChannels = 8;
+  const int kDstChannels = 38;
+  auto weights_shape = OHWI(kDstChannels, 3, 4, kSrcChannels);
+  WeightsDescription conv_weight_desc;
+  conv_weight_desc.output_group_size = 2;
+
+  Tensor<OHWI, DataType::FLOAT32> weights;
+  weights.shape = weights_shape;
+  weights.data.resize(weights_shape.DimensionsProduct());
+  for (int i = 0; i < weights.data.size(); ++i) {
+    weights.data[i] = half(static_cast<float>(i));
+  }
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      for (auto weights_layout : {WeightsLayout::kOSpatialIOGroupI4O4,
+                                  WeightsLayout::kOSpatialIOGroupO4I4}) {
+        conv_weight_desc.layout = weights_layout;
+        OperationDef op_def;
+        op_def.precision = precision;
+        op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
+        op_def.dst_tensors.push_back(
+            {data_type, TensorStorageType::BUFFER, Layout::UNKNOWN});
+        RETURN_IF_ERROR(ConvolutionWeightsConverterTest(
+            weights, conv_weight_desc, env, op_def));
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConverterToConvTransposedWeights4x4Test(
+    TestExecutionEnvironment* env) {
+  const int kSrcChannels = 7;
+  const int kDstChannels = 11;
+  auto weights_shape = OHWI(kDstChannels, 4, 4, kSrcChannels);
+  WeightsDescription weight_desc;
+  weight_desc.spatial_remap = {10, 11, 14, 15, 8, 9, 12, 13,
+                               2,  3,  6,  7,  0, 1, 4,  5};
+
+  Tensor<OHWI, DataType::FLOAT32> weights;
+  weights.shape = weights_shape;
+  weights.data.resize(weights_shape.DimensionsProduct());
+  for (int i = 0; i < weights.data.size(); ++i) {
+    weights.data[i] = half(static_cast<float>(i));
+  }
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      for (auto weights_layout : {WeightsLayout::kOICustomSpatialI4O4,
+                                  WeightsLayout::kOICustomSpatialO4I4}) {
+        weight_desc.layout = weights_layout;
+        OperationDef op_def;
+        op_def.precision = precision;
+        op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
+        op_def.dst_tensors.push_back(
+            {data_type, TensorStorageType::BUFFER, Layout::UNKNOWN});
+        RETURN_IF_ERROR(
+            ConvolutionWeightsConverterTest(weights, weight_desc, env, op_def));
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConverterToConvWeights4xTexturesTest(
+    TestExecutionEnvironment* env) {
+  const int src_channels = 9;
+  const int dst_channels = 17;
+  auto weights_shape = OHWI(dst_channels, 1, 1, src_channels);
+  WeightsDescription conv_weight_desc;
+  conv_weight_desc.output_group_size = 4;
+
+  Tensor<OHWI, DataType::FLOAT32> weights;
+  weights.shape = weights_shape;
+  weights.data.resize(weights_shape.DimensionsProduct());
+  for (int i = 0; i < weights.data.size(); ++i) {
+    weights.data[i] = half(static_cast<float>(i));
+  }
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      for (auto weights_layout :
+           {WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4,
+            WeightsLayout::k2DX4O4YIsSpatialIAndXIsOOGroupI4}) {
+        conv_weight_desc.layout = weights_layout;
+        OperationDef op_def;
+        op_def.precision = precision;
+        op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
+        op_def.dst_tensors.push_back(
+            {data_type, TensorStorageType::TEXTURE_2D, Layout::HWC});
+        op_def.dst_tensors.push_back(
+            {data_type, TensorStorageType::TEXTURE_2D, Layout::HWC});
+        op_def.dst_tensors.push_back(
+            {data_type, TensorStorageType::TEXTURE_2D, Layout::HWC});
+        op_def.dst_tensors.push_back(
+            {data_type, TensorStorageType::TEXTURE_2D, Layout::HWC});
+        RETURN_IF_ERROR(ConvolutionWeightsConverterTest(
+            weights, conv_weight_desc, env, op_def));
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/conv_weights_converter_test_util.h

@@ -0,0 +1,38 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_WEIGHTS_CONVERTER_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_WEIGHTS_CONVERTER_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConverterToConvWeights1x1OutX4Test(TestExecutionEnvironment* env);
+absl::Status ConverterToConvWeights1x1OutX4UnalignedTest(
+    TestExecutionEnvironment* env);
+absl::Status ConverterToConvWeights1x1OutX2Test(TestExecutionEnvironment* env);
+absl::Status ConverterToConvWeightsOutX2Test(TestExecutionEnvironment* env);
+absl::Status ConverterToConvTransposedWeights4x4Test(
+    TestExecutionEnvironment* env);
+absl::Status ConverterToConvWeights4xTexturesTest(
+    TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONV_WEIGHTS_CONVERTER_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/conversion.cc

@@ -0,0 +1,124 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/conversion.h"
+
+#include <memory>
+#include <string>
+
+#include "absl/strings/substitute.h"
+#include "tensorflow/lite/delegates/gpu/common/task/util.h"
+
+namespace tflite {
+namespace gpu {
+
+GPUOperation CreateTensorToTensorOp(const GpuInfo& gpu_info,
+                                    const TensorDescriptor& src_desc,
+                                    const TensorDescriptor& dst_desc) {
+  GPUOperation op;
+  op.args_.AddObjectRef("src_tensor", AccessType::READ,
+                        std::make_unique<TensorDescriptor>(src_desc));
+  op.args_.AddObjectRef("dst_tensor", AccessType::WRITE,
+                        std::make_unique<TensorDescriptor>(dst_desc));
+  op.code_ +=
+      R"(MAIN_FUNCTION($0) {
+  int linear_id = get_global_id(0);
+  int x = linear_id / args.dst_tensor.Batch();
+  int b = linear_id % args.dst_tensor.Batch();
+  int y = get_global_id(1);
+  int d = get_global_id(2);
+  if (x >= args.dst_tensor.Width() || y >= args.dst_tensor.Height() || d >= args.dst_tensor.Slices()) return;
+  args.src_tensor::type in_value = args.src_tensor.Read(x, y, d, b);
+)";
+  const std::string conversion = GetTypeConversion(
+      gpu_info, src_desc.GetDataType(), dst_desc.GetDataType(), 4);
+  op.code_ += "  args.dst_tensor::type out_value = " +
+              absl::Substitute(conversion, "in_value") + ";\n";
+  op.code_ += "args.dst_tensor.Write(out_value, x, y, d, b);\n";
+  op.code_ += "}\n";
+  return op;
+}
+
+GPUOperation CreateTensorToBhwcBufferOp(const GpuInfo& gpu_info,
+                                        const TensorDescriptor& src_desc,
+                                        const BufferDescriptor& dst_desc) {
+  GPUOperation op;
+  op.args_.AddObjectRef("tensor", AccessType::READ,
+                        std::make_unique<TensorDescriptor>(src_desc));
+  op.args_.AddObjectRef("buffer", AccessType::WRITE,
+                        std::make_unique<BufferDescriptor>(dst_desc));
+
+  op.code_ += R"(MAIN_FUNCTION($0) {
+  int linear_id = get_global_id(0);
+  int x = linear_id / args.tensor.Batch();
+  int b = linear_id % args.tensor.Batch();
+  int y = get_global_id(1);
+  int d = get_global_id(2);
+  if (x >= args.tensor.Width() || y >= args.tensor.Height() || d >= args.tensor.Slices()) return;
+  args.tensor::type in_value = args.tensor.Read(x, y, d, b);)";
+  const std::string conversion = GetTypeConversion(
+      gpu_info, src_desc.GetDataType(), dst_desc.element_type, 4);
+  op.code_ += "  " + GetTypeDeclaration(gpu_info, dst_desc.element_type, 4) +
+              " out_value = " + absl::Substitute(conversion, "in_value") +
+              ";\n";
+  op.code_ += R"(
+  int c = d * 4;
+  int index = ((b * args.tensor.Height() + y) * args.tensor.Width() + x) * args.tensor.Channels() + c;
+
+  args.buffer.Write(out_value.x, index);
+  if (c + 1 < args.tensor.Channels()) {
+    args.buffer.Write(out_value.y, index + 1);
+  }
+  if (c + 2 < args.tensor.Channels()) {
+    args.buffer.Write(out_value.z, index + 2);
+  }
+  if (c + 3 < args.tensor.Channels()) {
+    args.buffer.Write(out_value.w, index + 3);
+  }
+})";
+  return op;
+}
+
+GPUOperation CreateBhwcBufferToTensorOp(const GpuInfo& gpu_info,
+                                        const BufferDescriptor& src_desc,
+                                        const TensorDescriptor& dst_desc) {
+  GPUOperation op;
+  op.args_.AddObjectRef("buffer", AccessType::READ,
+                        std::make_unique<BufferDescriptor>(src_desc));
+  op.args_.AddObjectRef("tensor", AccessType::WRITE,
+                        std::make_unique<TensorDescriptor>(dst_desc));
+
+  op.code_ += R"(MAIN_FUNCTION($0) {
+  int linear_id = get_global_id(0);
+  int x = linear_id / args.tensor.Batch();
+  int b = linear_id % args.tensor.Batch();
+  int y = get_global_id(1);
+  int d = get_global_id(2);
+
+  if (x >= args.tensor.Width() || y >= args.tensor.Height() || d >= args.tensor.Slices()) return;
+  int c = d * 4;
+  int index = ((b * args.tensor.Height() + y) * args.tensor.Width() + x) * args.tensor.Channels() + c;
+  args.tensor::type result;
+  result.x = args.buffer.Read(index);
+  result.y = c + 1 < args.tensor.Channels() ? args.buffer.Read(index + 1) : 1;
+  result.z = c + 2 < args.tensor.Channels() ? args.buffer.Read(index + 2) : 2;
+  result.w = c + 3 < args.tensor.Channels() ? args.buffer.Read(index + 3) : 3;
+  args.tensor.Write(result, x, y, d, b);
+})";
+  return op;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/conversion.h

@@ -0,0 +1,39 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVERSION_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVERSION_H_
+
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+
+namespace tflite {
+namespace gpu {
+
+GPUOperation CreateTensorToTensorOp(const GpuInfo& gpu_info,
+                                    const TensorDescriptor& src_desc,
+                                    const TensorDescriptor& dst_desc);
+
+GPUOperation CreateTensorToBhwcBufferOp(const GpuInfo& gpu_info,
+                                        const TensorDescriptor& src_desc,
+                                        const BufferDescriptor& dst_desc);
+
+GPUOperation CreateBhwcBufferToTensorOp(const GpuInfo& gpu_info,
+                                        const BufferDescriptor& src_desc,
+                                        const TensorDescriptor& dst_desc);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVERSION_H_

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed.cc

@@ -0,0 +1,663 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/strings/substitute.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/weights_layout.h"
+#include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
+
+namespace tflite {
+namespace gpu {
+namespace {
+bool UseBufferForWeights(const GpuInfo& gpu_info) {
+  return gpu_info.IsMali() || gpu_info.IsApple() || gpu_info.IsAMD();
+}
+}  // namespace
+
+ConvolutionTransposed::ConvolutionTransposed(
+    const OperationDef& definition, const ConvolutionTransposedAttributes& attr,
+    const GpuInfo& gpu_info)
+    : GPUOperation(definition),
+      stride_(attr.stride.w, attr.stride.h, 1, 1),
+      block_size_(2, 2, 1, 2) {
+  if (UseBufferForWeights(gpu_info)) {
+    if (gpu_info.IsApple()) {
+      weights_layout_ = WeightsLayout::kOSpatialIOGroupO4I4;
+    } else {
+      weights_layout_ = WeightsLayout::kOSpatialIOGroupI4O4;
+    }
+  } else {
+    if (gpu_info.IsApple()) {
+      weights_layout_ = WeightsLayout::k2DX4O4YIsSpatialIAndXIsOOGroupI4;
+    } else {
+      weights_layout_ = WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4;
+    }
+  }
+  const bool is_f16 = definition.precision == CalculationsPrecision::F16;
+  if (gpu_info.IsMali()) {
+    if (gpu_info.mali_info.IsMidgard()) {
+      block_size_ = is_f16 ? int4(2, 1, 1, 2) : int4(2, 1, 1, 1);
+    } else {
+      block_size_ = is_f16 ? int4(2, 2, 1, 2) : int4(2, 2, 1, 1);
+    }
+    compiler_options_.push_back(CompilerOptions::kClFastRelaxedMath);
+  }
+  const int dst_depth = DivideRoundUp(attr.weights.shape.o, 4);
+  if (dst_depth == 1 || dst_depth == 3) {
+    if (!gpu_info.IsMali()) {
+      block_size_.y *= block_size_.w;
+    }
+    block_size_.w = 1;
+  }
+
+  args_.AddInt("stride_x", stride_.x);
+  args_.AddInt("stride_y", stride_.y);
+  args_.AddInt("padding_x", attr.padding.prepended.w);
+  args_.AddInt("padding_y", attr.padding.prepended.h);
+  args_.AddInt("kernel_size_x", attr.weights.shape.w);
+  args_.AddInt("kernel_size_y", attr.weights.shape.h);
+  code_ = GenerateConvolutionTransposedCode(definition_, gpu_info, block_size_);
+}
+
+ConvolutionTransposed::ConvolutionTransposed(
+    const OperationDef& definition,
+    const ConvolutionTransposed3DAttributes& attr, const GpuInfo& gpu_info)
+    : GPUOperation(definition),
+      stride_(attr.stride.w, attr.stride.h, attr.stride.d, 1),
+      block_size_(2, 2, 1, 2) {
+  if (UseBufferForWeights(gpu_info)) {
+    if (gpu_info.IsApple()) {
+      weights_layout_ = WeightsLayout::kOSpatialIOGroupO4I4;
+    } else {
+      weights_layout_ = WeightsLayout::kOSpatialIOGroupI4O4;
+    }
+  } else {
+    if (gpu_info.IsApple()) {
+      weights_layout_ = WeightsLayout::k2DX4O4YIsSpatialIAndXIsOOGroupI4;
+    } else {
+      weights_layout_ = WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4;
+    }
+  }
+  const bool is_f16 = definition.precision == CalculationsPrecision::F16;
+  if (gpu_info.IsMali()) {
+    if (gpu_info.mali_info.IsMidgard()) {
+      block_size_ = is_f16 ? int4(2, 1, 1, 2) : int4(2, 1, 1, 1);
+    } else {
+      block_size_ = is_f16 ? int4(2, 2, 1, 2) : int4(2, 2, 1, 1);
+    }
+    compiler_options_.push_back(CompilerOptions::kClFastRelaxedMath);
+  }
+  const int dst_depth = DivideRoundUp(attr.weights.shape.o, 4);
+  if (dst_depth == 1 || dst_depth == 3) {
+    if (!gpu_info.IsMali()) {
+      block_size_.y *= block_size_.w;
+    }
+    block_size_.w = 1;
+  }
+
+  args_.AddInt("stride_x", stride_.x);
+  args_.AddInt("stride_y", stride_.y);
+  args_.AddInt("stride_z", stride_.z);
+  args_.AddInt("padding_x", attr.padding.prepended.w);
+  args_.AddInt("padding_y", attr.padding.prepended.h);
+  args_.AddInt("padding_z", attr.padding.prepended.d);
+  args_.AddInt("kernel_size_x", attr.weights.shape.w);
+  args_.AddInt("kernel_size_y", attr.weights.shape.h);
+  args_.AddInt("kernel_size_z", attr.weights.shape.d);
+  args_.AddInt("grid_size_y");
+  code_ = GenerateConvolutionTransposedCode(definition_, gpu_info, block_size_);
+}
+
+std::string ConvolutionTransposed::GenerateConvolutionTransposedCode(
+    const OperationDef& op_def, const GpuInfo& gpu_info,
+    const int4& block_size) {
+  AddSrcTensor("src_tensor", op_def.src_tensors[0]);
+  AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
+
+  if (op_def.src_tensors.size() != 1) {
+    // dynamic weights
+    if (weights_layout_ == WeightsLayout::kOSpatialIOGroupI4O4 ||
+        weights_layout_ == WeightsLayout::kOSpatialIOGroupO4I4) {
+      BufferDescriptor desc;
+      desc.element_type = op_def.src_tensors[1].GetDataType();
+      desc.element_size = 4;
+      desc.memory_type = MemoryType::GLOBAL;
+      AddSrcBuffer("weights", desc);
+    } else {
+      for (int i = 0; i < 4; ++i) {
+        const std::string name = "weights" + std::to_string(i);
+        AddSrcTensor(name, definition_.src_tensors[1 + i]);
+      }
+    }
+  }
+
+  const auto& src_def = op_def.src_tensors[0];
+
+  std::string c;
+
+  const bool weights_are_buffer = UseBufferForWeights(gpu_info);
+  for (int s = 0; s < block_size.w; ++s) {
+    std::string f0, f1, f2, f3;
+    if (weights_are_buffer) {
+      if (gpu_info.SupportsPointersInKernels()) {
+        f0 = "weights_cache[" + std::to_string(s * 4 + 0) + "]";
+        f1 = "weights_cache[" + std::to_string(s * 4 + 1) + "]";
+        f2 = "weights_cache[" + std::to_string(s * 4 + 2) + "]";
+        f3 = "weights_cache[" + std::to_string(s * 4 + 3) + "]";
+      } else {
+        f0 = "f0";
+        f1 = "f1";
+        f2 = "f2";
+        f3 = "f3";
+      }
+    } else {
+      f0 = "f" + std::to_string(s * 4 + 0);
+      f1 = "f" + std::to_string(s * 4 + 1);
+      f2 = "f" + std::to_string(s * 4 + 2);
+      f3 = "f" + std::to_string(s * 4 + 3);
+    }
+    bool use_fma = gpu_info.IsAMD() && gpu_info.IsApiOpenCl();
+    if (GetWeightsDescription().IsI4O4()) {
+      switch (op_def.precision) {
+        case CalculationsPrecision::F32:
+        case CalculationsPrecision::F16:
+          if (use_fma) {
+            c += "#define CONV" + std::to_string(s) + "(R, S)    \\\n";
+            c += "R = fma(" + f0 + ", S.x, R); \\\n";
+            c += "R = fma(" + f1 + ", S.y, R); \\\n";
+            c += "R = fma(" + f2 + ", S.z, R); \\\n";
+            c += "R = fma(" + f3 + ", S.w, R);   \n";
+          } else {
+            c += "#define CONV" + std::to_string(s) + "(R, S)    \\\n";
+            c += "R += S.x * " + f0 + "; \\\n";
+            c += "R += S.y * " + f1 + "; \\\n";
+            c += "R += S.z * " + f2 + "; \\\n";
+            c += "R += S.w * " + f3 + ";   \n";
+          }
+          break;
+        case CalculationsPrecision::F32_F16:
+          c += "#define CONV" + std::to_string(s) + "(R, S) \\\n";
+          c += "R += TO_ACCUM_TYPE(S.x * " + f0 + " + S.y * " + f1 +
+               " + S.z * " + f2 + " + S.w * " + f3 + ");\n";
+          break;
+      }
+    } else {
+      // O4I4
+      c += "#define CONV" + std::to_string(s) + "(R, S)    \\\n";
+      c += "R.x += dot(S, " + f0 + "); \\\n";
+      c += "R.y += dot(S, " + f1 + "); \\\n";
+      c += "R.z += dot(S, " + f2 + "); \\\n";
+      c += "R.w += dot(S, " + f3 + ");   \n";
+    }
+  }
+
+  auto generate_id = [&](const std::string& x, const std::string& y,
+                         const std::string& z) {
+    std::string id;
+    if (src_def.HasAxis(Axis::WIDTH)) {
+      id += "_w" + x;
+    }
+    if (src_def.HasAxis(Axis::HEIGHT)) {
+      id += "_h" + y;
+    }
+    if (src_def.HasAxis(Axis::DEPTH)) {
+      id += "_d" + z;
+    }
+    return id;
+  };
+
+  auto generate_id_full = [&](const std::string& x, const std::string& y,
+                              const std::string& z, const std::string& s) {
+    return generate_id(x, y, z) + "_s" + s;
+  };
+
+  auto generate_check = [&](const std::string& x, const std::string& y,
+                            const std::string& z) {
+    std::string check;
+    const std::vector<Axis> axes{Axis::WIDTH, Axis::HEIGHT, Axis::DEPTH};
+    const std::vector<std::string> names{"in_x", "in_y", "in_z"};
+    const std::vector<std::string> coords{x, y, z};
+    for (int i = 0; i < axes.size(); ++i) {
+      const auto& axis = axes[i];
+      if (src_def.HasAxis(axis) && !src_def.SupportsZeroClamp(axis, gpu_info) &&
+          block_size[i] != 1) {
+        if (!check.empty()) {
+          check += " && ";
+        }
+        check += names[i] + coords[i];
+      }
+    }
+    return check;
+  };
+
+  c += "MAIN_FUNCTION($0) {\n";
+  if (op_def.IsBatchSupported()) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int dst_x = (linear_id / args.dst_tensor.Batch());\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int dst_x = GLOBAL_ID_0;\n";
+  }
+  c += "  int rem_x = dst_x % args.stride_x;\n";
+  c += "  int ceil_x = dst_x / args.stride_x;\n";
+  c += "  dst_x = ceil_x * args.stride_x * " + std::to_string(block_size.x) +
+       " + rem_x;\n";
+  if (src_def.HasAxis(Axis::DEPTH)) {
+    c += "  int linear_id_y = GLOBAL_ID_1;\n";
+    c += "  int dst_y = linear_id_y % args.grid_size_y;\n";
+    c += "  int dst_z = linear_id_y / args.grid_size_y;\n";
+    c += "  int rem_z = dst_z % args.stride_z;\n";
+    c += "  int ceil_z = dst_z / args.stride_z;\n";
+    c += "  dst_z = ceil_z * args.stride_z * " + std::to_string(block_size.z) +
+         " + rem_z;\n";
+    c += "  if (dst_z >= args.dst_tensor.Depth()) return;\n";
+  } else {
+    c += "  int dst_y = GLOBAL_ID_1;\n";
+  }
+  c += "  int rem_y = dst_y % args.stride_y;\n";
+  c += "  int ceil_y = dst_y / args.stride_y;\n";
+  c += "  dst_y = ceil_y * args.stride_y * " + std::to_string(block_size.y) +
+       " + rem_y;\n";
+  c += "  int dst_s = GLOBAL_ID_2 * " + std::to_string(block_size.w) + ";\n";
+  c += "  if (dst_x >= args.dst_tensor.Width() || dst_y >= "
+       "args.dst_tensor.Height() || dst_s >= "
+       "args.dst_tensor.Slices()) return;\n";
+  if (weights_are_buffer) {
+    c += "  int f_base = dst_s * args.src_tensor.Slices() * args.kernel_size_x "
+         "* args.kernel_size_y";
+    if (src_def.HasAxis(Axis::DEPTH)) {
+      c += " * args.kernel_size_z";
+    }
+    c += " * 4;\n";
+  }
+  for (int s = 0; s < block_size.w; ++s) {
+    const std::string sind = std::to_string(s);
+    for (int z = 0; z < block_size.z; ++z) {
+      const std::string zind = std::to_string(z);
+      for (int y = 0; y < block_size.y; ++y) {
+        const std::string yind = std::to_string(y);
+        for (int x = 0; x < block_size.x; ++x) {
+          const std::string xind = std::to_string(x);
+          c += "  ACCUM_FLT4 r" + generate_id_full(xind, yind, zind, sind) +
+               " = INIT_ACCUM_FLT4(0.0f);\n";
+        }
+      }
+    }
+  }
+  c += "  int kernel_first_dst_x = dst_x + args.padding_x;\n";
+  c += "  int kernel_first_dst_y = dst_y + args.padding_y;\n";
+  c += "  int kernel_last_dst_x = kernel_first_dst_x - args.kernel_size_x;\n";
+  c += "  int kernel_last_dst_y = kernel_first_dst_y - args.kernel_size_y;\n";
+  c += "  int offset_x = abs(args.padding_x);\n";
+  c += "  int offset_x_strided = offset_x * args.stride_x;\n";
+  c +=
+      "  int src_x = (kernel_first_dst_x + offset_x_strided) / args.stride_x - "
+      "offset_x;\n";
+  c += "  int offset_y = abs(args.padding_y);\n";
+  c += "  int offset_y_strided = offset_y * args.stride_y;\n";
+  c +=
+      "  int src_y = (kernel_first_dst_y + offset_y_strided) / args.stride_y - "
+      "offset_y;\n";
+  if (src_def.HasAxis(Axis::DEPTH)) {
+    c += "  int kernel_first_dst_z = dst_z + args.padding_z;\n";
+    c += "  int kernel_last_dst_z = kernel_first_dst_z - args.kernel_size_z;\n";
+    c += "  int offset_z = abs(args.padding_z);\n";
+    c += "  int offset_z_strided = offset_z * args.stride_z;\n";
+    c += "  int src_z = (kernel_first_dst_z + offset_z_strided) / "
+         "args.stride_z - offset_z;\n";
+    c += "  int src_as_dst_z = src_z * args.stride_z;\n";
+    c +=
+        "  for (;src_as_dst_z > kernel_last_dst_z; src_z -= 1, src_as_dst_z -= "
+        "args.stride_z) {\n";
+    for (int z = 0; z < block_size.z; ++z) {
+      const std::string zindex = std::to_string(z);
+      c += "    int sz" + zindex + " = src_z + " + zindex + ";\n";
+      if (!src_def.SupportsZeroClamp(Axis::DEPTH, gpu_info)) {
+        c += "    bool in_z" + zindex + " = sz" + zindex + " >= 0 && sz" +
+             zindex + " < args.src_tensor.Depth();\n";
+        if (!src_def.CanReadOutOfBorder(Axis::DEPTH)) {
+          c += "    sz" + zindex + " = clamp(sz" + zindex +
+               ", 0, args.src_tensor.Depth() - 1);\n";
+        }
+      }
+    }
+    if (block_size.z == 1 &&
+        !src_def.SupportsZeroClamp(Axis::DEPTH, gpu_info)) {
+      c += "    if (!in_z0) continue;\n";
+    }
+    c += "    int kernel_z = kernel_first_dst_z - src_as_dst_z;\n";
+    c += "    int src_as_dst_y = src_y * args.stride_y;\n";
+    c += "    int src_y_copy = src_y;\n";
+    c += "    for (;src_as_dst_y > kernel_last_dst_y; src_y_copy -= 1, "
+         "src_as_dst_y -= args.stride_y) {\n";
+  } else {
+    c += "  int src_as_dst_y = src_y * args.stride_y;\n";
+    c += "  for (;src_as_dst_y > kernel_last_dst_y; src_y -= 1, src_as_dst_y "
+         "-= args.stride_y) {\n";
+  }
+  for (int y = 0; y < block_size.y; ++y) {
+    const std::string yindex = std::to_string(y);
+    const std::string src_y =
+        src_def.HasAxis(Axis::DEPTH) ? "src_y_copy" : "src_y";
+    c += "    int sy" + yindex + " = " + src_y + " + " + yindex + ";\n";
+    if (!src_def.SupportsZeroClamp(Axis::HEIGHT, gpu_info)) {
+      c += "    bool in_y" + yindex + " = sy" + yindex + " >= 0 && sy" +
+           yindex + " < args.src_tensor.Height();\n";
+      if (!src_def.CanReadOutOfBorder(Axis::HEIGHT)) {
+        c += "    sy" + yindex + " = clamp(sy" + yindex +
+             ", 0, args.src_tensor.Height() - 1);\n";
+      }
+    }
+  }
+  if (block_size.y == 1 && !src_def.SupportsZeroClamp(Axis::HEIGHT, gpu_info)) {
+    c += "      if (!in_y0) continue;\n";
+  }
+  c += "    int kernel_y = kernel_first_dst_y - src_as_dst_y;\n";
+  c += "    int src_as_dst_x = src_x * args.stride_x;\n";
+  c += "    int src_x_copy = src_x;\n";
+  c += "    for (;src_as_dst_x > kernel_last_dst_x; src_x_copy -= 1, "
+       "src_as_dst_x "
+       "-= args.stride_x) {\n";
+  for (int x = 0; x < block_size.x; ++x) {
+    const std::string xindex = std::to_string(x);
+    c += "      int sx" + xindex + " = src_x_copy + " + xindex + ";\n";
+    if (!src_def.SupportsZeroClamp(Axis::WIDTH, gpu_info)) {
+      c += "      bool in_x" + xindex + " = sx" + xindex + " >= 0 && sx" +
+           xindex + " < args.src_tensor.Width();\n";
+      if (!src_def.CanReadOutOfBorder(Axis::WIDTH)) {
+        c += "      sx" + xindex + " = clamp(sx" + xindex +
+             ", 0, args.src_tensor.Width() - 1);\n";
+      }
+    }
+  }
+  if (block_size.x == 1 && !src_def.SupportsZeroClamp(Axis::WIDTH, gpu_info)) {
+    c += "      if (!in_x0) continue;\n";
+  }
+  for (int z = 0; z < block_size.z; ++z) {
+    const std::string zind = std::to_string(z);
+    for (int y = 0; y < block_size.y; ++y) {
+      const std::string yind = std::to_string(y);
+      for (int x = 0; x < block_size.x; ++x) {
+        const std::string xind = std::to_string(x);
+        const std::string id = generate_id(xind, yind, zind);
+        const std::string check = generate_check(xind, yind, zind);
+        std::string coords = "sx" + xind + ", sy" + yind;
+        if (src_def.HasAxis(Axis::DEPTH)) {
+          coords += ", sz" + zind;
+        }
+        if (src_def.IsLinear()) {
+          c += "      int addr" + id + " = args.src_tensor.GetAddress(" +
+               coords + ", 0);\n";
+          if (src_def.ReturnsZeroForNegOneRead(gpu_info)) {
+            c += "      addr" + id + " = select(-1, addr" + id + ", (" + check +
+                 "));\n";
+            c += "      int ds" + id +
+                 " = select(0, args.src_tensor.SliceStride(), (" + check +
+                 "));\n";
+          }
+        }
+      }
+    }
+  }
+  if (src_def.IsLinear() && !src_def.ReturnsZeroForNegOneRead(gpu_info)) {
+    c += "      int ds = args.src_tensor.SliceStride();\n";
+  }
+  c += "      int kernel_x = kernel_first_dst_x - src_as_dst_x;\n";
+  if (src_def.HasAxis(Axis::DEPTH)) {
+    c += "      int kernel_index = (kernel_z * args.kernel_size_y + kernel_y) "
+         "*  args.kernel_size_x + kernel_x;\n";
+  } else {
+    c += "      int kernel_index = kernel_y * args.kernel_size_x + kernel_x;\n";
+  }
+  if (weights_are_buffer) {
+    c += "      int f_offset = f_base + kernel_index * "
+         "args.src_tensor.Slices() * " +
+         std::to_string(block_size.w * 4) + ";\n";
+  } else {
+    c += "      int x_c = kernel_index * args.src_tensor.Slices();\n";
+  }
+  c += "      for (int s = 0; s < args.src_tensor.Slices(); ++s) {\n";
+  const bool conditional_read = gpu_info.IsMali();
+  for (int z = 0; z < block_size.z; ++z) {
+    const std::string zind = std::to_string(z);
+    for (int y = 0; y < block_size.y; ++y) {
+      const std::string yind = std::to_string(y);
+      for (int x = 0; x < block_size.x; ++x) {
+        const std::string xind = std::to_string(x);
+        const std::string id = generate_id(xind, yind, zind);
+        std::string address;
+        if (src_def.IsLinear()) {
+          address = "addr" + id;
+        } else {
+          address = "sx" + xind + ", sy" + yind;
+          if (src_def.HasAxis(Axis::DEPTH)) {
+            address += ", sz" + zind;
+          }
+          address += ", s";
+        }
+        if (src_def.ReturnsZeroForNegOneRead(gpu_info)) {
+          c += "        FLT4 src" + id + " = args.src_tensor.Read(" + address +
+               "); " + address + " += ds" + id + ";\n";
+        } else {
+          const std::string check = generate_check(xind, yind, zind);
+          if (!check.empty()) {
+            if (conditional_read) {
+              c += "        FLT4 src" + id + " = " + check +
+                   " ? args.src_tensor.Read(" + address +
+                   ") : INIT_FLT4(0.0f);\n";
+            } else {
+              c += "        FLT4 src" + id + " = args.src_tensor.Read(" +
+                   address + ") * INIT_FLT(" + check + ");\n";
+            }
+          } else {
+            c += "        FLT4 src" + id + " = args.src_tensor.Read(" +
+                 address + ");\n";
+          }
+          if (src_def.IsLinear()) {
+            c += "        addr" + id + " += ds;\n";
+          }
+        }
+      }
+    }
+  }
+  if (weights_are_buffer) {
+    if (gpu_info.SupportsPointersInKernels()) {
+      c += "        __global FLT4* weights_cache = "
+           "args.weights.GetPtr(f_offset);\n";
+    }
+  } else {
+    for (int s = 0; s < block_size.w; ++s) {
+      c += absl::Substitute(
+          R"(        FLT4 f$1 = args.weights0.Read(dst_s + $0, x_c);
+        FLT4 f$2 = args.weights1.Read(dst_s + $0, x_c);
+        FLT4 f$3 = args.weights2.Read(dst_s + $0, x_c);
+        FLT4 f$4 = args.weights3.Read(dst_s + $0, x_c);
+)",
+          s, s * 4 + 0, s * 4 + 1, s * 4 + 2, s * 4 + 3);
+    }
+    c += "        x_c++;\n";
+  }
+  if (weights_are_buffer && !gpu_info.SupportsPointersInKernels()) {
+    c += "      FLT4 f0, f1, f2, f3;\n";
+  }
+  for (int s = 0; s < block_size.w; ++s) {
+    if (weights_are_buffer && !gpu_info.SupportsPointersInKernels()) {
+      c += "        f0 = args.weights.Read(f_offset + " +
+           std::to_string(s * 4 + 0) + ");\n";
+      c += "        f1 = args.weights.Read(f_offset + " +
+           std::to_string(s * 4 + 1) + ");\n";
+      c += "        f2 = args.weights.Read(f_offset + " +
+           std::to_string(s * 4 + 2) + ");\n";
+      c += "        f3 = args.weights.Read(f_offset + " +
+           std::to_string(s * 4 + 3) + ");\n";
+    }
+    const std::string sind = std::to_string(s);
+    for (int z = 0; z < block_size.z; ++z) {
+      const std::string zind = std::to_string(z);
+      for (int y = 0; y < block_size.y; ++y) {
+        const std::string yind = std::to_string(y);
+        for (int x = 0; x < block_size.x; ++x) {
+          const std::string xind = std::to_string(x);
+          const std::string id = generate_id(xind, yind, zind);
+          const std::string full_id = generate_id_full(xind, yind, zind, sind);
+          c += "        CONV" + sind + "(r" + full_id + ", src" + id + ");\n";
+        }
+      }
+    }
+  }
+  if (weights_are_buffer) {
+    c += "        f_offset += " + std::to_string(block_size.w * 4) + ";\n";
+  }
+  c += "      }\n";
+  c += "    }\n";
+  c += "  }\n";
+  if (src_def.HasAxis(Axis::DEPTH)) {
+    c += "  }\n";
+  }
+  for (int s = 0; s < block_size.w; ++s) {
+    const std::string sind = std::to_string(s);
+    c += "  if (dst_s < args.dst_tensor.Slices()) {\n";
+    c += "    FLT4 bias_val = args.biases.Read(dst_s);\n";
+    for (int z = 0; z < block_size.z; ++z) {
+      const std::string zind = std::to_string(z);
+      for (int y = 0; y < block_size.y; ++y) {
+        const std::string yind = std::to_string(y);
+        for (int x = 0; x < block_size.x; ++x) {
+          const std::string xind = std::to_string(x);
+          const std::string id = generate_id_full(xind, yind, zind, sind);
+          std::string checks =
+              "xc < args.dst_tensor.Width() && yc < args.dst_tensor.Height()";
+          std::string coords = "xc, yc";
+          c += "    {\n";
+          c += "      int xc = dst_x + args.stride_x * " + xind + ";\n";
+          c += "      int yc = dst_y + args.stride_y * " + yind + ";\n";
+          if (src_def.HasAxis(Axis::DEPTH)) {
+            c += "      int zc = dst_z + args.stride_z * " + zind + ";\n";
+            checks += " && zc < args.dst_tensor.Depth()";
+            coords += ", zc";
+          }
+          c += "      if (" + checks + ") {\n";
+          c += "        FLT4 res = TO_FLT4(r" + id + ") + bias_val;\n";
+          c += "        args.dst_tensor.Write(res, " + coords + ", dst_s);\n";
+          c += "      }\n";
+          c += "    }\n";
+        }
+      }
+    }
+    c += "  }\n";
+    c += "  dst_s++;\n";
+  }
+  c += "}\n";
+  return c;
+}
+
+absl::Status ConvolutionTransposed::BindArguments(ArgumentsBinder* args) {
+  if (definition_.src_tensors[0].HasAxis(Axis::DEPTH)) {
+    const int aligned_h =
+        AlignByN(dst_[0]->Height(), stride_.y * block_size_.y);
+    RETURN_IF_ERROR(
+        args->SetInt("grid_size_y", DivideRoundUp(aligned_h, block_size_.y)));
+  }
+  return absl::OkStatus();
+}
+
+int3 ConvolutionTransposed::GetGridSize() const {
+  const int aligned_w = AlignByN(dst_[0]->Width(), stride_.x * block_size_.x);
+  const int aligned_h = AlignByN(dst_[0]->Height(), stride_.y * block_size_.y);
+  const int aligned_d = AlignByN(dst_[0]->Depth(), stride_.z * block_size_.z);
+  const int grid_x = DivideRoundUp(aligned_w, block_size_.x) * dst_[0]->Batch();
+  const int grid_y = DivideRoundUp(aligned_h, block_size_.y) *
+                     DivideRoundUp(aligned_d, block_size_.z);
+  const int grid_z = DivideRoundUp(dst_[0]->Slices(), block_size_.w);
+  return int3(grid_x, grid_y, grid_z);
+}
+
+void ConvolutionTransposed::GetPossibleKernelWorkGroups(
+    TuningType tuning_type, const GpuInfo& gpu_info,
+    const KernelInfo& kernel_info, std::vector<int3>* work_groups) const {
+  GetPossibleWorkGroupsConv(tuning_type, gpu_info, kernel_info, grid_size_,
+                            work_groups);
+}
+
+ConvolutionTransposed CreateConvolutionTransposed(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr) {
+  ConvolutionTransposed result(definition, attr, gpu_info);
+  result.UploadWeights(attr.weights, UseBufferForWeights(gpu_info));
+
+  TensorDescriptor bias_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), attr.bias);
+  result.args_.AddObject("biases", std::make_unique<TensorDescriptor>(
+                                       std::move(bias_tensor_desc)));
+  return result;
+}
+
+ConvolutionTransposed CreateConvolutionTransposed3D(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposed3DAttributes& attr) {
+  ConvolutionTransposed result(definition, attr, gpu_info);
+  result.UploadWeights(attr.weights, UseBufferForWeights(gpu_info));
+
+  TensorDescriptor bias_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), attr.bias);
+  result.args_.AddObject("biases", std::make_unique<TensorDescriptor>(
+                                       std::move(bias_tensor_desc)));
+  return result;
+}
+
+ConvolutionTransposed CreateConvolutionTransposedDynamicWeights(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr) {
+  OperationDef new_def = definition;
+  new_def.src_tensors = {
+      definition.src_tensors[0]};  // leaving only src_tensor def, weights defs
+                                   // will be added later
+  const DataType weights_type = definition.GetDataType();
+  if (UseBufferForWeights(gpu_info)) {
+    // add 1 src_tensor(buffer) for weights
+    new_def.src_tensors.push_back(
+        {weights_type, TensorStorageType::BUFFER, Layout::HWC});
+  } else {
+    // add 4 src_tensors(4X textures 2d) for weights
+    new_def.src_tensors.push_back(
+        {weights_type, TensorStorageType::TEXTURE_2D, Layout::HW});
+    new_def.src_tensors.push_back(
+        {weights_type, TensorStorageType::TEXTURE_2D, Layout::HW});
+    new_def.src_tensors.push_back(
+        {weights_type, TensorStorageType::TEXTURE_2D, Layout::HW});
+    new_def.src_tensors.push_back(
+        {weights_type, TensorStorageType::TEXTURE_2D, Layout::HW});
+  }
+  ConvolutionTransposed result(new_def, attr, gpu_info);
+
+  TensorDescriptor bias_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), attr.bias);
+  result.args_.AddObject("biases", std::make_unique<TensorDescriptor>(
+                                       std::move(bias_tensor_desc)));
+  return result;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed.h

@@ -0,0 +1,175 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_H_
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/buffer_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tensor_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/weights_conversion.h"
+#include "tensorflow/lite/delegates/gpu/common/task/weights_layout.h"
+#include "tensorflow/lite/delegates/gpu/common/tensor.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+class ConvolutionTransposed : public GPUOperation {
+ public:
+  ConvolutionTransposed() = default;
+  void GetPossibleKernelWorkGroups(
+      TuningType tuning_type, const GpuInfo& gpu_info,
+      const KernelInfo& kernel_info,
+      std::vector<int3>* work_groups) const override;
+  absl::Status BindArguments(ArgumentsBinder* args) override;
+  int3 GetGridSize() const override;
+
+  // Move only
+  ConvolutionTransposed(ConvolutionTransposed&& operation) = default;
+  ConvolutionTransposed& operator=(ConvolutionTransposed&& operation) = default;
+  ConvolutionTransposed(const ConvolutionTransposed&) = delete;
+  ConvolutionTransposed& operator=(const ConvolutionTransposed&) = delete;
+
+  WeightsDescription GetWeightsDescription() const {
+    WeightsDescription desc;
+    desc.type = DeduceDataTypeFromPrecision(definition_.precision);
+    desc.layout = weights_layout_;
+    desc.output_group_size = block_size_.w;
+    return desc;
+  }
+
+ private:
+  friend ConvolutionTransposed CreateConvolutionTransposed(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const ConvolutionTransposedAttributes& attr);
+  friend ConvolutionTransposed CreateConvolutionTransposed3D(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const ConvolutionTransposed3DAttributes& attr);
+  friend ConvolutionTransposed CreateConvolutionTransposedDynamicWeights(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const ConvolutionTransposedAttributes& attr);
+
+  ConvolutionTransposed(const OperationDef& definition,
+                        const ConvolutionTransposedAttributes& attr,
+                        const GpuInfo& gpu_info);
+  ConvolutionTransposed(const OperationDef& definition,
+                        const ConvolutionTransposed3DAttributes& attr,
+                        const GpuInfo& gpu_info);
+
+  template <DataType T>
+  void UploadWeights(const tflite::gpu::Tensor<OHWI, T>& weights,
+                     bool weights_are_buffer);
+
+  template <DataType T>
+  void UploadWeights(const tflite::gpu::Tensor<OHWDI, T>& weights,
+                     bool weights_are_buffer);
+
+  std::string GenerateConvolutionTransposedCode(const OperationDef& op_def,
+                                                const GpuInfo& gpu_info,
+                                                const int4& block_size);
+  int4 stride_;
+  int4 block_size_ = int4(1, 1, 1, 1);  // WHDS
+  WeightsLayout weights_layout_;
+};
+
+template <DataType T>
+void ConvolutionTransposed::UploadWeights(
+    const tflite::gpu::Tensor<OHWI, T>& weights, bool weights_are_buffer) {
+  const auto weights_desc = GetWeightsDescription();
+  const int flt_count =
+      GetTotalElementsCountForLayout(weights_desc, weights.shape);
+
+  std::vector<uint8_t> weights_data(flt_count * SizeOf(weights_desc.type));
+  RearrangeWeights(weights, weights_desc, absl::MakeSpan(weights_data));
+
+  if (weights_are_buffer) {
+    BufferDescriptor desc;
+    desc.element_type = weights_desc.type;
+    desc.element_size = 4;
+    desc.size = weights_data.size();
+    desc.data = std::move(weights_data);
+    args_.AddObject("weights",
+                    std::make_unique<BufferDescriptor>(std::move(desc)));
+  } else {
+    uint2 tex_size = Get2dResourceSize(weights_desc, weights.shape);
+    int sub_size = SizeOf(weights_desc.type) * 4 * tex_size.x * tex_size.y;
+    for (int i = 0; i < 4; ++i) {
+      TensorDescriptor desc = CreateConstantHWVec4TensorDescriptor(
+          weights_desc.type, TensorStorageType::TEXTURE_2D, tex_size.x,
+          tex_size.y, weights_data.data() + sub_size * i);
+      args_.AddObject("weights" + std::to_string(i),
+                      std::make_unique<TensorDescriptor>(std::move(desc)));
+    }
+  }
+}
+
+template <DataType T>
+void ConvolutionTransposed::UploadWeights(
+    const tflite::gpu::Tensor<OHWDI, T>& weights, bool weights_are_buffer) {
+  const auto weights_desc = GetWeightsDescription();
+  const int flt_count =
+      GetTotalElementsCountForLayout(weights_desc, weights.shape);
+
+  std::vector<uint8_t> weights_data(flt_count * SizeOf(weights_desc.type));
+  RearrangeWeights(weights, weights_desc, absl::MakeSpan(weights_data));
+
+  if (weights_are_buffer) {
+    BufferDescriptor desc;
+    desc.element_type = weights_desc.type;
+    desc.element_size = 4;
+    desc.size = weights_data.size();
+    desc.data = std::move(weights_data);
+    args_.AddObject("weights",
+                    std::make_unique<BufferDescriptor>(std::move(desc)));
+  } else {
+    uint2 tex_size = Get2dResourceSize(weights_desc, weights.shape);
+    int sub_size = SizeOf(weights_desc.type) * 4 * tex_size.x * tex_size.y;
+    for (int i = 0; i < 4; ++i) {
+      TensorDescriptor desc = CreateConstantHWVec4TensorDescriptor(
+          weights_desc.type, TensorStorageType::TEXTURE_2D, tex_size.x,
+          tex_size.y, weights_data.data() + sub_size * i);
+      args_.AddObject("weights" + std::to_string(i),
+                      std::make_unique<TensorDescriptor>(std::move(desc)));
+    }
+  }
+}
+
+ConvolutionTransposed CreateConvolutionTransposed(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr);
+
+ConvolutionTransposed CreateConvolutionTransposed3D(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposed3DAttributes& attr);
+
+ConvolutionTransposed CreateConvolutionTransposedDynamicWeights(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_H_

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_3x3.cc

@@ -0,0 +1,448 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_3x3.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
+
+namespace tflite {
+namespace gpu {
+
+ConvolutionTransposed3x3::ConvolutionTransposed3x3(
+    const OperationDef& definition, const GpuInfo& gpu_info, int2 padding)
+    : GPUOperation(definition), padding_(padding) {
+  work_group_size_ = int3(8, 4, 1);
+  work_group_launch_order_ = int3(2, 0, 1);
+  if (gpu_info.IsApple()) {
+    if (gpu_info.apple_info.IsBionic()) {
+      weights_upload_type_ = WeightsUploadType::GLOBAL_MEM;
+    } else {
+      weights_upload_type_ = WeightsUploadType::LOCAL_MEM_BY_THREADS;
+    }
+  } else if (gpu_info.IsPowerVR()) {
+    weights_upload_type_ = WeightsUploadType::LOCAL_MEM_ASYNC;
+  } else if (gpu_info.IsNvidia() || gpu_info.IsIntel()) {
+    weights_upload_type_ = WeightsUploadType::LOCAL_MEM_BY_THREADS;
+  } else if (gpu_info.IsAMD()) {
+    weights_upload_type_ = WeightsUploadType::CONSTANT_MEM;
+  } else {
+    weights_upload_type_ = WeightsUploadType::GLOBAL_MEM;
+  }
+  if (gpu_info.IsApple()) {
+    weights_layout_ = WeightsLayout::kOICustomSpatialO4I4;
+  } else {
+    weights_layout_ = WeightsLayout::kOICustomSpatialI4O4;
+  }
+  code_ = GenerateConvolutionTransposedCode(gpu_info, definition_,
+                                            weights_upload_type_, padding_,
+                                            work_group_launch_order_);
+  if (definition_.precision == CalculationsPrecision::F16 &&
+      gpu_info.IsPowerVR()) {
+    compiler_options_.push_back(CompilerOptions::kClFastRelaxedMath);
+  }
+}
+
+std::string ConvolutionTransposed3x3::GenerateConvolutionTransposedCode(
+    const GpuInfo& gpu_info, const OperationDef& op_def,
+    ConvolutionTransposed3x3::WeightsUploadType weights_upload_type,
+    int2 padding, int3 work_group_launch_order) {
+  auto src_desc = op_def.src_tensors[0];
+  AddSrcTensor("src_tensor", src_desc);
+  AddDstTensor("dst_tensor", op_def.src_tensors[0]);
+
+  if (op_def.src_tensors.size() == 2) {
+    // dynamic weights
+    BufferDescriptor desc;
+    desc.element_type = op_def.src_tensors[1].GetDataType();
+    desc.element_size = 4;
+    desc.memory_type =
+        weights_upload_type ==
+                ConvolutionTransposed3x3::WeightsUploadType::CONSTANT_MEM
+            ? MemoryType::CONSTANT
+            : MemoryType::GLOBAL;
+    AddSrcBuffer("weights", desc);
+  }
+
+  args_.AddInt("filter_offset");
+  args_.AddInt("padding_x");
+  args_.AddInt("padding_y");
+
+  const bool need_local_mem =
+      weights_upload_type ==
+          ConvolutionTransposed3x3::WeightsUploadType::LOCAL_MEM_BY_THREADS ||
+      weights_upload_type ==
+          ConvolutionTransposed3x3::WeightsUploadType::LOCAL_MEM_ASYNC;
+
+  std::string c;
+  if (GetWeightsDescription().IsI4O4()) {
+    switch (op_def.precision) {
+      case CalculationsPrecision::F32:
+      case CalculationsPrecision::F16:
+        c += "#define CONV(R, SRC, F) \\\n";
+        c += "  R += SRC.x * weights_cache[F]; \\\n";
+        c += "  R += SRC.y * weights_cache[F + 1]; \\\n";
+        c += "  R += SRC.z * weights_cache[F + 2]; \\\n";
+        c += "  R += SRC.w * weights_cache[F + 3];   \n";
+        break;
+      case CalculationsPrecision::F32_F16:
+        c += "#define CONV(R, SRC, F) \\\n";
+        c += "  R += TO_ACCUM_TYPE(SRC.x * weights_cache[F] + SRC.y * "
+             "weights_cache[F + 1] + SRC.z * weights_cache[F + 2] + SRC.w * "
+             "weights_cache[F + 3]);\n";
+        break;
+    }
+  } else {
+    // O4I4
+    c += "#define CONV(R, SRC, F) \\\n";
+    c += "  R.x += dot(SRC, weights_cache[F]); \\\n";
+    c += "  R.y += dot(SRC, weights_cache[F + 1]); \\\n";
+    c += "  R.z += dot(SRC, weights_cache[F + 2]); \\\n";
+    c += "  R.w += dot(SRC, weights_cache[F + 3]);   \n";
+  }
+
+  const int wg_total_size =
+      work_group_size_.x * work_group_size_.y * work_group_size_.z;
+  const std::string barrier =
+      wg_total_size == 32 && gpu_info.IsWaveSizeEqualTo32()
+          ? "SIMD_LOCAL_MEM_BARRIER"
+          : "LOCAL_MEM_BARRIER";
+  const std::string weights_space =
+      weights_upload_type ==
+              ConvolutionTransposed3x3::WeightsUploadType::CONSTANT_MEM
+          ? "__constant"
+          : "__global";
+
+  if (gpu_info.IsApiOpenCl()) {
+    c += "__attribute__((reqd_work_group_size(8, 4, 1)))\n";
+  }
+  c += "MAIN_FUNCTION($0) {\n";
+  int3 launch_remap;
+  launch_remap[work_group_launch_order.x] = 0;
+  launch_remap[work_group_launch_order.y] = 1;
+  launch_remap[work_group_launch_order.z] = 2;
+  auto GetGlobalID = [&](int id) {
+    std::string result;
+    const std::string sid = std::to_string(id);
+    if (work_group_launch_order[id] == id) {
+      return "GLOBAL_ID_" + sid;
+    } else {
+      return "GROUP_ID_" + std::to_string(launch_remap[id]) + " * GROUP_SIZE_" +
+             sid + " + LOCAL_ID_" + sid;
+    }
+  };
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = " + GetGlobalID(0) + ";\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = " + GetGlobalID(0) + ";\n";
+  }
+  c += "  int DST_X = X * 2;\n";
+  c += "  int SRC_X = X + args.padding_x;\n";
+  c += "  int Y = " + GetGlobalID(1) + ";\n";
+  c += "  int DST_Y = Y * 2;\n";
+  c += "  int SRC_Y = Y + args.padding_y;\n";
+  c += "  int Z = " + GetGlobalID(2) + ";\n";
+  if (!need_local_mem) {
+    c += "  if (DST_X >= args.dst_tensor.Width() || DST_Y >= "
+         "args.dst_tensor.Height() || Z >= args.dst_tensor.Slices()) return;\n";
+  }
+  c += "  ACCUM_FLT4 r0 = INIT_ACCUM_FLT4(0.0f);\n";
+  c += "  ACCUM_FLT4 r1 = INIT_ACCUM_FLT4(0.0f);\n";
+  c += "  ACCUM_FLT4 r2 = INIT_ACCUM_FLT4(0.0f);\n";
+  c += "  ACCUM_FLT4 r3 = INIT_ACCUM_FLT4(0.0f);\n";
+  c += "  int f_offset = Z * args.filter_offset;\n";
+  if (need_local_mem) {
+    c += "  __local FLT4 weights_cache[36];\n";
+  }
+  if (weights_upload_type ==
+      ConvolutionTransposed3x3::WeightsUploadType::LOCAL_MEM_BY_THREADS) {
+    c += "  int local_id = LOCAL_ID_1 * 8 + LOCAL_ID_0;\n";
+  }
+  if (!src_desc.SupportsZeroClamp(Axis::WIDTH, gpu_info)) {
+    c += "  bool in_x0 = SRC_X >= 0 && SRC_X < args.src_tensor.Width();\n";
+    c += "  bool in_x1 = SRC_X + 1 >= 0 && SRC_X + 1 < "
+         "args.src_tensor.Width();\n";
+  }
+  if (!src_desc.SupportsZeroClamp(Axis::HEIGHT, gpu_info)) {
+    c += "  bool in_y0 = SRC_Y >= 0 && SRC_Y < args.src_tensor.Height();\n";
+    c += "  bool in_y1 = SRC_Y + 1 >= 0 && SRC_Y + 1 < "
+         "args.src_tensor.Height();\n";
+  }
+  auto generate_check = [&](int x, int y) {
+    std::string check;
+    const std::vector<Axis> axes{Axis::WIDTH, Axis::HEIGHT};
+    const std::vector<std::string> names{"in_x" + std::to_string(x),
+                                         "in_y" + std::to_string(y)};
+    for (int i = 0; i < axes.size(); ++i) {
+      const auto& axis = axes[i];
+      if (src_desc.HasAxis(axis) &&
+          !src_desc.SupportsZeroClamp(axis, gpu_info)) {
+        if (!check.empty()) {
+          check += " && ";
+        }
+        check += names[i];
+      }
+    }
+    return check;
+  };
+  if (src_desc.IsLinear()) {
+    if (src_desc.ReturnsZeroForNegOneRead(gpu_info)) {
+      c += "  int addr_0 = args.src_tensor.GetAddress(SRC_X, SRC_Y, 0);\n";
+      c += "  int addr_1 = args.src_tensor.GetAddress(SRC_X + 1, SRC_Y, 0);\n";
+      c += "  int addr_2 = args.src_tensor.GetAddress(SRC_X, SRC_Y + 1, 0);\n";
+      c += "  int addr_3 = args.src_tensor.GetAddress(SRC_X+1, SRC_Y+1, 0);\n";
+      c += "  addr_0 = select(-1, addr_0, (in_x0 && in_y0));\n";
+      c += "  addr_1 = select(-1, addr_1, (in_x1 && in_y0));\n";
+      c += "  addr_2 = select(-1, addr_2, (in_x0 && in_y1));\n";
+      c += "  addr_3 = select(-1, addr_3, (in_x1 && in_y1));\n";
+      c += "  int dz_0 = select(0, args.src_tensor.SliceStride(), (in_x0 && "
+           "in_y0));\n";
+      c += "  int dz_1 = select(0, args.src_tensor.SliceStride(), (in_x1 && "
+           "in_y0));\n";
+      c += "  int dz_2 = select(0, args.src_tensor.SliceStride(), (in_x0 && "
+           "in_y1));\n";
+      c += "  int dz_3 = select(0, args.src_tensor.SliceStride(), (in_x1 && "
+           "in_y1));\n";
+    } else {
+      c += "  int xc0 = clamp(SRC_X, 0, args.src_tensor.Width() - 1);\n";
+      c += "  int xc1 = clamp(SRC_X + 1, 0, args.src_tensor.Width() - 1);\n";
+      c += "  int yc0 = clamp(SRC_Y, 0, args.src_tensor.Height() - 1);\n";
+      c += "  int yc1 = clamp(SRC_Y + 1, 0, args.src_tensor.Height() - 1);\n";
+      c += "  int addr_0 = args.src_tensor.GetAddress(xc0, yc0, 0);\n";
+      c += "  int addr_1 = args.src_tensor.GetAddress(xc1, yc0, 0);\n";
+      c += "  int addr_2 = args.src_tensor.GetAddress(xc0, yc1, 0);\n";
+      c += "  int addr_3 = args.src_tensor.GetAddress(xc1, yc1, 0);\n";
+      c += "  int dz = args.src_tensor.SliceStride();\n";
+    }
+  }
+  auto read_src = [&](int x, int y) {
+    if (src_desc.IsLinear()) {
+      const std::string id = std::to_string(y * 2 + x);
+      const std::string addr = "addr_" + std::to_string(y * 2 + x);
+      if (src_desc.ReturnsZeroForNegOneRead(gpu_info)) {
+        return "args.src_tensor.Read(" + addr + "); " + addr + " += dz_" + id +
+               ";\n";
+      } else {
+        return "args.src_tensor.Read(" + addr + ") * INIT_FLT(in_x" +
+               std::to_string(x) + " && in_y" + std::to_string(y) + "); " +
+               addr + " += dz;\n";
+      }
+    } else {
+      std::string check = generate_check(x, y);
+      if (!check.empty()) {
+        check = " * INIT_FLT(" + check + ")";
+      }
+      return "args.src_tensor.Read(SRC_X + " + std::to_string(x) +
+             ", SRC_Y + " + std::to_string(y) + ", s)" + check + ";\n";
+    }
+  };
+  const int padding_x_rem = abs(padding.x) % 2;
+  const int padding_y_rem = abs(padding.y) % 2;
+  std::vector<std::pair<int, int>> permutation;
+  if (padding_x_rem == 1 && padding_y_rem == 1) {
+    permutation = {{0, 0}, {1, 0}, {1, 1}, {2, 0}, {2, 2},
+                   {3, 0}, {3, 1}, {3, 2}, {3, 3}};
+  } else if (padding_x_rem == 0 && padding_y_rem == 1) {
+    permutation = {{0, 0}, {0, 1}, {1, 1}, {2, 0}, {2, 1},
+                   {2, 2}, {2, 3}, {3, 1}, {3, 3}};
+  } else if (padding_x_rem == 1 && padding_y_rem == 0) {
+    permutation = {{0, 0}, {0, 2}, {1, 0}, {1, 1}, {1, 2},
+                   {1, 3}, {2, 2}, {3, 2}, {3, 3}};
+  } else {  // padding_x_rem == 0 && padding_y_rem == 0
+    permutation = {{0, 0}, {0, 1}, {0, 2}, {0, 3}, {1, 1},
+                   {1, 3}, {2, 2}, {2, 3}, {3, 3}};
+  }
+  c += "  for (int s = 0; s < args.src_tensor.Slices(); ++s) {\n";
+  if (need_local_mem) {
+    c += "    " + barrier + ";\n";
+  }
+  if (weights_upload_type ==
+      ConvolutionTransposed3x3::WeightsUploadType::LOCAL_MEM_ASYNC) {
+    c += "    async_work_group_copy(weights_cache, "
+         "args.weights.GetPtr(f_offset), 36, "
+         "0);\n";
+  } else if (weights_upload_type ==
+             ConvolutionTransposed3x3::WeightsUploadType::
+                 LOCAL_MEM_BY_THREADS) {
+    c += "    weights_cache[local_id] = args.weights.Read(f_offset + "
+         "local_id);\n";
+    c += "    if (local_id < 4) {\n";
+    c += "      weights_cache[local_id + 32] = args.weights.Read(f_offset + "
+         "local_id + "
+         "32);\n";
+    c += "    };\n";
+  } else {  // GLOBAL_MEM/CONSTANT_MEM
+    c += "    " + weights_space +
+         " FLT4* weights_cache = args.weights.GetPtr(f_offset);\n";
+  }
+  c += "    FLT4 src0 = " + read_src(0, 0);
+  c += "    FLT4 src1 = " + read_src(1, 0);
+  c += "    FLT4 src2 = " + read_src(0, 1);
+  c += "    FLT4 src3 = " + read_src(1, 1);
+  c += "    f_offset += 36;\n";
+  if (need_local_mem) {
+    c += "    " + barrier + ";\n";
+  }
+  for (int i = 0; i < 9; ++i) {
+    const std::string r_name = "r" + std::to_string(permutation[i].first);
+    const std::string s_name = "src" + std::to_string(permutation[i].second);
+    const std::string w_name = std::to_string(i * 4);
+    c += "    CONV(" + r_name + ", " + s_name + ", " + w_name + ");\n";
+  }
+  c += "  }\n";
+  if (need_local_mem) {
+    c += "  if (DST_X >= args.dst_tensor.Width() || DST_Y >= "
+         "args.dst_tensor.Height() || Z >= args.dst_tensor.Slices()) return;\n";
+  }
+  c += "  FLT4 bias_val = args.biases.Read(Z);\n";
+  for (int y = 0; y < 2; ++y) {
+    for (int x = 0; x < 2; ++x) {
+      const std::string s_x = std::to_string(x);
+      const std::string s_y = std::to_string(y);
+      const std::string id = std::to_string(y * 2 + x);
+      const std::string x_c = "DST_X + " + s_x;
+      const std::string y_c = "DST_Y + " + s_y;
+      c += "  if (" + x_c + " < args.dst_tensor.Width() && " + y_c +
+           " < args.dst_tensor.Height()) {\n";
+      c += "    FLT4 res0 = TO_FLT4(r" + id + ") + bias_val;\n";
+      c += "    args.dst_tensor.Write(res0, " + x_c + ", " + y_c + ", Z);\n";
+      c += "  }\n";
+    }
+  }
+  c += "}\n";
+  return c;
+}
+
+absl::Status ConvolutionTransposed3x3::BindArguments(ArgumentsBinder* args) {
+  RETURN_IF_ERROR(args->SetInt("filter_offset", 4 * 9 * src_[0]->Slices()));
+  const int padding_x =
+      padding_.x >= 1 ? (padding_.x - 1) / 2 : (padding_.x - 2) / 2;
+  const int padding_y =
+      padding_.y >= 1 ? (padding_.y - 1) / 2 : (padding_.y - 2) / 2;
+  RETURN_IF_ERROR(args->SetInt("padding_x", padding_x));
+  return args->SetInt("padding_y", padding_y);
+}
+
+void ConvolutionTransposed3x3::GetPossibleKernelWorkGroups(
+    TuningType tuning_type, const GpuInfo& gpu_info,
+    const KernelInfo& kernel_info, std::vector<int3>* work_groups) const {
+  if (weights_upload_type_ == WeightsUploadType::LOCAL_MEM_ASYNC ||
+      weights_upload_type_ == WeightsUploadType::LOCAL_MEM_BY_THREADS) {
+    work_groups->push_back(work_group_size_);
+    return;
+  }
+  GetPossibleWorkGroupsConv(tuning_type, gpu_info, kernel_info, grid_size_,
+                            work_groups);
+}
+
+int3 ConvolutionTransposed3x3::GetGridSize() const {
+  const int grid_x = DivideRoundUp(dst_[0]->Width(), 2) * dst_[0]->Batch();
+  const int grid_y = DivideRoundUp(dst_[0]->Height(), 2);
+  const int grid_z = dst_[0]->Slices();
+  return int3(grid_x, grid_y, grid_z);
+}
+
+std::vector<int> ConvolutionTransposed3x3::GetSpatialWeightsRemap() const {
+  const int padding_x_rem = abs(padding_.x) % 2;
+  const int padding_y_rem = abs(padding_.y) % 2;
+
+  std::vector<int> remap;
+  if (padding_x_rem == 1 && padding_y_rem == 1) {
+    return std::vector<int>{4, 5, 3, 7, 1, 8, 6, 2, 0};
+  } else if (padding_x_rem == 0 && padding_y_rem == 1) {
+    return std::vector<int>{5, 3, 4, 8, 6, 2, 0, 7, 1};
+  } else if (padding_x_rem == 1 && padding_y_rem == 0) {
+    return std::vector<int>{7, 1, 8, 6, 2, 0, 4, 5, 3};
+  } else {  // padding_x_rem == 0 && padding_y_rem == 0
+    return std::vector<int>{8, 6, 2, 0, 7, 1, 5, 3, 4};
+  }
+}
+
+void ConvolutionTransposed3x3::UploadWeights(
+    const tflite::gpu::Tensor<OHWI, DataType::FLOAT32>& weights) {
+  const auto weights_desc = GetWeightsDescription();
+  const int flt_count =
+      GetTotalElementsCountForLayout(weights_desc, weights.shape);
+
+  BufferDescriptor desc;
+  desc.element_type = weights_desc.type;
+  desc.element_size = 4;
+  desc.memory_type =
+      weights_upload_type_ ==
+              ConvolutionTransposed3x3::WeightsUploadType::CONSTANT_MEM
+          ? MemoryType::CONSTANT
+          : MemoryType::GLOBAL;
+  desc.size = flt_count * SizeOf(desc.element_type);
+  desc.data.resize(desc.size);
+
+  RearrangeWeights(weights, weights_desc, absl::MakeSpan(desc.data));
+
+  args_.AddObject("weights",
+                  std::make_unique<BufferDescriptor>(std::move(desc)));
+}
+
+bool IsConvolutionTransposed3x3Supported(
+    const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr) {
+  return attr.weights.shape.w == 3 && attr.weights.shape.h == 3 &&
+         attr.stride.w == 2 && attr.stride.h == 2;
+}
+
+ConvolutionTransposed3x3 CreateConvolutionTransposed3x3(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr) {
+  const int2 padding = int2(attr.padding.prepended.w, attr.padding.prepended.h);
+  ConvolutionTransposed3x3 result(definition, gpu_info, padding);
+  result.UploadWeights(attr.weights);
+
+  TensorDescriptor bias_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), attr.bias);
+  result.args_.AddObject("biases", std::make_unique<TensorDescriptor>(
+                                       std::move(bias_tensor_desc)));
+  return result;
+}
+
+ConvolutionTransposed3x3 CreateConvolutionTransposed3x3DynamicWeights(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr) {
+  OperationDef new_def = definition;
+  new_def.src_tensors = {
+      definition.src_tensors[0]};  // leaving only src_tensor def, weights defs
+                                   // will be added later
+  const DataType weights_type = definition.GetDataType();
+  // add 1 src_tensor(buffer) for weights
+  new_def.src_tensors.push_back(
+      {weights_type, TensorStorageType::BUFFER, Layout::HWC});
+
+  const int2 padding = int2(attr.padding.prepended.w, attr.padding.prepended.h);
+  ConvolutionTransposed3x3 result(new_def, gpu_info, padding);
+
+  TensorDescriptor bias_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), attr.bias);
+  result.args_.AddObject("biases", std::make_unique<TensorDescriptor>(
+                                       std::move(bias_tensor_desc)));
+  return result;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_3x3.h

@@ -0,0 +1,109 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_3X3_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_3X3_H_
+
+#include <string>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/buffer_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tensor_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/weights_conversion.h"
+#include "tensorflow/lite/delegates/gpu/common/task/weights_layout.h"
+#include "tensorflow/lite/delegates/gpu/common/tensor.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+class ConvolutionTransposed3x3 : public GPUOperation {
+ public:
+  ConvolutionTransposed3x3() = default;
+  void GetPossibleKernelWorkGroups(
+      TuningType tuning_type, const GpuInfo& gpu_info,
+      const KernelInfo& kernel_info,
+      std::vector<int3>* work_groups) const override;
+  absl::Status BindArguments(ArgumentsBinder* args) override;
+  int3 GetGridSize() const override;
+
+  // Move only
+  ConvolutionTransposed3x3(ConvolutionTransposed3x3&& operation) = default;
+  ConvolutionTransposed3x3& operator=(ConvolutionTransposed3x3&& operation) =
+      default;
+  ConvolutionTransposed3x3(const ConvolutionTransposed3x3&) = delete;
+  ConvolutionTransposed3x3& operator=(const ConvolutionTransposed3x3&) = delete;
+
+  WeightsDescription GetWeightsDescription() const {
+    WeightsDescription desc;
+    desc.type = DeduceDataTypeFromPrecision(definition_.precision);
+    desc.layout = weights_layout_;
+    desc.spatial_remap = GetSpatialWeightsRemap();
+    return desc;
+  }
+
+  enum class WeightsUploadType {
+    LOCAL_MEM_ASYNC,
+    LOCAL_MEM_BY_THREADS,
+    GLOBAL_MEM,
+    CONSTANT_MEM,
+  };
+
+ private:
+  ConvolutionTransposed3x3(const OperationDef& definition,
+                           const GpuInfo& gpu_info, int2 padding);
+  friend ConvolutionTransposed3x3 CreateConvolutionTransposed3x3(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const ConvolutionTransposedAttributes& attr);
+  friend ConvolutionTransposed3x3 CreateConvolutionTransposed3x3DynamicWeights(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const ConvolutionTransposedAttributes& attr);
+
+  void UploadWeights(
+      const tflite::gpu::Tensor<OHWI, DataType::FLOAT32>& weights);
+
+  std::vector<int> GetSpatialWeightsRemap() const;
+
+  std::string GenerateConvolutionTransposedCode(
+      const GpuInfo& gpu_info, const OperationDef& op_def,
+      ConvolutionTransposed3x3::WeightsUploadType weights_upload_type,
+      int2 padding, int3 work_group_launch_order);
+
+  int2 padding_;
+  WeightsUploadType weights_upload_type_;
+  WeightsLayout weights_layout_;
+};
+
+bool IsConvolutionTransposed3x3Supported(
+    const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr);
+
+ConvolutionTransposed3x3 CreateConvolutionTransposed3x3(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr);
+
+ConvolutionTransposed3x3 CreateConvolutionTransposed3x3DynamicWeights(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_3X3_H_

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_3x3_test_util.cc

@@ -0,0 +1,68 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_3x3_test_util.h"
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_3x3.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConvolutionTransposed3x3Test(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 1);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f};
+
+  ConvolutionTransposedAttributes attr;
+  attr.padding.prepended = HW(1, 1);
+  attr.padding.appended = HW(0, 0);
+  attr.stride = HW(2, 2);
+  attr.weights.shape = OHWI(1, 3, 3, 1);
+  attr.weights.data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f};
+  attr.bias.shape = Linear(1);
+  attr.bias.data = {0.0f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      ConvolutionTransposed3x3 operation =
+          CreateConvolutionTransposed3x3(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor,
+          std::make_unique<ConvolutionTransposed3x3>(std::move(operation)),
+          BHWC(1, 4, 4, 1), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({0.0f, 1.0f, 1.0f, 1.0f, 2.0f, 6.0f, 4.0f, 4.0f, 2.0f,
+                         5.0f, 3.0f, 3.0f, 2.0f, 5.0f, 3.0f, 3.0f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_3x3_test_util.h

@@ -0,0 +1,30 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_3X3_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_3X3_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConvolutionTransposed3x3Test(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_3X3_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_3x3_thin.cc

@@ -0,0 +1,280 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_3x3_thin.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/strings/substitute.h"
+#include "tensorflow/lite/delegates/gpu/common/precision.h"
+#include "tensorflow/lite/delegates/gpu/common/task/buffer_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/weights_conversion.h"
+
+namespace tflite {
+namespace gpu {
+namespace {
+std::string ConvInstr(CalculationsPrecision precision, bool is_i4_o4,
+                      const std::string& dst_name, const std::string& src_name,
+                      int weights_offset) {
+  std::string c;
+  if (is_i4_o4) {
+    switch (precision) {
+      case CalculationsPrecision::F32:
+      case CalculationsPrecision::F16:
+        c += "  $0 += $1.x * args.weights.Read($2); \n";
+        c += "  $0 += $1.y * args.weights.Read($3); \n";
+        c += "  $0 += $1.z * args.weights.Read($4); \n";
+        c += "  $0 += $1.w * args.weights.Read($5); \n";
+        break;
+      case CalculationsPrecision::F32_F16:
+        c += "  $0 += TO_ACCUM_TYPE($1.x * args.weights.Read($2) + $1.y * "
+             "args.weights.Read($3) + $1.z * args.weights.Read($4) + $1.w * "
+             "args.weights.Read($5)); \n";
+        break;
+    }
+  } else {
+    // O4I4
+    c += "  $0.x += dot($1, args.weights.Read($2)); \n";
+    c += "  $0.y += dot($1, args.weights.Read($3)); \n";
+    c += "  $0.z += dot($1, args.weights.Read($4)); \n";
+    c += "  $0.w += dot($1, args.weights.Read($5)); \n";
+  }
+  return absl::Substitute(c, dst_name, src_name, weights_offset,
+                          weights_offset + 1, weights_offset + 2,
+                          weights_offset + 3);
+}
+}  // namespace
+
+ConvolutionTransposed3x3Thin::ConvolutionTransposed3x3Thin(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr)
+    : GPUOperation(definition) {
+  if (gpu_info.IsApple()) {
+    weights_layout_ = WeightsLayout::kOICustomSpatialO4I4;
+  } else {
+    weights_layout_ = WeightsLayout::kOICustomSpatialI4O4;
+  }
+  code_ = GenerateConvolutionTransposedCode(
+      definition_, gpu_info, DivideRoundUp(attr.weights.shape.i, 4),
+      DivideRoundUp(attr.weights.shape.o, 4));
+}
+
+std::string ConvolutionTransposed3x3Thin::GenerateConvolutionTransposedCode(
+    const OperationDef& op_def, const GpuInfo& gpu_info, int src_depth,
+    int dst_depth) {
+  AddSrcTensor("src_tensor", op_def.src_tensors[0]);
+  AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
+
+  if (op_def.src_tensors.size() == 2) {
+    // dynamic weights
+    BufferDescriptor desc;
+    desc.element_type = op_def.src_tensors[1].GetDataType();
+    desc.element_size = 4;
+    desc.memory_type = MemoryType::CONSTANT;
+    AddSrcBuffer("weights", desc);
+  }
+
+  std::string c;
+
+  c += "MAIN_FUNCTION($0) {\n";
+  if (op_def.IsBatchSupported()) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  c += "  int Y = GLOBAL_ID_1;\n";
+  c += "  if (X >= args.src_tensor.Width() || Y >= args.src_tensor.Height()) "
+       "return;\n";
+  for (int d = 0; d < dst_depth; ++d) {
+    const std::string layer = std::to_string(d);
+    c += "  ACCUM_FLT4 r" + layer + "[2][2];\n";
+    c += "  r" + layer + "[0][0] = INIT_ACCUM_FLT4(0.0f);\n";
+    c += "  r" + layer + "[0][1] = INIT_ACCUM_FLT4(0.0f);\n";
+    c += "  r" + layer + "[1][0] = INIT_ACCUM_FLT4(0.0f);\n";
+    c += "  r" + layer + "[1][1] = INIT_ACCUM_FLT4(0.0f);\n";
+  }
+  for (int s = 0; s < src_depth; ++s) {
+    const std::string z = std::to_string(s);
+    c += "  {\n";
+    if (op_def.src_tensors[0].SupportsZeroClamp(Axis::WIDTH, gpu_info) &&
+        op_def.src_tensors[0].SupportsZeroClamp(Axis::HEIGHT, gpu_info)) {
+      c += "  FLT4 src0 = args.src_tensor.Read(X, Y, " + z + ");\n";
+      c += "  FLT4 src1 = args.src_tensor.Read(X + 1, Y, " + z + ");\n";
+      c += "  FLT4 src2 = args.src_tensor.Read(X, Y + 1, " + z + ");\n";
+      c += "  FLT4 src3 = args.src_tensor.Read(X + 1, Y + 1, " + z + ");\n";
+    } else if (op_def.src_tensors[0].IsLinear() &&
+               op_def.src_tensors[0].ReturnsZeroForNegOneRead(gpu_info)) {
+      c += "  int c0 = args.src_tensor.GetAddress(X, Y, " + z + ");\n";
+      c += "  int c1 = args.src_tensor.GetAddress(X + 1, Y, " + z + ");\n";
+      c += "  int c2 = args.src_tensor.GetAddress(X, Y + 1, " + z + ");\n";
+      c += "  int c3 = args.src_tensor.GetAddress(X + 1, Y + 1, " + z + ");\n";
+      c += "  bool x_in = X + 1 < args.src_tensor.Width();\n";
+      c += "  bool y_in = Y + 1 < args.src_tensor.Height();\n";
+      c += "  c1 = select(-1, c1, x_in);\n";
+      c += "  c2 = select(-1, c2, y_in);\n";
+      c += "  c3 = select(-1, c3, x_in && y_in);\n";
+      c += "  FLT4 src0 = args.src_tensor.Read(c0);\n";
+      c += "  FLT4 src1 = args.src_tensor.Read(c1);\n";
+      c += "  FLT4 src2 = args.src_tensor.Read(c2);\n";
+      c += "  FLT4 src3 = args.src_tensor.Read(c3);\n";
+    } else {
+      // Manual zero clamp
+      c += "  bool x_in = X + 1 < args.src_tensor.Width();\n";
+      c += "  bool y_in = Y + 1 < args.src_tensor.Height();\n";
+      c += "  FLT4 src0 = args.src_tensor.Read(X, Y, " + z + ");\n";
+      c += "  FLT4 src1 = INIT_FLT4(0.0);\n";
+      c += "  FLT4 src2 = INIT_FLT4(0.0);\n";
+      c += "  FLT4 src3 = INIT_FLT4(0.0);\n";
+      c += "  if (x_in) {\n";
+      c += "    src1 = args.src_tensor.Read(X + 1, Y, " + z + ");\n";
+      c += "  }\n";
+      c += "  if (y_in) {\n";
+      c += "    src2 = args.src_tensor.Read(X, Y + 1, " + z + ");\n";
+      c += "  }\n";
+      c += "  if (x_in && y_in) {\n";
+      c += "    src3 = args.src_tensor.Read(X + 1, Y + 1, " + z + ");\n";
+      c += "  }\n";
+    }
+    for (int d = 0; d < dst_depth; ++d) {
+      const std::string layer = std::to_string(d);
+      const int filters_index = (s * dst_depth + d) * 36;
+      const bool is_i4_o4 = GetWeightsDescription().IsI4O4();
+      c += ConvInstr(op_def.precision, is_i4_o4, "r" + layer + "[0][0]", "src0",
+                     filters_index);
+      c += ConvInstr(op_def.precision, is_i4_o4, "r" + layer + "[0][1]", "src0",
+                     filters_index + 4);
+      c += ConvInstr(op_def.precision, is_i4_o4, "r" + layer + "[0][1]", "src1",
+                     filters_index + 8);
+      c += ConvInstr(op_def.precision, is_i4_o4, "r" + layer + "[1][0]", "src0",
+                     filters_index + 12);
+      c += ConvInstr(op_def.precision, is_i4_o4, "r" + layer + "[1][0]", "src2",
+                     filters_index + 16);
+      c += ConvInstr(op_def.precision, is_i4_o4, "r" + layer + "[1][1]", "src0",
+                     filters_index + 20);
+      c += ConvInstr(op_def.precision, is_i4_o4, "r" + layer + "[1][1]", "src1",
+                     filters_index + 24);
+      c += ConvInstr(op_def.precision, is_i4_o4, "r" + layer + "[1][1]", "src2",
+                     filters_index + 28);
+      c += ConvInstr(op_def.precision, is_i4_o4, "r" + layer + "[1][1]", "src3",
+                     filters_index + 32);
+    }
+    c += "  }\n";
+  }
+  c += "  X *= 2;\n";
+  c += "  Y *= 2;\n";
+  for (int d = 0; d < dst_depth; ++d) {
+    const std::string layer = std::to_string(d);
+    c += "  {\n";
+    c += "  FLT4 bias_val = args.biases.Read(" + layer + ");\n";
+    for (int y = 0; y < 2; ++y) {
+      for (int x = 0; x < 2; ++x) {
+        const std::string x_coord = "X + " + std::to_string(x);
+        const std::string y_coord = "Y + " + std::to_string(y);
+        c += "  {\n";
+        c += "    FLT4 result = TO_FLT4(r" + layer + "[" + std::to_string(y) +
+             "][" + std::to_string(x) + "]) + bias_val;\n";
+        c += "    args.dst_tensor.Write(result, " + x_coord + ", " + y_coord +
+             ", " + layer + ");\n";
+        c += "  }\n";
+      }
+    }
+    c += "  }\n";
+  }
+  c += "}\n";
+
+  return c;
+}
+
+int3 ConvolutionTransposed3x3Thin::GetGridSize() const {
+  const int grid_x = src_[0]->Width() * dst_[0]->Batch();
+  const int grid_y = src_[0]->Height();
+  const int grid_z = 1;
+  return int3(grid_x, grid_y, grid_z);
+}
+
+std::vector<int> ConvolutionTransposed3x3Thin::GetSpatialWeightsRemap() const {
+  return std::vector<int>{4, 5, 3, 7, 1, 8, 6, 2, 0};
+}
+
+void ConvolutionTransposed3x3Thin::UploadWeights(
+    const tflite::gpu::Tensor<OHWI, DataType::FLOAT32>& weights) {
+  const auto weights_desc = GetWeightsDescription();
+  const int flt_count =
+      GetTotalElementsCountForLayout(weights_desc, weights.shape);
+
+  BufferDescriptor desc;
+  desc.element_type = weights_desc.type;
+  desc.element_size = 4;
+  desc.memory_type = MemoryType::CONSTANT;
+  desc.size = flt_count * SizeOf(desc.element_type);
+  desc.data.resize(desc.size);
+
+  RearrangeWeights(weights, weights_desc, absl::MakeSpan(desc.data));
+
+  args_.AddObject("weights",
+                  std::make_unique<BufferDescriptor>(std::move(desc)));
+}
+
+bool IsConvolutionTransposed3x3ThinSupported(
+    const ConvolutionTransposedAttributes& attr) {
+  return attr.weights.shape.o <= 8 && attr.weights.shape.w == 3 &&
+         attr.weights.shape.h == 3 && attr.stride.w == 2 &&
+         attr.stride.h == 2 && attr.padding.prepended.w == 1 &&
+         attr.padding.prepended.h == 1 && attr.padding.appended.w == 1 &&
+         attr.padding.appended.h == 1;
+}
+
+ConvolutionTransposed3x3Thin CreateConvolutionTransposed3x3Thin(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr) {
+  ConvolutionTransposed3x3Thin result(gpu_info, definition, attr);
+  result.UploadWeights(attr.weights);
+
+  TensorDescriptor bias_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), attr.bias);
+  result.args_.AddObject("biases", std::make_unique<TensorDescriptor>(
+                                       std::move(bias_tensor_desc)));
+  return result;
+}
+
+ConvolutionTransposed3x3Thin CreateConvolutionTransposed3x3ThinDynamicWeights(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr) {
+  OperationDef new_def = definition;
+  new_def.src_tensors = {
+      definition.src_tensors[0]};  // leaving only src_tensor def, weights defs
+                                   // will be added later
+  const DataType weights_type = definition.GetDataType();
+  // add 1 src_tensor(buffer) for weights
+  new_def.src_tensors.push_back(
+      {weights_type, TensorStorageType::BUFFER, Layout::HWC});
+  ConvolutionTransposed3x3Thin result(gpu_info, new_def, attr);
+
+  TensorDescriptor bias_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), attr.bias);
+  result.args_.AddObject("biases", std::make_unique<TensorDescriptor>(
+                                       std::move(bias_tensor_desc)));
+  return result;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_3x3_thin.h

@@ -0,0 +1,95 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_3X3_THIN_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_3X3_THIN_H_
+
+#include <string>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/weights_layout.h"
+#include "tensorflow/lite/delegates/gpu/common/tensor.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+class ConvolutionTransposed3x3Thin : public GPUOperation {
+ public:
+  ConvolutionTransposed3x3Thin() = default;
+  int3 GetGridSize() const override;
+
+  // Move only
+  ConvolutionTransposed3x3Thin(ConvolutionTransposed3x3Thin&& operation) =
+      default;
+  ConvolutionTransposed3x3Thin& operator=(
+      ConvolutionTransposed3x3Thin&& operation) = default;
+  ConvolutionTransposed3x3Thin(const ConvolutionTransposed3x3Thin&) = delete;
+  ConvolutionTransposed3x3Thin& operator=(const ConvolutionTransposed3x3Thin&) =
+      delete;
+
+  WeightsDescription GetWeightsDescription() const {
+    WeightsDescription desc;
+    desc.type = DeduceDataTypeFromPrecision(definition_.precision);
+    desc.layout = weights_layout_;
+    desc.spatial_remap = GetSpatialWeightsRemap();
+    return desc;
+  }
+
+ private:
+  ConvolutionTransposed3x3Thin(const GpuInfo& gpu_info,
+                               const OperationDef& definition,
+                               const ConvolutionTransposedAttributes& attr);
+
+  friend ConvolutionTransposed3x3Thin CreateConvolutionTransposed3x3Thin(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const ConvolutionTransposedAttributes& attr);
+  friend ConvolutionTransposed3x3Thin
+  CreateConvolutionTransposed3x3ThinDynamicWeights(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const ConvolutionTransposedAttributes& attr);
+
+  void UploadWeights(
+      const tflite::gpu::Tensor<OHWI, DataType::FLOAT32>& weights);
+
+  std::vector<int> GetSpatialWeightsRemap() const;
+
+  std::string GenerateConvolutionTransposedCode(const OperationDef& op_def,
+                                                const GpuInfo& gpu_info,
+                                                int src_depth, int dst_depth);
+
+  WeightsLayout weights_layout_;
+};
+
+bool IsConvolutionTransposed3x3ThinSupported(
+    const ConvolutionTransposedAttributes& attr);
+
+ConvolutionTransposed3x3Thin CreateConvolutionTransposed3x3Thin(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr);
+
+ConvolutionTransposed3x3Thin CreateConvolutionTransposed3x3ThinDynamicWeights(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_3X3_THIN_H_

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_3x3_thin_test_util.cc

@@ -0,0 +1,107 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_3x3_thin_test_util.h"
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_3x3_thin.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConvolutionTransposed3x3ThinSimpleWeightsTest(
+    TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 1);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f};
+
+  ConvolutionTransposedAttributes attr;
+  attr.padding.prepended = HW(1, 1);
+  attr.padding.appended = HW(1, 1);
+  attr.stride = HW(2, 2);
+  attr.weights.shape = OHWI(1, 3, 3, 1);
+  attr.weights.data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f};
+  attr.bias.shape = Linear(2);
+  attr.bias.data = {0.0f, 0.0f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      ConvolutionTransposed3x3Thin operation =
+          CreateConvolutionTransposed3x3Thin(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor,
+          std::make_unique<ConvolutionTransposed3x3Thin>(std::move(operation)),
+          BHWC(1, 4, 4, 1), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({0.0f, 1.0f, 1.0f, 1.0f, 2.0f, 6.0f, 4.0f, 4.0f, 2.0f,
+                         5.0f, 3.0f, 3.0f, 2.0f, 5.0f, 3.0f, 3.0f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConvolutionTransposed3x3ThinTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 1);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f};
+
+  ConvolutionTransposedAttributes attr;
+  attr.padding.prepended = HW(1, 1);
+  attr.padding.appended = HW(1, 1);
+  attr.stride = HW(2, 2);
+  attr.weights.shape = OHWI(1, 3, 3, 1);
+  attr.weights.data = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f};
+  attr.bias.shape = Linear(1);
+  attr.bias.data = {0.5f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      ConvolutionTransposed3x3Thin operation =
+          CreateConvolutionTransposed3x3Thin(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor,
+          std::make_unique<ConvolutionTransposed3x3Thin>(std::move(operation)),
+          BHWC(1, 4, 4, 1), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear(
+          {0.5f, 4.5f, 5.5f, 6.5f, 4.5f, 16.5f, 14.5f, 18.5f, 10.5f, 24.5f,
+           15.5f, 18.5f, 16.5f, 39.5f, 24.5f, 27.5f},
+          dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_3x3_thin_test_util.h

@@ -0,0 +1,33 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_3X3_THIN_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_3X3_THIN_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConvolutionTransposed3x3ThinSimpleWeightsTest(
+    TestExecutionEnvironment* env);
+
+absl::Status ConvolutionTransposed3x3ThinTest(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_3X3_THIN_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_4x4.cc

@@ -0,0 +1,437 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_4x4.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
+
+namespace tflite {
+namespace gpu {
+
+namespace {
+ConvolutionTransposed4x4::WeightsUploadType GetBestWeightsUploadType(
+    const GpuInfo& gpu_info) {
+  ConvolutionTransposed4x4::WeightsUploadType weights_upload_type =
+      ConvolutionTransposed4x4::WeightsUploadType::GLOBAL_MEM;
+  if (gpu_info.IsApple()) {
+    if (gpu_info.apple_info.IsBionic()) {
+      weights_upload_type =
+          ConvolutionTransposed4x4::WeightsUploadType::GLOBAL_MEM;
+    } else {
+      weights_upload_type =
+          ConvolutionTransposed4x4::WeightsUploadType::LOCAL_MEM_BY_THREADS;
+    }
+  } else if (gpu_info.IsPowerVR()) {
+    weights_upload_type =
+        ConvolutionTransposed4x4::WeightsUploadType::LOCAL_MEM_ASYNC;
+  } else if (gpu_info.IsNvidia() || gpu_info.IsIntel()) {
+    weights_upload_type =
+        ConvolutionTransposed4x4::WeightsUploadType::LOCAL_MEM_BY_THREADS;
+  } else if (gpu_info.IsAMD()) {
+    weights_upload_type =
+        ConvolutionTransposed4x4::WeightsUploadType::CONSTANT_MEM;
+  } else {
+    weights_upload_type =
+        ConvolutionTransposed4x4::WeightsUploadType::GLOBAL_MEM;
+  }
+  return weights_upload_type;
+}
+}  // namespace
+
+ConvolutionTransposed4x4::ConvolutionTransposed4x4(
+    const OperationDef& definition, const GpuInfo& gpu_info)
+    : GPUOperation(definition) {
+  work_group_size_ = int3(8, 4, 1);
+  if (gpu_info.IsApple()) {
+    work_group_launch_order_ = int3(2, 0, 1);
+  }
+
+  if (gpu_info.IsApple()) {
+    weights_layout_ = WeightsLayout::kOICustomSpatialO4I4;
+  } else {
+    weights_layout_ = WeightsLayout::kOICustomSpatialI4O4;
+  }
+
+  code_ = GenerateConvolutionTransposedCode(gpu_info, definition_,
+                                            GetBestWeightsUploadType(gpu_info));
+  if (definition_.precision == CalculationsPrecision::F16 &&
+      gpu_info.IsPowerVR()) {
+    compiler_options_.push_back(CompilerOptions::kClFastRelaxedMath);
+  }
+}
+
+std::string ConvolutionTransposed4x4::GenerateConvolutionTransposedCode(
+    const GpuInfo& gpu_info, const OperationDef& op_def,
+    WeightsUploadType weights_upload_type) {
+  auto src_desc = op_def.src_tensors[0];
+  AddSrcTensor("src_tensor", src_desc);
+  AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
+
+  if (op_def.src_tensors.size() == 2) {
+    // dynamic weights
+    BufferDescriptor desc;
+    desc.element_type = op_def.src_tensors[1].GetDataType();
+    desc.element_size = 4;
+    desc.memory_type =
+        weights_upload_type ==
+                ConvolutionTransposed4x4::WeightsUploadType::CONSTANT_MEM
+            ? MemoryType::CONSTANT
+            : MemoryType::GLOBAL;
+    AddSrcBuffer("weights", desc);
+  }
+
+  args_.AddInt("filter_offset");
+
+  const bool need_local_mem =
+      weights_upload_type ==
+          ConvolutionTransposed4x4::WeightsUploadType::LOCAL_MEM_BY_THREADS ||
+      weights_upload_type ==
+          ConvolutionTransposed4x4::WeightsUploadType::LOCAL_MEM_ASYNC;
+
+  const int wg_total_size =
+      work_group_size_.x * work_group_size_.y * work_group_size_.z;
+  const std::string barrier =
+      wg_total_size == 32 && gpu_info.IsWaveSizeEqualTo32()
+          ? "SIMD_LOCAL_MEM_BARRIER"
+          : "LOCAL_MEM_BARRIER";
+
+  std::string c;
+  if (GetWeightsDescription().IsI4O4()) {
+    switch (op_def.precision) {
+      case CalculationsPrecision::F32:
+      case CalculationsPrecision::F16:
+        c += "#define CONV(R, SRC, F) \\\n";
+        c += "  R += SRC.x * weights_cache[F]; \\\n";
+        c += "  R += SRC.y * weights_cache[F + 1]; \\\n";
+        c += "  R += SRC.z * weights_cache[F + 2]; \\\n";
+        c += "  R += SRC.w * weights_cache[F + 3];   \n";
+        break;
+      case CalculationsPrecision::F32_F16:
+        c += "#define CONV(R, SRC, F) \\\n";
+        c += "  R += TO_ACCUM_TYPE(SRC.x * weights_cache[F] + SRC.y * "
+             "weights_cache[F + 1] + SRC.z * weights_cache[F + 2] + SRC.w * "
+             "weights_cache[F + 3]);\n";
+        break;
+    }
+  } else {
+    // O4I4
+    c += "#define CONV(R, SRC, F) \\\n";
+    c += "  R.x += dot(SRC, weights_cache[F]); \\\n";
+    c += "  R.y += dot(SRC, weights_cache[F + 1]); \\\n";
+    c += "  R.z += dot(SRC, weights_cache[F + 2]); \\\n";
+    c += "  R.w += dot(SRC, weights_cache[F + 3]);   \n";
+  }
+
+  const std::string weights_space =
+      weights_upload_type ==
+              ConvolutionTransposed4x4::WeightsUploadType::CONSTANT_MEM
+          ? "__constant"
+          : "__global";
+
+  if (gpu_info.IsApiOpenCl()) {
+    c += "__attribute__((reqd_work_group_size(8, 4, 1)))\n";
+  }
+  c += "MAIN_FUNCTION($0) {\n";
+  std::string grid_coords[3];
+  int3 launch_remap;
+  launch_remap[work_group_launch_order_.x] = 0;
+  launch_remap[work_group_launch_order_.y] = 1;
+  launch_remap[work_group_launch_order_.z] = 2;
+  if (work_group_launch_order_[0] == 0) {
+    grid_coords[0] = "GLOBAL_ID_0";
+  } else {
+    grid_coords[0] = "(GROUP_ID_" + std::to_string(launch_remap[0]) +
+                     " * GROUP_SIZE_0 + LOCAL_ID_0);\n";
+  }
+  if (work_group_launch_order_[1] == 1) {
+    grid_coords[1] = "GLOBAL_ID_1";
+  } else {
+    grid_coords[1] = "(GROUP_ID_" + std::to_string(launch_remap[1]) +
+                     " * GROUP_SIZE_1 + LOCAL_ID_1);\n";
+  }
+  if (work_group_launch_order_[2] == 2) {
+    grid_coords[2] = "GLOBAL_ID_2";
+  } else {
+    grid_coords[2] = "(GROUP_ID_" + std::to_string(launch_remap[2]) +
+                     " * GROUP_SIZE_2 + LOCAL_ID_2);\n";
+  }
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = " + grid_coords[0] + ";\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = " + grid_coords[0] + ";\n";
+  }
+  c += "  int Y = " + grid_coords[1] + ";\n";
+  c += "  int Z = " + grid_coords[2] + ";\n";
+  if (!need_local_mem) {
+    c += "  if (X * 2 > args.dst_tensor.Width() || Y * 2 > "
+         "args.dst_tensor.Height() || Z >= args.dst_tensor.Slices()) "
+         "return;\n";
+  }
+  c += "  ACCUM_FLT4 r0 = INIT_ACCUM_FLT4(0.0f);\n";
+  c += "  ACCUM_FLT4 r1 = INIT_ACCUM_FLT4(0.0f);\n";
+  c += "  ACCUM_FLT4 r2 = INIT_ACCUM_FLT4(0.0f);\n";
+  c += "  ACCUM_FLT4 r3 = INIT_ACCUM_FLT4(0.0f);\n";
+  c += "  int f_offset = Z * args.filter_offset;\n";
+  if (need_local_mem) {
+    c += "  __local FLT4 weights_cache[64];\n";
+  }
+  if (weights_upload_type ==
+      ConvolutionTransposed4x4::WeightsUploadType::LOCAL_MEM_BY_THREADS) {
+    c += "  int local_id = LOCAL_ID_1 * 8 + LOCAL_ID_0;\n";
+  }
+  if (!src_desc.SupportsZeroClamp(Axis::WIDTH, gpu_info)) {
+    c += "  bool in_x0 = X - 1 >= 0 && X - 1 < args.src_tensor.Width();\n";
+    c += "  bool in_x1 = X >= 0 && X < args.src_tensor.Width();\n";
+  }
+  if (!src_desc.SupportsZeroClamp(Axis::HEIGHT, gpu_info)) {
+    c += "  bool in_y0 = Y - 1 >= 0 && Y - 1 < args.src_tensor.Height();\n";
+    c += "  bool in_y1 = Y >= 0 && Y < args.src_tensor.Height();\n";
+  }
+  auto generate_check = [&](int x, int y) {
+    std::string check;
+    const std::vector<Axis> axes{Axis::WIDTH, Axis::HEIGHT};
+    const std::vector<std::string> names{"in_x" + std::to_string(x),
+                                         "in_y" + std::to_string(y)};
+    for (int i = 0; i < axes.size(); ++i) {
+      const auto& axis = axes[i];
+      if (src_desc.HasAxis(axis) &&
+          !src_desc.SupportsZeroClamp(axis, gpu_info)) {
+        if (!check.empty()) {
+          check += " && ";
+        }
+        check += names[i];
+      }
+    }
+    return check;
+  };
+  if (src_desc.IsLinear()) {
+    if (src_desc.ReturnsZeroForNegOneRead(gpu_info)) {
+      c += "  int addr_0 = args.src_tensor.GetAddress(X - 1, Y - 1, 0);\n";
+      c += "  int addr_1 = args.src_tensor.GetAddress(X, Y - 1, 0);\n";
+      c += "  int addr_2 = args.src_tensor.GetAddress(X - 1, Y, 0);\n";
+      c += "  int addr_3 = args.src_tensor.GetAddress(X, Y, 0);\n";
+      c += "  addr_0 = select(-1, addr_0, (in_x0 && in_y0));\n";
+      c += "  addr_1 = select(-1, addr_1, (in_x1 && in_y0));\n";
+      c += "  addr_2 = select(-1, addr_2, (in_x0 && in_y1));\n";
+      c += "  addr_3 = select(-1, addr_3, (in_x1 && in_y1));\n";
+      c += "  int dz_0 = select(0, args.src_tensor.SliceStride(), (in_x0 && "
+           "in_y0));\n";
+      c += "  int dz_1 = select(0, args.src_tensor.SliceStride(), (in_x1 && "
+           "in_y0));\n";
+      c += "  int dz_2 = select(0, args.src_tensor.SliceStride(), (in_x0 && "
+           "in_y1));\n";
+      c += "  int dz_3 = select(0, args.src_tensor.SliceStride(), (in_x1 && "
+           "in_y1));\n";
+    } else {
+      c += "  int xc0 = clamp(X - 1, 0, args.src_tensor.Width() - 1);\n";
+      c += "  int xc1 = clamp(X, 0, args.src_tensor.Width() - 1);\n";
+      c += "  int yc0 = clamp(Y - 1, 0, args.src_tensor.Height() - 1);\n";
+      c += "  int yc1 = clamp(Y, 0, args.src_tensor.Height() - 1);\n";
+      c += "  int addr_0 = args.src_tensor.GetAddress(xc0, yc0, 0);\n";
+      c += "  int addr_1 = args.src_tensor.GetAddress(xc1, yc0, 0);\n";
+      c += "  int addr_2 = args.src_tensor.GetAddress(xc0, yc1, 0);\n";
+      c += "  int addr_3 = args.src_tensor.GetAddress(xc1, yc1, 0);\n";
+      c += "  int dz = args.src_tensor.SliceStride();\n";
+    }
+  }
+  auto read_src = [&](int x, int y) {
+    if (src_desc.IsLinear()) {
+      const std::string id = std::to_string(y * 2 + x);
+      const std::string addr = "addr_" + std::to_string(y * 2 + x);
+      if (src_desc.ReturnsZeroForNegOneRead(gpu_info)) {
+        return "args.src_tensor.Read(" + addr + "); " + addr + " += dz_" + id +
+               ";";
+      } else {
+        return "args.src_tensor.Read(" + addr + ") * INIT_FLT(in_x" +
+               std::to_string(x) + " && in_y" + std::to_string(y) + "); " +
+               addr + " += dz;";
+      }
+    } else {
+      std::string check = generate_check(x, y);
+      if (!check.empty()) {
+        check = " * INIT_FLT(" + check + ")";
+      }
+      return "args.src_tensor.Read(X + " + std::to_string(x - 1) + ", Y + " +
+             std::to_string(y - 1) + ", s)" + check + ";";
+    }
+  };
+  c += "  for (int s = 0; s < args.src_tensor.Slices(); ++s) {\n";
+  if (need_local_mem) {
+    c += "    " + barrier + ";\n";
+  }
+  if (weights_upload_type ==
+      ConvolutionTransposed4x4::WeightsUploadType::LOCAL_MEM_ASYNC) {
+    c += "    async_work_group_copy(weights_cache, "
+         "args.weights.GetPtr(f_offset), 64, "
+         "0);\n";
+  } else if (weights_upload_type ==
+             ConvolutionTransposed4x4::WeightsUploadType::
+                 LOCAL_MEM_BY_THREADS) {
+    c += "    weights_cache[local_id] = args.weights.Read(f_offset + "
+         "local_id);\n";
+    c += "    weights_cache[local_id + 32] = args.weights.Read(f_offset + "
+         "local_id + "
+         "32);\n";
+  } else {  // GLOBAL_MEM
+    c += "    " + weights_space +
+         " FLT4* weights_cache = args.weights.GetPtr(f_offset);\n";
+  }
+  c += "    FLT4 src0 = " + read_src(0, 0) + ";\n";
+  c += "    FLT4 src1 = " + read_src(1, 0) + ";\n";
+  c += "    FLT4 src2 = " + read_src(0, 1) + ";\n";
+  c += "    FLT4 src3 = " + read_src(1, 1) + ";\n";
+  c += "    f_offset += 64;\n";
+  if (need_local_mem) {
+    c += "    " + barrier + ";\n";
+  }
+  c += "    CONV(r0, src0, 0);\n";
+  c += "    CONV(r1, src0, 4);\n";
+  c += "    CONV(r2, src0, 8);\n";
+  c += "    CONV(r3, src0, 12);\n";
+  c += "    CONV(r0, src1, 16);\n";
+  c += "    CONV(r1, src1, 20);\n";
+  c += "    CONV(r2, src1, 24);\n";
+  c += "    CONV(r3, src1, 28);\n";
+  c += "    CONV(r0, src2, 32);\n";
+  c += "    CONV(r1, src2, 36);\n";
+  c += "    CONV(r2, src2, 40);\n";
+  c += "    CONV(r3, src2, 44);\n";
+  c += "    CONV(r0, src3, 48);\n";
+  c += "    CONV(r1, src3, 52);\n";
+  c += "    CONV(r2, src3, 56);\n";
+  c += "    CONV(r3, src3, 60);\n";
+  c += "  }\n";
+  c += "\n";
+  if (need_local_mem) {
+    c += "  if (X * 2 > args.dst_tensor.Width() || Y * 2 > "
+         "args.dst_tensor.Height() || Z >= args.dst_tensor.Slices()) "
+         "return;\n";
+  }
+  c += "  X = X * 2 - 1;\n";
+  c += "  Y = Y * 2 - 1;\n";
+  c += "\n";
+  c += "  FLT4 bias_val = args.biases.Read(Z);\n";
+  c += "  if (X >= 0 && Y >= 0) {\n";
+  c += "    FLT4 result = TO_FLT4(r0) + bias_val;\n";
+  c += "    args.dst_tensor.Write(result, X, Y, Z);\n";
+  c += "  }\n";
+  c += "  if (X + 1 < args.dst_tensor.Width() && Y >= 0) {\n";
+  c += "    FLT4 result = TO_FLT4(r1) + bias_val;\n";
+  c += "    args.dst_tensor.Write(result, X + 1, Y, Z);\n";
+  c += "  }\n";
+  c += "  if (X >= 0 && Y + 1 < args.dst_tensor.Height()) {\n";
+  c += "    FLT4 result = TO_FLT4(r2) + bias_val;\n";
+  c += "    args.dst_tensor.Write(result, X, Y + 1, Z);\n";
+  c += "  }\n";
+  c += "  if (X + 1 < args.dst_tensor.Width() && Y + 1 < "
+       "args.dst_tensor.Height()) {\n";
+  c += "    FLT4 result = TO_FLT4(r3) + bias_val;\n";
+  c += "    args.dst_tensor.Write(result, X + 1, Y + 1, Z);\n";
+  c += "  }\n";
+  c += "}\n";
+  return c;
+}
+
+absl::Status ConvolutionTransposed4x4::BindArguments(ArgumentsBinder* args) {
+  return args->SetInt("filter_offset", 4 * 16 * src_[0]->Slices());
+}
+
+int3 ConvolutionTransposed4x4::GetGridSize() const {
+  const int grid_x = DivideRoundUp(dst_[0]->Width() + 2, 2) * dst_[0]->Batch();
+  const int grid_y = DivideRoundUp(dst_[0]->Height() + 2, 2);
+  const int grid_z = dst_[0]->Slices();
+  return int3(grid_x, grid_y, grid_z);
+}
+
+std::vector<int> ConvolutionTransposed4x4::GetSpatialWeightsRemap() const {
+  return std::vector<int>{10, 11, 14, 15, 8, 9, 12, 13, 2, 3, 6, 7, 0, 1, 4, 5};
+}
+
+void ConvolutionTransposed4x4::UploadWeights(
+    const tflite::gpu::Tensor<OHWI, DataType::FLOAT32>& weights,
+    WeightsUploadType weights_upload_type) {
+  const auto weights_desc = GetWeightsDescription();
+  const int flt_count =
+      GetTotalElementsCountForLayout(weights_desc, weights.shape);
+
+  BufferDescriptor desc;
+  desc.element_type = weights_desc.type;
+  desc.element_size = 4;
+  desc.memory_type =
+      weights_upload_type ==
+              ConvolutionTransposed4x4::WeightsUploadType::CONSTANT_MEM
+          ? MemoryType::CONSTANT
+          : MemoryType::GLOBAL;
+  desc.size = flt_count * SizeOf(desc.element_type);
+  desc.data.resize(desc.size);
+
+  RearrangeWeights(weights, weights_desc, absl::MakeSpan(desc.data));
+  args_.AddObject("weights",
+                  std::make_unique<BufferDescriptor>(std::move(desc)));
+}
+
+bool IsConvolutionTransposed4x4Supported(
+    const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr) {
+  return attr.weights.shape.w == 4 && attr.weights.shape.h == 4 &&
+         attr.stride.w == 2 && attr.stride.h == 2 &&
+         attr.padding.prepended.w == 1 && attr.padding.prepended.h == 1;
+}
+
+ConvolutionTransposed4x4 CreateConvolutionTransposed4x4(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr) {
+  ConvolutionTransposed4x4 result(definition, gpu_info);
+  result.UploadWeights(attr.weights, GetBestWeightsUploadType(gpu_info));
+
+  TensorDescriptor bias_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), attr.bias);
+  result.args_.AddObject("biases", std::make_unique<TensorDescriptor>(
+                                       std::move(bias_tensor_desc)));
+  return result;
+}
+
+ConvolutionTransposed4x4 CreateConvolutionTransposed4x4DynamicWeights(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr) {
+  OperationDef new_def = definition;
+  new_def.src_tensors = {
+      definition.src_tensors[0]};  // leaving only src_tensor def, weights defs
+                                   // will be added later
+  const DataType weights_type = definition.GetDataType();
+  // add 1 src_tensor(buffer) for weights
+  new_def.src_tensors.push_back(
+      {weights_type, TensorStorageType::BUFFER, Layout::HWC});
+
+  ConvolutionTransposed4x4 result(new_def, gpu_info);
+
+  TensorDescriptor bias_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), attr.bias);
+  result.args_.AddObject("biases", std::make_unique<TensorDescriptor>(
+                                       std::move(bias_tensor_desc)));
+  return result;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_4x4.h

@@ -0,0 +1,110 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_4X4_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_4X4_H_
+
+#include <string>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/buffer_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tensor_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/weights_conversion.h"
+#include "tensorflow/lite/delegates/gpu/common/task/weights_layout.h"
+#include "tensorflow/lite/delegates/gpu/common/tensor.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+class ConvolutionTransposed4x4 : public GPUOperation {
+ public:
+  ConvolutionTransposed4x4() = default;
+  void GetPossibleKernelWorkGroups(
+      TuningType tuning_type, const GpuInfo& gpu_info,
+      const KernelInfo& kernel_info,
+      std::vector<int3>* work_groups) const override {
+    work_groups->push_back(work_group_size_);
+  }
+  absl::Status BindArguments(ArgumentsBinder* args) override;
+  int3 GetGridSize() const override;
+
+  // Move only
+  ConvolutionTransposed4x4(ConvolutionTransposed4x4&& operation) = default;
+  ConvolutionTransposed4x4& operator=(ConvolutionTransposed4x4&& operation) =
+      default;
+  ConvolutionTransposed4x4(const ConvolutionTransposed4x4&) = delete;
+  ConvolutionTransposed4x4& operator=(const ConvolutionTransposed4x4&) = delete;
+
+  WeightsDescription GetWeightsDescription() const {
+    WeightsDescription desc;
+    desc.type = DeduceDataTypeFromPrecision(definition_.precision);
+    desc.layout = weights_layout_;
+    desc.spatial_remap = GetSpatialWeightsRemap();
+    return desc;
+  }
+
+  enum class WeightsUploadType {
+    LOCAL_MEM_ASYNC,
+    LOCAL_MEM_BY_THREADS,
+    GLOBAL_MEM,
+    CONSTANT_MEM,
+  };
+
+ private:
+  ConvolutionTransposed4x4(const OperationDef& definition,
+                           const GpuInfo& gpu_info);
+
+  friend ConvolutionTransposed4x4 CreateConvolutionTransposed4x4(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const ConvolutionTransposedAttributes& attr);
+  friend ConvolutionTransposed4x4 CreateConvolutionTransposed4x4DynamicWeights(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const ConvolutionTransposedAttributes& attr);
+
+  void UploadWeights(
+      const tflite::gpu::Tensor<OHWI, DataType::FLOAT32>& weights,
+      WeightsUploadType weights_upload_type);
+
+  std::vector<int> GetSpatialWeightsRemap() const;
+
+  std::string GenerateConvolutionTransposedCode(
+      const GpuInfo& gpu_info, const OperationDef& op_def,
+      WeightsUploadType weights_upload_type);
+
+  WeightsLayout weights_layout_;
+};
+
+bool IsConvolutionTransposed4x4Supported(
+    const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr);
+
+ConvolutionTransposed4x4 CreateConvolutionTransposed4x4(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr);
+
+ConvolutionTransposed4x4 CreateConvolutionTransposed4x4DynamicWeights(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_4X4_H_

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_4x4_test_util.cc

@@ -0,0 +1,70 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_4x4_test_util.h"
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_4x4.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConvolutionTransposed4x4SimpleWeightsTest(
+    TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 1);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f};
+
+  ConvolutionTransposedAttributes attr;
+  attr.padding.prepended = HW(1, 1);
+  attr.padding.appended = HW(0, 0);
+  attr.stride = HW(2, 2);
+  attr.weights.shape = OHWI(1, 4, 4, 1);
+  attr.weights.data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f,
+                       1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f};
+  attr.bias.shape = Linear(1);
+  attr.bias.data = {0.0f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      ConvolutionTransposed4x4 operation =
+          CreateConvolutionTransposed4x4(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor,
+          std::make_unique<ConvolutionTransposed4x4>(std::move(operation)),
+          BHWC(1, 4, 4, 1), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({0.0f, 1.0f, 1.0f, 1.0f, 2.0f, 6.0f, 6.0f, 4.0f, 2.0f,
+                         6.0f, 6.0f, 4.0f, 2.0f, 5.0f, 5.0f, 3.0f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_4x4_test_util.h

@@ -0,0 +1,31 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_4X4_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_4X4_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConvolutionTransposed4x4SimpleWeightsTest(
+    TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_4X4_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_test_util.cc

@@ -0,0 +1,131 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_test_util.h"
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/precision.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed.h"
+#include "tensorflow/lite/delegates/gpu/common/tensor.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConvolutionTransposedSimpleWeightsTest(
+    TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  ConvolutionTransposedAttributes attr;
+  attr.padding.prepended = HW(0, 0);
+  attr.padding.appended = HW(0, 0);
+  attr.stride = HW(2, 2);
+  attr.weights.shape = OHWI(2, 2, 2, 2);
+  attr.weights.data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f,
+                       1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f};
+  attr.bias.shape = Linear(2);
+  attr.bias.data = {0.0f, 0.0f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      ConvolutionTransposed operation =
+          CreateConvolutionTransposed(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor,
+          std::make_unique<ConvolutionTransposed>(std::move(operation)),
+          BHWC(1, 4, 4, 2), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({1.0f, 1.0f, 1.0f, 1.0f, 5.0f,  5.0f,  5.0f,  5.0f,
+                         1.0f, 1.0f, 1.0f, 1.0f, 5.0f,  5.0f,  5.0f,  5.0f,
+                         9.0f, 9.0f, 9.0f, 9.0f, 13.0f, 13.0f, 13.0f, 13.0f,
+                         9.0f, 9.0f, 9.0f, 9.0f, 13.0f, 13.0f, 13.0f, 13.0f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConvolutionTransposedTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  ConvolutionTransposedAttributes attr;
+  attr.padding.prepended = HW(0, 0);
+  attr.padding.appended = HW(0, 0);
+  attr.stride = HW(2, 2);
+  attr.weights.shape = OHWI(1, 2, 2, 2);
+  attr.weights.data = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f};
+  attr.bias.shape = Linear(1);
+  attr.bias.data = {0.5f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      ConvolutionTransposed operation =
+          CreateConvolutionTransposed(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor,
+          std::make_unique<ConvolutionTransposed>(std::move(operation)),
+          BHWC(1, 4, 4, 1), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear(
+          {2.5f, 4.5f, 8.5f, 18.5f, 6.5f, 8.5f, 28.5f, 38.5f, 14.5f, 32.5f,
+           20.5f, 46.5f, 50.5f, 68.5f, 72.5f, 98.5f},
+          dst_tensor.data, eps));
+
+      // Odd dimensions test (cropped output)
+      {
+        ConvolutionTransposedAttributes attr_odd = attr;
+        attr_odd.padding.appended = HW(1, 1);
+        TensorFloat32 dst_tensor_odd;
+        ConvolutionTransposed operation_odd =
+            CreateConvolutionTransposed(env->GetGpuInfo(), op_def, attr_odd);
+        RETURN_IF_ERROR(env->ExecuteGPUOperation(
+            src_tensor,
+            std::make_unique<ConvolutionTransposed>(std::move(operation_odd)),
+            BHWC(1, 3, 3, 1), &dst_tensor_odd));
+        RETURN_IF_ERROR(PointWiseNear(
+            {2.5f, 4.5f, 8.5f, 6.5f, 8.5f, 28.5f, 14.5f, 32.5f, 20.5f},
+            dst_tensor_odd.data, eps));
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_test_util.h

@@ -0,0 +1,32 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConvolutionTransposedSimpleWeightsTest(
+    TestExecutionEnvironment* env);
+absl::Status ConvolutionTransposedTest(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_thin.cc

@@ -0,0 +1,174 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_thin.h"
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
+
+namespace tflite {
+namespace gpu {
+
+ConvolutionTransposedThin::ConvolutionTransposedThin(
+    const OperationDef& definition, const ConvolutionTransposedAttributes& attr,
+    const GpuInfo& gpu_info)
+    : GPUOperation(definition) {
+  code_ = GenerateConvolutionTransposedCode(
+      definition_, DivideRoundUp(attr.weights.shape.i, 4), attr.weights.shape.o,
+      int2(attr.weights.shape.w, attr.weights.shape.h));
+  if (definition_.precision == CalculationsPrecision::F16 &&
+      gpu_info.IsAdreno() && gpu_info.adreno_info.IsAdreno3xx()) {
+    compiler_options_.push_back(CompilerOptions::kAdrenoFullSimd);
+  }
+}
+
+ConvolutionTransposedThin::ConvolutionTransposedThin(
+    ConvolutionTransposedThin&& operation)
+    : GPUOperation(std::move(operation)) {}
+
+ConvolutionTransposedThin& ConvolutionTransposedThin::operator=(
+    ConvolutionTransposedThin&& operation) {
+  if (this != &operation) {
+    GPUOperation::operator=(std::move(operation));
+  }
+  return *this;
+}
+
+std::string ConvolutionTransposedThin::GenerateConvolutionTransposedCode(
+    const OperationDef& op_def, int src_depth, int dst_channels,
+    const int2& kernel_size) {
+  AddSrcTensor("src_tensor", op_def.src_tensors[0]);
+  AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
+
+  const std::string channel_x = dst_channels == 1 ? "" : ".x";
+  const std::vector<std::string> postfix = {channel_x, ".y", ".z", ".w"};
+  const std::vector<std::string> channel = {".x", ".y", ".z", ".w"};
+
+  const std::string type_postfix =
+      dst_channels == 1 ? "" : std::to_string(dst_channels);
+
+  std::string accum_type;
+
+  switch (op_def.precision) {
+    case CalculationsPrecision::F32:
+    case CalculationsPrecision::F32_F16:
+      accum_type = "float" + type_postfix;
+      break;
+    case CalculationsPrecision::F16:
+      accum_type = "half" + type_postfix;
+      break;
+  }
+
+  std::string c;
+  c += "MAIN_FUNCTION($0) {\n";
+  if (op_def.IsBatchSupported()) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  c += "  int Y = GLOBAL_ID_1;\n";
+  c += "  if (X >= args.src_tensor.Width() || Y >= args.src_tensor.Height()) "
+       "return;\n";
+  c += "  " + accum_type + " r[" + std::to_string(kernel_size.y) + "][" +
+       std::to_string(kernel_size.x) + "];\n";
+  c += "  {\n";
+  c += "  FLT4 src = args.src_tensor.Read(X, Y, 0);\n";
+  int index = 0;
+  for (int y = 0; y < kernel_size.y; ++y) {
+    for (int x = 0; x < kernel_size.x; ++x) {
+      std::string r_s =
+          "  r[" + std::to_string(y) + "][" + std::to_string(x) + "]";
+      for (int d = 0; d < dst_channels; ++d) {
+        c += r_s + postfix[d] + " = dot(src, args.weights.Read(" +
+             std::to_string(index) + "));\n";
+        index++;
+      }
+    }
+  }
+  c += "  }\n";
+  for (int i = 1; i < src_depth; ++i) {
+    c += "  if (X > " + std::to_string(-i) +
+         ") {  // always true, to reduce registers usage\n";
+    c +=
+        "  FLT4 src = args.src_tensor.Read(X, Y, " + std::to_string(i) + ");\n";
+    for (int y = 0; y < kernel_size.y; ++y) {
+      for (int x = 0; x < kernel_size.x; ++x) {
+        std::string r_s =
+            "  r[" + std::to_string(y) + "][" + std::to_string(x) + "]";
+        for (int d = 0; d < dst_channels; ++d) {
+          c += r_s + postfix[d] + " += dot(src, args.weights.Read(" +
+               std::to_string(index) + "));\n";
+          index++;
+        }
+      }
+    }
+    c += "  }\n";
+  }
+  c += "  X *= " + std::to_string(kernel_size.x) + ";\n";
+  c += "  Y *= " + std::to_string(kernel_size.y) + ";\n";
+  for (int y = 0; y < kernel_size.y; ++y) {
+    for (int x = 0; x < kernel_size.x; ++x) {
+      const std::string x_coord = "X + " + std::to_string(x);
+      const std::string y_coord = "Y + " + std::to_string(y);
+      c += "  if (" + x_coord + " < args.dst_tensor.Width() && " + y_coord +
+           " < args.dst_tensor.Height()) {\n";
+      c += "    FLT4 result = args.weights.Read(" + std::to_string(index) +
+           ");\n";
+      for (int d = 0; d < dst_channels; ++d) {
+        c += "    result" + channel[d] + " += r[" + std::to_string(y) + "][" +
+             std::to_string(x) + "]" + postfix[d] + ";\n";
+      }
+      c += "    args.dst_tensor.Write(result, " + x_coord + ", " + y_coord +
+           ", 0);\n";
+      c += "  }\n";
+    }
+  }
+  c += "}\n";
+
+  return c;
+}
+
+int3 ConvolutionTransposedThin::GetGridSize() const {
+  const int grid_x = src_[0]->Width() * dst_[0]->Batch();
+  const int grid_y = src_[0]->Height();
+  const int grid_z = 1;
+  return int3(grid_x, grid_y, grid_z);
+}
+
+bool IsConvolutionTransposedThinSupported(
+    const ConvolutionTransposedAttributes& attr) {
+  return attr.weights.shape.o <= 4 && attr.weights.shape.w == attr.stride.w &&
+         attr.weights.shape.h == attr.stride.h &&
+         attr.padding.prepended.w == 0 && attr.padding.prepended.h == 0 &&
+         attr.padding.appended.w == 0 && attr.padding.appended.h == 0;
+}
+
+ConvolutionTransposedThin CreateConvolutionTransposedThin(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr) {
+  ConvolutionTransposedThin result(definition, attr, gpu_info);
+  result.UploadData(attr.weights, attr.bias);
+  return result;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_thin.h

@@ -0,0 +1,150 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_THIN_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_THIN_H_
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/buffer_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tensor_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/tensor.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+class ConvolutionTransposedThin : public GPUOperation {
+ public:
+  ConvolutionTransposedThin() = default;
+  int3 GetGridSize() const override;
+
+  // Move only
+  ConvolutionTransposedThin(ConvolutionTransposedThin&& operation);
+  ConvolutionTransposedThin& operator=(ConvolutionTransposedThin&& operation);
+  ConvolutionTransposedThin(const ConvolutionTransposedThin&) = delete;
+  ConvolutionTransposedThin& operator=(const ConvolutionTransposedThin&) =
+      delete;
+
+ private:
+  friend ConvolutionTransposedThin CreateConvolutionTransposedThin(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const ConvolutionTransposedAttributes& attr);
+  ConvolutionTransposedThin(const OperationDef& definition,
+                            const ConvolutionTransposedAttributes& attr,
+                            const GpuInfo& gpu_info);
+  template <DataType T>
+  void UploadData(const tflite::gpu::Tensor<OHWI, T>& weights,
+                  const tflite::gpu::Tensor<Linear, T>& biases);
+
+  template <DataType S, typename T>
+  void RearrangeWeightsData(const tflite::gpu::Tensor<OHWI, S>& weights,
+                            absl::Span<T> dst);
+  std::string GenerateConvolutionTransposedCode(const OperationDef& op_def,
+                                                int src_depth, int dst_channels,
+                                                const int2& kernel_size);
+};
+
+template <DataType T>
+void ConvolutionTransposedThin::UploadData(
+    const tflite::gpu::Tensor<OHWI, T>& weights,
+    const tflite::gpu::Tensor<Linear, T>& biases) {
+  const int src_depth = DivideRoundUp(weights.shape.i, 4);
+  const int flt4_count =
+      weights.shape.w * weights.shape.h * src_depth * weights.shape.o;
+
+  const bool f32_weights = definition_.precision == CalculationsPrecision::F32;
+  const int flt4_size = f32_weights ? sizeof(float4) : sizeof(half4);
+
+  BufferDescriptor desc;
+  desc.element_type = f32_weights ? DataType::FLOAT32 : DataType::FLOAT16;
+  desc.element_size = 4;
+  desc.memory_type = MemoryType::CONSTANT;
+  desc.size = flt4_size * (flt4_count + 1);
+  desc.data.resize(desc.size);
+
+  if (f32_weights) {
+    float4* gpu_data = reinterpret_cast<float4*>(desc.data.data());
+    RearrangeWeightsData(weights, absl::MakeSpan(gpu_data, flt4_count));
+    float4 bias_value(0.0f);
+    for (int i = 0; i < weights.shape.o; ++i) {
+      bias_value[i] = biases.data[i];
+    }
+    gpu_data[flt4_count] = bias_value;
+  } else {
+    half4* gpu_data = reinterpret_cast<half4*>(desc.data.data());
+    RearrangeWeightsData(weights, absl::MakeSpan(gpu_data, flt4_count));
+    half4 bias_value(0.0f);
+    for (int i = 0; i < weights.shape.o; ++i) {
+      bias_value[i] = biases.data[i];
+    }
+    gpu_data[flt4_count] = bias_value;
+  }
+
+  args_.AddObject("weights",
+                  std::make_unique<BufferDescriptor>(std::move(desc)));
+}
+
+template <DataType S, typename T>
+void ConvolutionTransposedThin::RearrangeWeightsData(
+    const tflite::gpu::Tensor<OHWI, S>& weights, absl::Span<T> dst) {
+  const int src_depth = DivideRoundUp(weights.shape.i, 4);
+  const int kernel_x = weights.shape.w;
+  const int kernel_y = weights.shape.h;
+
+  int counter = 0;
+  for (int s = 0; s < src_depth; ++s) {
+    for (int y = 0; y < kernel_y; ++y) {
+      for (int x = 0; x < kernel_x; ++x) {
+        std::vector<T> filters(weights.shape.o);
+        for (int j = 0; j < weights.shape.o; ++j) {
+          for (int i = 0; i < 4; ++i) {
+            const int s_ch = s * 4 + i;
+            const int d_ch = j;
+            if (s_ch < weights.shape.i && d_ch < weights.shape.o) {
+              const int f_index = weights.shape.LinearIndex({d_ch, y, x, s_ch});
+              filters[j][i] = weights.data[f_index];
+            } else {
+              filters[j][i] = 0.0f;
+            }
+          }
+        }
+        for (int j = 0; j < weights.shape.o; ++j) {
+          dst[counter++] = filters[j];
+        }
+      }
+    }
+  }
+}
+
+bool IsConvolutionTransposedThinSupported(
+    const ConvolutionTransposedAttributes& attr);
+
+ConvolutionTransposedThin CreateConvolutionTransposedThin(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const ConvolutionTransposedAttributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_THIN_H_

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_thin_test_util.cc

@@ -0,0 +1,110 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_thin_test_util.h"
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_thin.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConvolutionTransposedThinSimpleWeightsTest(
+    TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  ConvolutionTransposedAttributes attr;
+  attr.padding.prepended = HW(0, 0);
+  attr.padding.appended = HW(0, 0);
+  attr.stride = HW(2, 2);
+  attr.weights.shape = OHWI(2, 2, 2, 2);
+  attr.weights.data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f,
+                       1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f};
+  attr.bias.shape = Linear(2);
+  attr.bias.data = {0.0f, 0.0f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      ConvolutionTransposedThin operation =
+          CreateConvolutionTransposedThin(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor,
+          std::make_unique<ConvolutionTransposedThin>(std::move(operation)),
+          BHWC(1, 4, 4, 2), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({1.0f, 1.0f, 1.0f, 1.0f, 5.0f,  5.0f,  5.0f,  5.0f,
+                         1.0f, 1.0f, 1.0f, 1.0f, 5.0f,  5.0f,  5.0f,  5.0f,
+                         9.0f, 9.0f, 9.0f, 9.0f, 13.0f, 13.0f, 13.0f, 13.0f,
+                         9.0f, 9.0f, 9.0f, 9.0f, 13.0f, 13.0f, 13.0f, 13.0f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConvolutionTransposedThinTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  ConvolutionTransposedAttributes attr;
+  attr.padding.prepended = HW(0, 0);
+  attr.padding.appended = HW(0, 0);
+  attr.stride = HW(2, 2);
+  attr.weights.shape = OHWI(1, 2, 2, 2);
+  attr.weights.data = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f};
+  attr.bias.shape = Linear(1);
+  attr.bias.data = {0.5f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      ConvolutionTransposedThin operation =
+          CreateConvolutionTransposedThin(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor,
+          std::make_unique<ConvolutionTransposedThin>(std::move(operation)),
+          BHWC(1, 4, 4, 1), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear(
+          {2.5f, 4.5f, 8.5f, 18.5f, 6.5f, 8.5f, 28.5f, 38.5f, 14.5f, 32.5f,
+           20.5f, 46.5f, 50.5f, 68.5f, 72.5f, 98.5f},
+          dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/convolution_transposed_thin_test_util.h

@@ -0,0 +1,33 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_THIN_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_THIN_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status ConvolutionTransposedThinSimpleWeightsTest(
+    TestExecutionEnvironment* env);
+
+absl::Status ConvolutionTransposedThinTest(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CONVOLUTION_TRANSPOSED_THIN_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/cumsum.cc

@@ -0,0 +1,140 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/gpu/common/tasks/cumsum.h"
+
+#include <string>
+#include <utility>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+
+namespace tflite {
+namespace gpu {
+
+void Cumsum::GetCumsumCode(const OperationDef& op_def) {
+  AddSrcTensor("src_tensor", op_def.src_tensors[0]);
+  AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
+  std::map<Axis, std::string> task_sizes = {
+      {Axis::WIDTH, "args.src_tensor.Width()"},
+      {Axis::HEIGHT, "args.src_tensor.Height()"},
+      {Axis::DEPTH, "args.src_tensor.Depth()"},
+      {Axis::CHANNELS, "args.src_tensor.Slices()"},
+      {Axis::BATCH, "args.src_tensor.Batch()"},
+  };
+  std::string limit = task_sizes[axis_];
+  task_sizes[axis_] = "1";
+  std::map<Axis, std::string> index_name = {
+      {Axis::WIDTH, "X"},    {Axis::HEIGHT, "Y"}, {Axis::DEPTH, "Z"},
+      {Axis::CHANNELS, "S"}, {Axis::BATCH, "B"},
+  };
+  std::string indexes = "X, Y";
+  std::string c;
+  c += "MAIN_FUNCTION($0) {\n";
+  if (definition_.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    indexes += ", Z";
+    c += "  int linear_id = GLOBAL_ID_1;\n";
+    c += "  int Y = linear_id % " + task_sizes[Axis::HEIGHT] + ";\n";
+    c += "  int D = linear_id / " + task_sizes[Axis::HEIGHT] + ";\n";
+    c += "  if (D >= " + task_sizes[Axis::DEPTH] + ") return;\n";
+  } else {
+    c += "  int Y = GLOBAL_ID_1;\n";
+    c += "  if (Y >= " + task_sizes[Axis::HEIGHT] + ") return;\n";
+  }
+  indexes += ", S";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    indexes += ", B";
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / " + task_sizes[Axis::BATCH] + ";\n";
+    c += "  int B = linear_id % " + task_sizes[Axis::BATCH] + ";\n";
+    c += "  if (X >= " + task_sizes[Axis::WIDTH] + ") return;\n";
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+    c += "  if (X >= " + task_sizes[Axis::WIDTH] + ") return;\n";
+  }
+  c += "  int S = GLOBAL_ID_2;\n";
+  c += "  if (S >= " + task_sizes[Axis::CHANNELS] + ") return;\n";
+  if (definition_.precision == CalculationsPrecision::F16) {
+    c += "  float4 res = TO_FLOAT4(args.src_tensor::zero_value);\n";
+  } else {
+    c += "  args.src_tensor::type res = args.src_tensor::zero_value;\n";
+  }
+  c += "  for (; " + index_name[axis_] + " < " + limit + "; " +
+       index_name[axis_] + "++) {\n";
+  if (definition_.precision == CalculationsPrecision::F16) {
+    c +=
+        "    float4 curr = TO_FLOAT4(args.src_tensor.Read(" + indexes + "));\n";
+  } else {
+    c += "    args.src_tensor::type curr = args.src_tensor.Read(" + indexes +
+         ");\n";
+  }
+
+  if (axis_ == Axis::CHANNELS) {
+    if (definition_.precision == CalculationsPrecision::F16) {
+      c += "    res.x = res.w + curr.x;\n";
+      c += "    res.y = res.x + curr.y;\n";
+      c += "    res.z = res.y + curr.z;\n";
+      c += "    res.w = res.z + curr.w;\n";
+    } else {
+      c += "    res.x = res.w + curr.x;\n";
+      c += "    res.y = res.x + curr.y;\n";
+      c += "    res.z = res.y + curr.z;\n";
+      c += "    res.w = res.z + curr.w;\n";
+    }
+  } else {
+    c += "    res += curr;\n";
+  }
+  if (definition_.precision == CalculationsPrecision::F16) {
+    c += "    args.dst_tensor.Write(TO_FLT4(res), " + indexes + ");\n";
+  } else {
+    c += "    args.dst_tensor.Write(res, " + indexes + ");\n";
+  }
+
+  c += "  }\n";
+  c += "}\n";
+  code_ = c;
+}
+
+int3 Cumsum::GetGridSize() const {
+  const int width = axis_ == Axis::WIDTH ? 1 : src_[0]->Width();
+  const int height = axis_ == Axis::HEIGHT ? 1 : src_[0]->Height();
+  const int depth = axis_ == Axis::DEPTH ? 1 : src_[0]->Depth();
+  const int batch = axis_ == Axis::BATCH ? 1 : src_[0]->Batch();
+  const int slices = axis_ == Axis::CHANNELS ? 1 : src_[0]->Slices();
+  const int grid_x = width * batch;
+  const int grid_y = height * depth;
+  const int grid_z = slices;
+  return int3(grid_x, grid_y, grid_z);
+}
+
+Cumsum::Cumsum(Cumsum&& operation)
+    : GPUOperation(std::move(operation)), axis_(operation.axis_) {}
+
+Cumsum& Cumsum::operator=(Cumsum&& operation) {
+  if (this != &operation) {
+    axis_ = operation.axis_;
+    GPUOperation::operator=(std::move(operation));
+  }
+  return *this;
+}
+
+Cumsum CreateCumsum(const OperationDef& definition,
+                    const CumsumAttributes& attr) {
+  Cumsum op(definition, attr.axis);
+  op.GetCumsumCode(definition);
+  return op;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/cumsum.h

@@ -0,0 +1,53 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CUMSUM_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CUMSUM_H_
+
+#include <string>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+class Cumsum : public GPUOperation {
+ public:
+  Cumsum() = default;
+  explicit Cumsum(const OperationDef& definition, Axis axis)
+      : GPUOperation(definition), axis_(axis) {}
+  int3 GetGridSize() const override;
+
+  // Move only
+  Cumsum(Cumsum&& operation);
+  Cumsum& operator=(Cumsum&& operation);
+  Cumsum(const Cumsum&) = delete;
+  Cumsum& operator=(const Cumsum&) = delete;
+  void GetCumsumCode(const OperationDef& op_def);
+
+ private:
+  Axis axis_;
+};
+
+Cumsum CreateCumsum(const OperationDef& definition,
+                    const CumsumAttributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CUMSUM_H_

tensorflow/lite/delegates/gpu/common/tasks/cumsum_test_util.cc

@@ -0,0 +1,225 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/cumsum_test_util.h"
+
+#include <memory>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/cumsum.h"
+
+namespace tflite {
+namespace gpu {
+
+template <Axis axis>
+absl::Status CumsumHWC(TestExecutionEnvironment* env) {
+  Tensor<HWC, DataType::FLOAT32> src_tensor;
+  src_tensor.shape = HWC(8, 6, 4);
+  BHWC shape = BHWC(1, 8, 6, 4);
+  src_tensor.data = {
+      7,  -3, -5, 8,  1,  4,  7,  1,  4,  -9, 6,  -5, 9,  -1, 5,  3,  -6, -2,
+      -3, 3,  -1, -8, -9, -1, 2,  3,  8,  -4, -9, -2, 7,  8,  6,  9,  1,  -1,
+      -9, -6, -9, 1,  -5, 3,  3,  9,  7,  6,  -6, -9, 4,  -9, -6, -7, 6,  -5,
+      -8, 4,  4,  5,  -6, -1, 9,  -3, 3,  9,  7,  4,  -6, -5, -2, 7,  -6, -5,
+      -9, 0,  -7, -3, -3, 6,  -6, -2, -4, 4,  2,  -2, 0,  -6, -4, 5,  -6, -3,
+      1,  -6, 0,  5,  -2, -8, 6,  4,  -6, -5, 8,  3,  -5, -2, -9, -4, 5,  -2,
+      -7, -7, -6, 1,  5,  2,  9,  1,  -2, 9,  1,  -2, 6,  -5, -3, -6, 3,  7,
+      5,  -6, 6,  3,  -3, 4,  6,  8,  7,  -3, 3,  -3, 3,  -4, -1, -4, -7, -4,
+      1,  -5, -4, -5, -1, -9, 9,  -3, 7,  3,  4,  -5, 8,  9,  -6, 5,  -5, -9,
+      2,  3,  4,  5,  -9, 0,  2,  0,  -2, 8,  8,  -5, -6, -6, -7, -1, 7,  -3,
+      2,  6,  4,  -9, 1,  -9, -5, -6, -5, -1, 8,  2};
+  std::map<Axis, std::vector<float>> expected = {
+      {Axis::HEIGHT,
+       {7,   -3,  -5,  8,   1,   4,   7,   1,   4,   -9,  6,   -5,  9,   -1,
+        5,   3,   -6,  -2,  -3,  3,   -1,  -8,  -9,  -1,  9,   0,   3,   4,
+        -8,  2,   14,  9,   10,  0,   7,   -6,  0,   -7,  -4,  4,   -11, 1,
+        0,   12,  6,   -2,  -15, -10, 13,  -9,  -3,  -3,  -2,  -3,  6,   13,
+        14,  5,   1,   -7,  9,   -10, -1,  13,  -4,  5,   -6,  7,   4,   5,
+        -21, -15, 4,   -9,  -10, -6,  -5,  3,   0,   11,  10,  9,   3,   -9,
+        9,   -16, -5,  18,  -10, 2,   -5,  1,   4,   10,  -23, -23, 10,  -5,
+        -16, -11, 3,   6,   -5,  9,   1,   5,   8,   -11, 2,   -23, -11, 19,
+        -5,  4,   4,   2,   2,   19,  -22, -25, 16,  -10, -19, -17, 6,   13,
+        0,   3,   7,   8,   5,   -7,  8,   -15, -4,  16,  -2,  1,   7,   -2,
+        1,   15,  -29, -29, 17,  -15, -23, -22, 5,   4,   9,   0,   14,  11,
+        9,   -12, 16,  -6,  -10, 21,  -7,  -8,  9,   1,   5,   20,  -38, -29,
+        19,  -15, -25, -14, 13,  -1,  3,   -6,  7,   10,  16,  -15, 18,  0,
+        -6,  12,  -6,  -17, 4,   -5,  0,   19,  -30, -27}},
+      {Axis::WIDTH,
+       {7,   -3,  -5,  8,   8,   1,   2,   9,   12,  -8,  8,   4,   21,  -9,
+        13,  7,   15,  -11, 10,  10,  14,  -19, 1,   9,   2,   3,   8,   -4,
+        -7,  1,   15,  4,   -1,  10,  16,  3,   -10, 4,   7,   4,   -15, 7,
+        10,  13,  -8,  13,  4,   4,   4,   -9,  -6,  -7,  10,  -14, -14, -3,
+        14,  -9,  -20, -4,  23,  -12, -17, 5,   30,  -8,  -23, 0,   28,  -1,
+        -29, -5,  -9,  0,   -7,  -3,  -12, 6,   -13, -5,  -16, 10,  -11, -7,
+        -16, 4,   -15, -2,  -22, 1,   -14, -8,  -22, 6,   -16, -16, 6,   4,
+        -6,  -5,  14,  7,   -11, -7,  5,   3,   -6,  -9,  -2,  -4,  -12, -8,
+        3,   -2,  -3,  -7,  1,   7,   -2,  -9,  6,   -5,  -3,  -6,  9,   2,
+        2,   -12, 15,  5,   -1,  -8,  21,  13,  6,   -11, 24,  10,  9,   -15,
+        23,  6,   2,   -19, 1,   -5,  -4,  -5,  0,   -14, 5,   -8,  7,   -11,
+        9,   -13, 15,  -2,  3,   -8,  10,  -11, 5,   -5,  14,  -6,  -4,  -5,
+        2,   0,   -2,  8,   10,  -5,  -8,  2,   3,   -6,  -1,  -1,  5,   0,
+        3,   -10, 6,   -9,  -2,  -16, 1,   -10, 6,   -14}},
+      {Axis::CHANNELS,
+       {7,   4,   -1,  7,   1,   5,   12,  13,  4,   -5,  1,  -4,  9,   8,  13,
+        16,  -6,  -8,  -11, -8,  -1,  -9,  -18, -19, 2,   5,  13,  9,   -9, -11,
+        -4,  4,   6,   15,  16,  15,  -9,  -15, -24, -23, -5, -2,  1,   10, 7,
+        13,  7,   -2,  4,   -5,  -11, -18, 6,   1,   -7,  -3, 4,   9,   3,  2,
+        9,   6,   9,   18,  7,   11,  5,   0,   -2,  5,   -1, -6,  -9,  -9, -16,
+        -19, -3,  3,   -3,  -5,  -4,  0,   2,   0,   0,   -6, -10, -5,  -6, -9,
+        -8,  -14, 0,   5,   3,   -5,  6,   10,  4,   -1,  8,  11,  6,   4,  -9,
+        -13, -8,  -10, -7,  -14, -20, -19, 5,   7,   16,  17, -2,  7,   8,  6,
+        6,   1,   -2,  -8,  3,   10,  15,  9,   6,   9,   6,  10,  6,   14, 21,
+        18,  3,   0,   3,   -1,  -1,  -5,  -12, -16, 1,   -4, -8,  -13, -1, -10,
+        -1,  -4,  7,   10,  14,  9,   8,   17,  11,  16,  -5, -14, -12, -9, 4,
+        9,   0,   0,   2,   2,   0,   8,   8,   3,   -3,  -9, -7,  -8,  -1, -4,
+        2,   8,   12,  3,   1,   -8,  -13, -19, -5,  -6,  2,  4}}};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      OperationDef op_def;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.precision = precision;
+      TensorDescriptor& src = op_def.src_tensors[0];
+      TensorDescriptor& dst = op_def.dst_tensors[0];
+      CumsumAttributes attr = {axis};
+      Cumsum operation = CreateCumsum(op_def, attr);
+      dst.SetBHWCShape(shape);
+      src.UploadData(src_tensor);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {&src}, {&dst}, std::make_unique<Cumsum>(std::move(operation))));
+      TensorFloat32 dst_tensor;
+      dst.DownloadData(&dst_tensor);
+      RETURN_IF_ERROR(PointWiseNear(expected[axis], dst_tensor.data, 0.0f));
+    }
+  }
+  return absl::OkStatus();
+}
+
+template <Axis axis>
+absl::Status CumsumBHWC(TestExecutionEnvironment* env) {
+  Tensor<BHWC, DataType::FLOAT32> src_tensor;
+  src_tensor.shape = BHWC(6, 8, 1, 4);
+  BHWC shape = BHWC(6, 8, 1, 4);
+  src_tensor.data = {
+      -6, -9, -9, 3,  0,  -6, -7, 6,  -3, 5,  -2, 0,  -6, -7, -1, 0,  6,  -7,
+      5,  -7, -1, 0,  9,  2,  -9, 7,  -9, 1,  -8, 1,  -4, -8, -4, -7, -5, 0,
+      -2, 0,  9,  -7, 2,  7,  4,  1,  0,  0,  -7, 7,  2,  1,  6,  3,  7,  -4,
+      -5, -9, 5,  9,  -5, 3,  4,  5,  -2, -1, -5, 9,  -3, 6,  3,  -1, -2, 3,
+      -8, -4, -8, -4, -2, 8,  -5, 2,  -1, 8,  5,  1,  8,  5,  7,  -6, -8, -3,
+      9,  3,  -6, -9, -6, 4,  0,  0,  1,  -5, 3,  6,  -6, 2,  -2, 1,  9,  2,
+      -5, -8, -8, -4, 8,  -1, 7,  -3, 2,  -9, 2,  -2, -1, -6, 7,  2,  -5, -3,
+      2,  0,  3,  -4, -9, 7,  -6, 9,  9,  5,  -1, 5,  8,  8,  1,  -4, -6, -7,
+      2,  1,  -2, -5, 8,  -5, -4, 9,  6,  3,  3,  -9, -2, 1,  -7, 5,  1,  6,
+      0,  -3, -5, 1,  -3, -5, -1, -9, -5, 7,  -7, -3, 5,  6,  8,  -5, -9, -7,
+      1,  -6, -4, -3, -3, 2,  7,  6,  3,  8,  -8, -2};
+  std::map<Axis, std::vector<float>> expected = {
+      {Axis::BATCH,
+       {-6,  -9,  -9,  3,   0,   -6,  -7,  6,   -3,  5,   -2,  0,   -6,  -7,
+        -1,  0,   6,   -7,  5,   -7,  -1,  0,   9,   2,   -9,  7,   -9,  1,
+        -8,  1,   -4,  -8,  -10, -16, -14, 3,   -2,  -6,  2,   -1,  -1,  12,
+        2,   1,   -6,  -7,  -8,  7,   8,   -6,  11,  -4,  6,   -4,  4,   -7,
+        -4,  16,  -14, 4,   -4,  6,   -6,  -9,  -15, -7,  -17, 9,   1,   -7,
+        0,   2,   -9,  8,   -6,  -3,  -8,  1,   -13, 9,   7,   2,   16,  -3,
+        14,  1,   11,  -13, -12, 13,  -5,  7,   -10, -3,  -12, -5,  -15, -7,
+        -16, 4,   4,   -1,  -6,  4,   -11, 9,   3,   -1,  -13, -7,  -21, 5,
+        15,  1,   23,  -6,  16,  -8,  13,  -15, -13, 7,   2,   9,   -15, -6,
+        -10, -5,  -12, -11, -25, 11,  -2,  8,   3,   9,   -12, 14,  11,  7,
+        -12, -11, -27, -2,  17,  2,   21,  -11, 24,  -13, 9,   -6,  -7,  10,
+        5,   0,   -17, -5,  -17, 0,   -11, -5,  -25, 8,   -7,  9,   0,   4,
+        -13, 5,   6,   14,  -19, -14, -22, 4,   25,  -3,  12,  -18, 25,  -19,
+        5,   -9,  -10, 12,  12,  6,   -14, 3,   -25, -2}},
+      {Axis::HEIGHT,
+       {-6,  -9,  -9,  3,   -6, -15, -16, 9,   -9, -10, -18, 9,   -15, -17, -19,
+        9,   -9,  -24, -14, 2,  -10, -24, -5,  4,  -19, -17, -14, 5,   -27, -16,
+        -18, -3,  -4,  -7,  -5, 0,   -6,  -7,  4,  -7,  -4,  0,   8,   -6,  -4,
+        0,   1,   1,   -2,  1,  7,   4,   5,   -3, 2,   -5,  10,  6,   -3,  -2,
+        14,  11,  -5,  -3,  -5, 9,   -3,  6,   -2, 8,   -5,  9,   -10, 4,   -13,
+        5,   -12, 12,  -18, 7,  -13, 20,  -13, 8,  -5,  25,  -6,  2,   -13, 22,
+        3,   5,   -19, 13,  -3, 9,   0,   0,   1,  -5,  3,   6,   -5,  -3,  1,
+        7,   4,   -1,  -4,  -1, -4,  -5,  4,   -2, 3,   -8,  6,   -11, 5,   -10,
+        5,   -17, 12,  -8,  0,  -20, 14,  -8,  3,  -4,  -9,  7,   -3,  5,   0,
+        12,  -4,  10,  8,   20, -3,  6,   2,   13, -1,  7,   0,   8,   7,   2,
+        -4,  17,  13,  5,   -1, 8,   11,  6,   -8, 13,  1,   6,   0,   -3,  -4,
+        7,   -3,  -8,  -5,  -2, -8,  -1,  -12, -5, -3,  5,   -4,  -10, -12, -2,
+        -3,  -16, -16, -5,  -6, -14, -9,  1,   -3, -6,  -17, -1}},
+      {Axis::WIDTH,
+       {-6, -9, -9, 3,  0,  -6, -7, 6,  -3, 5,  -2, 0,  -6, -7, -1, 0,  6,  -7,
+        5,  -7, -1, 0,  9,  2,  -9, 7,  -9, 1,  -8, 1,  -4, -8, -4, -7, -5, 0,
+        -2, 0,  9,  -7, 2,  7,  4,  1,  0,  0,  -7, 7,  2,  1,  6,  3,  7,  -4,
+        -5, -9, 5,  9,  -5, 3,  4,  5,  -2, -1, -5, 9,  -3, 6,  3,  -1, -2, 3,
+        -8, -4, -8, -4, -2, 8,  -5, 2,  -1, 8,  5,  1,  8,  5,  7,  -6, -8, -3,
+        9,  3,  -6, -9, -6, 4,  0,  0,  1,  -5, 3,  6,  -6, 2,  -2, 1,  9,  2,
+        -5, -8, -8, -4, 8,  -1, 7,  -3, 2,  -9, 2,  -2, -1, -6, 7,  2,  -5, -3,
+        2,  0,  3,  -4, -9, 7,  -6, 9,  9,  5,  -1, 5,  8,  8,  1,  -4, -6, -7,
+        2,  1,  -2, -5, 8,  -5, -4, 9,  6,  3,  3,  -9, -2, 1,  -7, 5,  1,  6,
+        0,  -3, -5, 1,  -3, -5, -1, -9, -5, 7,  -7, -3, 5,  6,  8,  -5, -9, -7,
+        1,  -6, -4, -3, -3, 2,  7,  6,  3,  8,  -8, -2}},
+      {Axis::CHANNELS,
+       {-6,  -15, -24, -21, 0,   -6,  -13, -7, -3,  2,  0,   0,   -6,  -13, -14,
+        -14, 6,   -1,  4,   -3,  -1,  -1,  8,  10,  -9, -2,  -11, -10, -8,  -7,
+        -11, -19, -4,  -11, -16, -16, -2,  -2, 7,   0,  2,   9,   13,  14,  0,
+        0,   -7,  0,   2,   3,   9,   12,  7,  3,   -2, -11, 5,   14,  9,   12,
+        4,   9,   7,   6,   -5,  4,   1,   7,  3,   2,  0,   3,   -8,  -12, -20,
+        -24, -2,  6,   1,   3,   -1,  7,   12, 13,  8,  13,  20,  14,  -8,  -11,
+        -2,  1,   -6,  -15, -21, -17, 0,   0,  1,   -4, 3,   9,   3,   5,   -2,
+        -1,  8,   10,  -5,  -13, -21, -25, 8,  7,   14, 11,  2,   -7,  -5,  -7,
+        -1,  -7,  0,   2,   -5,  -8,  -6,  -6, 3,   -1, -10, -3,  -6,  3,   12,
+        17,  -1,  4,   12,  20,  1,   -3,  -9, -16, 2,  3,   1,   -4,  8,   3,
+        -1,  8,   6,   9,   12,  3,   -2,  -1, -8,  -3, 1,   7,   7,   4,   -5,
+        -4,  -7,  -12, -1,  -10, -15, -8,  -7, -10, -5, 1,   8,   3,   -6,  -13,
+        1,   -5,  -9,  -12, -3,  -1,  6,   12, 3,   11, 3,   1}}};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      OperationDef op_def;
+      op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::BHWC});
+      op_def.precision = precision;
+      TensorDescriptor& src = op_def.src_tensors[0];
+      TensorDescriptor& dst = op_def.dst_tensors[0];
+      CumsumAttributes attr = {axis};
+      Cumsum operation = CreateCumsum(op_def, attr);
+      dst.SetBHWCShape(shape);
+      src.UploadData(src_tensor);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {&src}, {&dst}, std::make_unique<Cumsum>(std::move(operation))));
+      TensorFloat32 dst_tensor;
+      dst.DownloadData(&dst_tensor);
+      RETURN_IF_ERROR(PointWiseNear(expected[axis], dst_tensor.data, 0.0f));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status CumsumHWCTest(TestExecutionEnvironment* env) {
+  RETURN_IF_ERROR(CumsumHWC<Axis::HEIGHT>(env));
+  RETURN_IF_ERROR(CumsumHWC<Axis::WIDTH>(env));
+  RETURN_IF_ERROR(CumsumHWC<Axis::CHANNELS>(env));
+  return absl::OkStatus();
+}
+
+absl::Status CumsumBHWCTest(TestExecutionEnvironment* env) {
+  RETURN_IF_ERROR(CumsumBHWC<Axis::BATCH>(env));
+  RETURN_IF_ERROR(CumsumBHWC<Axis::HEIGHT>(env));
+  RETURN_IF_ERROR(CumsumBHWC<Axis::WIDTH>(env));
+  RETURN_IF_ERROR(CumsumBHWC<Axis::CHANNELS>(env));
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/cumsum_test_util.h

@@ -0,0 +1,32 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CUMSUM_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CUMSUM_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status CumsumHWCTest(TestExecutionEnvironment* env);
+
+absl::Status CumsumBHWCTest(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_CUMSUM_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv.cc

@@ -0,0 +1,488 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/task/tensor_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
+#include "tensorflow/lite/delegates/gpu/common/util.h"
+
+namespace tflite {
+namespace gpu {
+
+namespace {
+
+bool IsSpecializedCase(int channel_multiplier) {
+  return channel_multiplier == 1 || channel_multiplier == 2 ||
+         channel_multiplier == 4;
+}
+
+void AppendToBack(const std::string& value, const std::string& delimeter,
+                  std::string* result) {
+  if (!result->empty()) {
+    *result += delimeter;
+  }
+  *result += value;
+}
+
+std::string GetSrcValue(int channel_multiplier,
+                        const std::vector<std::string>& coords,
+                        const std::string& value_name) {
+  std::string coords_str;
+  for (const auto& coord : coords) {
+    AppendToBack(coord, ", ", &coords_str);
+  }
+  std::string c;
+  if (channel_multiplier == 1) {
+    c += "    " + value_name + " = args.src_tensor.Read(" + coords_str +
+         ", S);\n";
+  } else if (channel_multiplier == 2) {
+    c += "    {int s_layer = S / 2;\n";
+    c += "    FLT4 src = args.src_tensor.Read(" + coords_str + ", s_layer);\n";
+    c += "    FLT2 t0 = S % 2 == 0 ? src.xy : src.zw;\n";
+    c += "    " + value_name + " = INIT_FLT4v4(t0.x, t0.x, t0.y, t0.y);}\n";
+  } else if (channel_multiplier == 4) {
+    c += "    {int s_layer = S / 4;\n";
+    c += "    FLT4 src = args.src_tensor.Read(" + coords_str + ", s_layer);\n";
+    c += "    FLT t0 = src.x;\n";
+    c += "    int reminder = S % 4;\n";
+    c += "    if (reminder == 1) t0 = src.y;\n";
+    c += "    if (reminder == 2) t0 = src.z;\n";
+    c += "    if (reminder == 3) t0 = src.w;\n";
+    c += "    " + value_name + " = INIT_FLT4v4(t0, t0, t0, t0);}\n";
+  } else {
+    c += "    {int s_layer = S / args.ch_multiplier;\n";
+    c += "    FLT4 src = args.src_tensor.Read(" + coords_str + ", s_layer);\n";
+    c += "    int s_offset = (S % args.ch_multiplier) * 4;\n";
+    c += "    FLT temp_arr[4] = {src.x, src.y, src.z, src.w};\n";
+    c += "    src.x = temp_arr[(s_offset + 0) / args.ch_multiplier];\n";
+    c += "    src.y = temp_arr[(s_offset + 1) / args.ch_multiplier];\n";
+    c += "    src.z = temp_arr[(s_offset + 2) / args.ch_multiplier];\n";
+    c += "    src.w = temp_arr[(s_offset + 3) / args.ch_multiplier];\n";
+    c += "    " + value_name + " = src;}\n";
+  }
+
+  return c;
+}
+
+std::string GetSrcXYCheck(const GpuInfo& gpu_info,
+                          const TensorDescriptor& src_desc,
+                          const std::string& x_coord,
+                          const std::string& y_coord) {
+  std::string result;
+  if (!src_desc.SupportsZeroClamp(Axis::WIDTH, gpu_info)) {
+    const std::string x_check =
+        x_coord + " >= 0 && " + x_coord + " < args.src_tensor.Width()";
+    AppendToBack(x_check, " && ", &result);
+  }
+  if (!src_desc.SupportsZeroClamp(Axis::HEIGHT, gpu_info)) {
+    const std::string y_check =
+        y_coord + " >= 0 && " + y_coord + " < args.src_tensor.Height()";
+    AppendToBack(y_check, " && ", &result);
+  }
+  return result;
+}
+
+bool UseBuffersForWeights(const GpuInfo& gpu_info) {
+  if (gpu_info.IsApple() &&
+      gpu_info.apple_info.IsFamilyOrLower(AppleInfo::Family::kApple2)) {
+    return false;
+  }
+  return !gpu_info.SupportsImages() || gpu_info.IsMali() ||
+         gpu_info.IsApple() || gpu_info.IsAMD();
+}
+}  // namespace
+
+DepthwiseConv::DepthwiseConv(const OperationDef& definition,
+                             const DepthwiseConvParams& params)
+    : GPUOperation(definition), params_(params) {
+  if (params.UseLocalMem()) {
+    work_group_size_ = params.work_group_size;
+  }
+}
+
+int3 DepthwiseConv::GetGridSize() const {
+  const int grid_x = dst_[0]->Width() * dst_[0]->Batch();
+  const int grid_y = dst_[0]->Height() * dst_[0]->Depth();
+  const int grid_z = dst_[0]->Slices();
+  return int3(grid_x, grid_y, grid_z);
+}
+
+void DepthwiseConv::GetPossibleKernelWorkGroups(
+    TuningType tuning_type, const GpuInfo& gpu_info,
+    const KernelInfo& kernel_info, std::vector<int3>* work_groups) const {
+  if (params_.UseLocalMem()) {
+    work_groups->push_back(work_group_size_);
+    return;
+  }
+  GetPossibleWorkGroups(tuning_type, gpu_info, kernel_info, grid_size_,
+                        work_groups);
+}
+
+std::string DepthwiseConv::GenerateSrcUpload(const GpuInfo& gpu_info) {
+  int cache_size_x = params_.work_group_size.x +
+                     params_.x_kernel_size * params_.x_dilation_size - 1;
+  int cache_size_y = params_.work_group_size.y +
+                     params_.y_kernel_size * params_.y_dilation_size - 1;
+  int groups_x = DivideRoundUp(cache_size_x, params_.work_group_size.x);
+  int groups_y = DivideRoundUp(cache_size_y, params_.work_group_size.y);
+  std::string c;
+  c += "  __local FLT4 spatial_cache[" + std::to_string(cache_size_y) + "][" +
+       std::to_string(cache_size_x) + "];\n";
+  for (int gr_y = 0; gr_y < groups_y; ++gr_y) {
+    std::string y_offset = std::to_string(params_.work_group_size.y * gr_y);
+    std::string ys = "(y_src + " + y_offset + ")";
+    std::string ly = "(LOCAL_ID_1 + " + y_offset + ")";
+    for (int gr_x = 0; gr_x < groups_x; ++gr_x) {
+      std::string x_offset = std::to_string(params_.work_group_size.x * gr_x);
+      std::string xs = "(x_src + " + x_offset + ")";
+      std::string lx = "(LOCAL_ID_0 + " + x_offset + ")";
+      std::string value = "spatial_cache[" + ly + "][" + lx + "]";
+      std::string src_value_read_instructions =
+          GetSrcValue(params_.channel_multiplier, {xs, ys}, value);
+      std::string check =
+          GetSrcXYCheck(gpu_info, definition_.src_tensors[0], xs, ys);
+      c += "  if (" + lx + " < " + std::to_string(cache_size_x) + " && " + ly +
+           " < " + std::to_string(cache_size_y) + ") {\n";
+      if (check.empty()) {
+        c += src_value_read_instructions;
+      } else {
+        c += "    if (" + check + ") {\n";
+        c += src_value_read_instructions;
+        c += "    } else {\n";
+        c += "      " + value + " = INIT_FLT4(0.0f);\n";
+        c += "    }\n";
+      }
+      c += "  }\n";
+    }
+  }
+  return c;
+}
+
+std::string DepthwiseConv::GenerateWeightsUpload(const GpuInfo& gpu_info) {
+  const bool weights_are_buffer = UseBuffersForWeights(gpu_info);
+  auto read_weight = [](bool weights_are_buffer, const std::string& lid,
+                        int work_group_total_size) {
+    if (weights_are_buffer) {
+      return "args.weights.Read(S * args.kernels_total_size + " + lid + ")";
+    } else {
+      return "args.weights.Read(" + lid + ", S)";
+    }
+  };
+  std::string c;
+  const int work_group_total_size = params_.GetWorkGroupTotalSize();
+  c += "  __local FLT4 weights_cache[" +
+       std::to_string(params_.GetKernelsTotalSize()) + "];\n";
+  c += "  int linear_local_id = (LOCAL_ID_2 * GROUP_SIZE_1 + LOCAL_ID_1) * "
+       "GROUP_SIZE_0 + LOCAL_ID_0;\n";
+  const int groups = params_.GetKernelsTotalSize() / work_group_total_size;
+  const int reminder = params_.GetKernelsTotalSize() % work_group_total_size;
+  for (int i = 0; i < groups; ++i) {
+    const std::string lid =
+        "linear_local_id + " + std::to_string(work_group_total_size * i);
+    c += "  weights_cache[" + lid +
+         "] = " + read_weight(weights_are_buffer, lid, work_group_total_size) +
+         ";\n";
+  }
+  if (reminder != 0) {
+    const std::string lid =
+        "linear_local_id + " + std::to_string(work_group_total_size * groups);
+    c += "  if (linear_local_id < " + std::to_string(reminder) + ") {\n";
+    c += "    weights_cache[" + lid +
+         "] = " + read_weight(weights_are_buffer, lid, work_group_total_size) +
+         ";\n";
+    c += "  }\n";
+  }
+  return c;
+}
+
+std::string DepthwiseConv::GenerateCode(const GpuInfo& gpu_info) {
+  const bool weights_are_buffer = UseBuffersForWeights(gpu_info);
+  const bool dynamic_weights = definition_.src_tensors.size() == 2;
+  AddSrcTensor("src_tensor", definition_.src_tensors[0]);
+  if (dynamic_weights) {
+    AddSrcTensor("weights", definition_.src_tensors[1]);
+  }
+  AddDstTensor("dst_tensor", definition_.dst_tensors[0]);
+
+  std::string c;
+
+  const auto& src_desc = definition_.src_tensors[0];
+  c += "MAIN_FUNCTION($0) {\n";
+  if (src_desc.HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  if (src_desc.HasAxis(Axis::DEPTH)) {
+    c += "  int linear_id_1 = GLOBAL_ID_1;\n";
+    c += "  int Y = linear_id_1 / args.dst_tensor.Depth();\n";
+    c += "  int Z = linear_id_1 % args.dst_tensor.Depth();\n";
+  } else {
+    c += "  int Y = GLOBAL_ID_1;\n";
+  }
+  c += "  int S = GLOBAL_ID_2;\n";
+  c += "  int x_src = X * args.stride_x + args.padding_x;\n";
+  c += "  int y_src = Y * args.stride_y + args.padding_y;\n";
+  if (src_desc.HasAxis(Axis::DEPTH)) {
+    c += "  int z_src = Z * args.stride_z + args.padding_z;\n";
+  }
+  if (params_.use_spatial_caching) {
+    c += GenerateSrcUpload(gpu_info);
+  }
+  if (params_.use_weights_caching) {
+    c += GenerateWeightsUpload(gpu_info);
+  }
+  if (params_.UseLocalMem()) {
+    c += "  LOCAL_MEM_BARRIER;\n";
+  }
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "S >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
+  c += "  ACCUM_FLT4 r = INIT_ACCUM_FLT4(0.0f);\n";
+  if (!dynamic_weights && !params_.use_weights_caching) {
+    if (weights_are_buffer) {
+      c += "  int fx_c = S * args.kernels_total_size;\n";
+    } else {
+      c += "  int fx_c = 0;\n";
+    }
+  }
+  std::string kernel_size_x =
+      dynamic_weights ? "args.weights.Width()" : "args.kernel_size_x";
+  std::string kernel_size_y =
+      dynamic_weights ? "args.weights.Height()" : "args.kernel_size_y";
+  std::string kernel_size_z =
+      dynamic_weights ? "args.weights.Depth()" : "args.kernel_size_z";
+  if (params_.UseLocalMem()) {
+    kernel_size_x = std::to_string(params_.x_kernel_size);
+    kernel_size_y = std::to_string(params_.y_kernel_size);
+    kernel_size_z = std::to_string(params_.z_kernel_size);
+  }
+
+  std::string check;
+  std::vector<std::string> coords;
+  if (src_desc.HasAxis(Axis::DEPTH)) {
+    c += "  for (int kz = 0; kz < " + kernel_size_z + "; ++kz) {\n";
+    if (!params_.use_spatial_caching) {
+      c += "    int z_c = z_src + kz * args.dilation_z;\n";
+      coords.insert(coords.begin(), "z_c");
+      if (!src_desc.SupportsZeroClamp(Axis::DEPTH, gpu_info)) {
+        c += "    bool inside_z = z_c >= 0 && z_c < args.src_tensor.Depth();\n";
+        c += "    z_c = clamp(z_c, 0, args.src_tensor.Depth() - 1);\n";
+        AppendToBack("inside_z", " && ", &check);
+      }
+    }
+  }
+  if (src_desc.HasAxis(Axis::HEIGHT)) {
+    c += "  for (int ky = 0; ky < " + kernel_size_y + "; ++ky) {\n";
+    if (!params_.use_spatial_caching) {
+      c += "    int y_c = y_src + ky * args.dilation_y;\n";
+      coords.insert(coords.begin(), "y_c");
+      if (!src_desc.SupportsZeroClamp(Axis::HEIGHT, gpu_info)) {
+        c +=
+            "    bool inside_y = y_c >= 0 && y_c < args.src_tensor.Height();\n";
+        c += "    y_c = clamp(y_c, 0, args.src_tensor.Height() - 1);\n";
+        AppendToBack("inside_y", " && ", &check);
+      }
+    }
+  }
+  if (src_desc.HasAxis(Axis::WIDTH)) {
+    c += "  for (int kx = 0; kx < " + kernel_size_x + "; ++kx) {\n";
+    if (!params_.use_spatial_caching) {
+      c += "    int x_c = x_src + kx * args.dilation_x;\n";
+      coords.insert(coords.begin(), "x_c");
+      if (!src_desc.SupportsZeroClamp(Axis::WIDTH, gpu_info)) {
+        c += "    bool inside_x = x_c >= 0 && x_c < args.src_tensor.Width();\n";
+        c += "    x_c = clamp(x_c, 0, args.src_tensor.Width() - 1);\n";
+        AppendToBack("inside_x", " && ", &check);
+      }
+    }
+  }
+  std::string weight_value;
+  if (params_.use_weights_caching) {
+    std::string weight_index = "ky";
+    if (src_desc.HasAxis(Axis::DEPTH)) {
+      weight_index =
+          "(kz * " + std::to_string(params_.y_kernel_size) + " + ky)";
+    }
+    weight_value = "weights_cache[" + weight_index + " * " +
+                   std::to_string(params_.x_kernel_size) + " + kx]";
+  } else {
+    weight_value = "f";
+    if (dynamic_weights) {
+      c += "    FLT4 f = args.weights.Read(kx, ky, S);\n";
+    } else {
+      if (weights_are_buffer) {
+        c += "    FLT4 f = args.weights.Read(fx_c);\n";
+      } else {
+        c += "    FLT4 f = args.weights.Read(fx_c, S);\n";
+      }
+    }
+  }
+  std::string src_value;
+  if (params_.use_spatial_caching) {
+    std::string loc_x = params_.x_dilation_size == 1
+                            ? "kx"
+                            : "kx * " + std::to_string(params_.x_dilation_size);
+    std::string loc_y = params_.y_dilation_size == 1
+                            ? "ky"
+                            : "ky * " + std::to_string(params_.y_dilation_size);
+    src_value =
+        "spatial_cache[LOCAL_ID_1 + " + loc_y + "][LOCAL_ID_0 + " + loc_x + "]";
+  } else {
+    c += "    FLT4 src_final;\n";
+    src_value = "src_final";
+    c += GetSrcValue(params_.channel_multiplier, coords, src_value);
+    if (!check.empty()) {
+      c += "    src_final = src_final * INIT_FLT(" + check + ");\n";
+    }
+  }
+  c += "    r += TO_ACCUM_TYPE(" + src_value + " * " + weight_value + ");\n";
+  if (!dynamic_weights && !params_.use_weights_caching) {
+    c += "    fx_c++;\n";
+  }
+  if (src_desc.HasAxis(Axis::WIDTH)) {
+    c += "  }\n";
+  }
+  if (src_desc.HasAxis(Axis::HEIGHT)) {
+    c += "  }\n";
+  }
+  if (src_desc.HasAxis(Axis::DEPTH)) {
+    c += "  }\n";
+  }
+  c += "  FLT4 res0 = TO_FLT4(r) + args.biases.Read(S);\n";
+  if (src_desc.HasAxis(Axis::DEPTH)) {
+    c += "  args.dst_tensor.Write(res0, X, Y, Z, S);\n";
+  } else {
+    c += "  args.dst_tensor.Write(res0, X, Y, S);\n";
+  }
+  c += "}\n";
+  return c;
+}
+
+DepthwiseConv CreateDepthwiseConvolution2D(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const DepthwiseConvolution2DAttributes& attr) {
+  const bool weights_are_buffer = UseBuffersForWeights(gpu_info);
+  DepthwiseConv::DepthwiseConvParams params;
+  params.channel_multiplier = attr.weights.shape.o;
+  if (gpu_info.IsAMD()) {
+    if (attr.strides.w == 1 && attr.strides.h == 1 && attr.dilations.w == 1 &&
+        attr.dilations.h == 1 &&
+        attr.weights.shape.w * attr.weights.shape.h >= 10) {
+      params.use_weights_caching = true;
+      params.use_spatial_caching = true;
+      params.x_kernel_size = attr.weights.shape.w;
+      params.y_kernel_size = attr.weights.shape.h;
+      params.x_dilation_size = attr.dilations.w;
+      params.y_dilation_size = attr.dilations.h;
+      params.work_group_size = int3(16, 16, 1);
+    }
+  }
+  DepthwiseConv op(definition, params);
+  op.args_.AddInt("kernel_size_x", attr.weights.shape.w);
+  op.args_.AddInt("stride_x", attr.strides.w);
+  op.args_.AddInt("padding_x", -attr.padding.prepended.w);
+  op.args_.AddInt("dilation_x", attr.dilations.w);
+  op.args_.AddInt("kernel_size_y", attr.weights.shape.h);
+  op.args_.AddInt("stride_y", attr.strides.h);
+  op.args_.AddInt("padding_y", -attr.padding.prepended.h);
+  op.args_.AddInt("dilation_y", attr.dilations.h);
+  op.args_.AddInt("kernels_total_size",
+                  attr.weights.shape.w * attr.weights.shape.h);
+  if (!IsSpecializedCase(attr.weights.shape.o)) {
+    op.args_.AddInt("ch_multiplier", attr.weights.shape.o);
+  }
+  op.code_ = op.GenerateCode(gpu_info);
+  op.UploadWeightsForDWConv2D(attr.weights, weights_are_buffer);
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+
+  TensorDescriptor bias_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), attr.bias);
+  op.args_.AddObject("biases", std::make_unique<TensorDescriptor>(
+                                   std::move(bias_tensor_desc)));
+  return op;
+}
+
+DepthwiseConv CreateDepthwiseConvolution2DDynamicWeights(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const DepthwiseConvolution2DAttributes& attr) {
+  DepthwiseConv::DepthwiseConvParams params;
+  params.channel_multiplier = 1;
+  DepthwiseConv op(definition, params);
+  op.args_.AddInt("stride_x", attr.strides.w);
+  op.args_.AddInt("padding_x", -attr.padding.prepended.w);
+  op.args_.AddInt("dilation_x", attr.dilations.w);
+  op.args_.AddInt("stride_y", attr.strides.h);
+  op.args_.AddInt("padding_y", -attr.padding.prepended.h);
+  op.args_.AddInt("dilation_y", attr.dilations.h);
+  op.code_ = op.GenerateCode(gpu_info);
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+
+  TensorDescriptor bias_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), attr.bias);
+  op.args_.AddObject("biases", std::make_unique<TensorDescriptor>(
+                                   std::move(bias_tensor_desc)));
+  return op;
+}
+
+DepthwiseConv CreateDepthwiseConvolution3D(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const DepthwiseConvolution3DAttributes& attr) {
+  const bool weights_are_buffer = UseBuffersForWeights(gpu_info);
+  DepthwiseConv::DepthwiseConvParams params;
+  params.channel_multiplier = attr.weights.shape.o;
+  DepthwiseConv op(definition, params);
+  op.args_.AddInt("kernel_size_x", attr.weights.shape.w);
+  op.args_.AddInt("stride_x", attr.strides.w);
+  op.args_.AddInt("padding_x", -attr.padding.prepended.w);
+  op.args_.AddInt("dilation_x", attr.dilations.w);
+  op.args_.AddInt("kernel_size_y", attr.weights.shape.h);
+  op.args_.AddInt("stride_y", attr.strides.h);
+  op.args_.AddInt("padding_y", -attr.padding.prepended.h);
+  op.args_.AddInt("dilation_y", attr.dilations.h);
+  op.args_.AddInt("kernel_size_z", attr.weights.shape.d);
+  op.args_.AddInt("stride_z", attr.strides.d);
+  op.args_.AddInt("padding_z", -attr.padding.prepended.d);
+  op.args_.AddInt("dilation_z", attr.dilations.d);
+  op.args_.AddInt(
+      "kernels_total_size",
+      attr.weights.shape.w * attr.weights.shape.h * attr.weights.shape.d);
+  if (!IsSpecializedCase(attr.weights.shape.o)) {
+    op.args_.AddInt("ch_multiplier", attr.weights.shape.o);
+  }
+  op.code_ = op.GenerateCode(gpu_info);
+  op.UploadWeightsForDWConv3D(attr.weights, weights_are_buffer);
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+
+  TensorDescriptor bias_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), attr.bias);
+  op.args_.AddObject("biases", std::make_unique<TensorDescriptor>(
+                                   std::move(bias_tensor_desc)));
+  return op;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv.h

@@ -0,0 +1,271 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_H_
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/buffer_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/tensor.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+class DepthwiseConv : public GPUOperation {
+ public:
+  int3 GetGridSize() const override;
+  void GetPossibleKernelWorkGroups(
+      TuningType tuning_type, const GpuInfo& gpu_info,
+      const KernelInfo& kernel_info,
+      std::vector<int3>* work_groups) const override;
+
+  // Move only
+  DepthwiseConv(DepthwiseConv&& operation) = default;
+  DepthwiseConv& operator=(DepthwiseConv&& operation) = default;
+  DepthwiseConv(const DepthwiseConv&) = delete;
+  DepthwiseConv& operator=(const DepthwiseConv&) = delete;
+
+  friend DepthwiseConv CreateDepthwiseConvolution2D(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const DepthwiseConvolution2DAttributes& attr);
+
+  friend DepthwiseConv CreateDepthwiseConvolution2DDynamicWeights(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const DepthwiseConvolution2DAttributes& attr);
+
+  friend DepthwiseConv CreateDepthwiseConvolution3D(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const DepthwiseConvolution3DAttributes& attr);
+
+ private:
+  struct DepthwiseConvParams {
+    bool UseLocalMem() const {
+      return use_weights_caching || use_spatial_caching;
+    }
+    int GetKernelsTotalSize() const {
+      return x_kernel_size * y_kernel_size * z_kernel_size;
+    }
+    int GetWorkGroupTotalSize() const {
+      return work_group_size.x * work_group_size.y * work_group_size.z;
+    }
+    int channel_multiplier;
+    // Supportd only tensors with Width & Height spatial dimensions
+    // optional, if true, spatial dims will be uploaded to local mem
+    bool use_spatial_caching = false;
+    // optional, if true, weights will be uploaded to local memory
+    bool use_weights_caching = false;
+    // optional, if UsesLocalMem() return true this field must be initialized
+    int3 work_group_size = int3(1, 1, 1);
+
+    // optional, if UsesLocalMem() return true this field must be initialized
+    int x_kernel_size = 1;
+    // optional, if UsesLocalMem() return true this field must be initialized
+    int y_kernel_size = 1;
+    // optional, if UsesLocalMem() return true this field must be initialized
+    int z_kernel_size = 1;
+
+    // optional, if use_spatial_caching true this field must be initialized
+    int x_dilation_size = 1;
+    // optional, if use_spatial_caching true this field must be initialized
+    int y_dilation_size = 1;
+    // optional, if use_spatial_caching true this field must be initialized
+    int z_dilation_size = 1;
+  };
+
+  explicit DepthwiseConv(const OperationDef& definition,
+                         const DepthwiseConvParams& params);
+
+  std::string GenerateSrcUpload(const GpuInfo& gpu_info);
+  std::string GenerateWeightsUpload(const GpuInfo& gpu_info);
+  std::string GenerateCode(const GpuInfo& gpu_info);
+
+  template <DataType T>
+  void UploadWeightsForDWConv2D(const tflite::gpu::Tensor<OHWI, T>& weights,
+                                bool weights_are_buffer);
+
+  template <DataType T>
+  void UploadWeightsForDWConv3D(const tflite::gpu::Tensor<OHWDI, T>& weights,
+                                bool weights_are_buffer);
+
+  DepthwiseConvParams params_;
+};
+
+template <DataType S, typename T>
+void RearrangeWeightsForDWConv2D(const tflite::gpu::Tensor<OHWI, S>& weights,
+                                 absl::Span<T> dst) {
+  const int dst_channels = weights.shape.i * weights.shape.o;
+  const int dst_depth = DivideRoundUp(dst_channels, 4);
+  const int kernel_x = weights.shape.w;
+  const int kernel_y = weights.shape.h;
+
+  int counter = 0;
+  for (int d = 0; d < dst_depth; ++d) {
+    for (int y = 0; y < kernel_y; ++y) {
+      for (int x = 0; x < kernel_x; ++x) {
+        T filter_val;
+        for (int i = 0; i < 4; ++i) {
+          const int d_ch = d * 4 + i;
+          if (d_ch < dst_channels) {
+            const int f_index = weights.shape.LinearIndex(
+                {d_ch % weights.shape.o, y, x, d_ch / weights.shape.o});
+            filter_val[i] = weights.data[f_index];
+          } else {
+            filter_val[i] = 0.0f;
+          }
+        }
+        dst[counter++] = filter_val;
+      }
+    }
+  }
+}
+
+template <DataType T>
+void DepthwiseConv::UploadWeightsForDWConv2D(
+    const tflite::gpu::Tensor<OHWI, T>& weights, bool weights_are_buffer) {
+  const int dst_channels = weights.shape.i * weights.shape.o;
+  const int dst_slices = DivideRoundUp(dst_channels, 4);
+  const int kernel_x = weights.shape.w;
+  const int kernel_y = weights.shape.h;
+
+  const int elements_count = kernel_x * kernel_y * dst_slices;
+
+  const bool fp32_weights = definition_.precision == CalculationsPrecision::F32;
+  const int float4_size = fp32_weights ? 16 : 8;
+
+  std::vector<uint8_t> data(float4_size * elements_count);
+
+  if (fp32_weights) {
+    float4* ptr = reinterpret_cast<float4*>(data.data());
+    RearrangeWeightsForDWConv2D(weights, absl::MakeSpan(ptr, elements_count));
+  } else {
+    half4* ptr = reinterpret_cast<half4*>(data.data());
+    RearrangeWeightsForDWConv2D(weights, absl::MakeSpan(ptr, elements_count));
+  }
+
+  if (weights_are_buffer) {
+    BufferDescriptor desc;
+    desc.element_type = fp32_weights ? DataType::FLOAT32 : DataType::FLOAT16;
+    desc.element_size = 4;
+    desc.size = float4_size * elements_count;
+    desc.data = std::move(data);
+    args_.AddObject("weights", std::make_unique<BufferDescriptor>(desc));
+  } else {
+    TensorDescriptor desc = CreateConstantHWVec4TensorDescriptor(
+        fp32_weights ? DataType::FLOAT32 : DataType::FLOAT16,
+        TensorStorageType::TEXTURE_2D, kernel_x * kernel_y, dst_slices,
+        data.data());
+    args_.AddObject("weights", std::make_unique<TensorDescriptor>(desc));
+  }
+}
+
+template <DataType S, typename T>
+void RearrangeWeightsForDWConv3D(const tflite::gpu::Tensor<OHWDI, S>& weights,
+                                 absl::Span<T> dst) {
+  const int dst_channels = weights.shape.i * weights.shape.o;
+  const int dst_slices = DivideRoundUp(dst_channels, 4);
+  const int kernel_x = weights.shape.w;
+  const int kernel_y = weights.shape.h;
+  const int kernel_z = weights.shape.d;
+
+  int counter = 0;
+  for (int d = 0; d < dst_slices; ++d) {
+    for (int z = 0; z < kernel_z; ++z) {
+      for (int y = 0; y < kernel_y; ++y) {
+        for (int x = 0; x < kernel_x; ++x) {
+          T filter_val;
+          for (int i = 0; i < 4; ++i) {
+            const int d_ch = d * 4 + i;
+            if (d_ch < dst_channels) {
+              const int f_index = weights.shape.LinearIndex(
+                  {d_ch % weights.shape.o, y, x, z, d_ch / weights.shape.o});
+              filter_val[i] = weights.data[f_index];
+            } else {
+              filter_val[i] = 0.0f;
+            }
+          }
+          dst[counter++] = filter_val;
+        }
+      }
+    }
+  }
+}
+
+template <DataType T>
+void DepthwiseConv::UploadWeightsForDWConv3D(
+    const tflite::gpu::Tensor<OHWDI, T>& weights, bool weights_are_buffer) {
+  const int dst_channels = weights.shape.i * weights.shape.o;
+  const int dst_slices = DivideRoundUp(dst_channels, 4);
+  const int kernel_x = weights.shape.w;
+  const int kernel_y = weights.shape.h;
+  const int kernel_z = weights.shape.d;
+
+  const int elements_count = kernel_x * kernel_y * kernel_z * dst_slices;
+
+  const bool fp32_weights = definition_.precision == CalculationsPrecision::F32;
+  const int float4_size = fp32_weights ? 16 : 8;
+
+  std::vector<uint8_t> data(float4_size * elements_count);
+
+  if (fp32_weights) {
+    float4* ptr = reinterpret_cast<float4*>(data.data());
+    RearrangeWeightsForDWConv3D(weights, absl::MakeSpan(ptr, elements_count));
+  } else {
+    half4* ptr = reinterpret_cast<half4*>(data.data());
+    RearrangeWeightsForDWConv3D(weights, absl::MakeSpan(ptr, elements_count));
+  }
+
+  if (weights_are_buffer) {
+    BufferDescriptor desc;
+    desc.element_type = fp32_weights ? DataType::FLOAT32 : DataType::FLOAT16;
+    desc.element_size = 4;
+    desc.size = float4_size * elements_count;
+    desc.data = std::move(data);
+    args_.AddObject("weights",
+                    std::make_unique<BufferDescriptor>(std::move(desc)));
+  } else {
+    TensorDescriptor desc = CreateConstantHWVec4TensorDescriptor(
+        fp32_weights ? DataType::FLOAT32 : DataType::FLOAT16,
+        TensorStorageType::TEXTURE_2D, kernel_x * kernel_y * kernel_z,
+        dst_slices, data.data());
+    args_.AddObject("weights", std::make_unique<TensorDescriptor>(desc));
+  }
+}
+
+DepthwiseConv CreateDepthwiseConvolution2D(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const DepthwiseConvolution2DAttributes& attr);
+
+DepthwiseConv CreateDepthwiseConvolution2DDynamicWeights(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const DepthwiseConvolution2DAttributes& attr);
+
+DepthwiseConv CreateDepthwiseConvolution3D(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const DepthwiseConvolution3DAttributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_H_

tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_3x3.cc

@@ -0,0 +1,352 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_3x3.h"
+
+#include <string>
+#include <utility>
+
+#include "absl/strings/match.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
+
+namespace tflite {
+namespace gpu {
+
+DepthwiseConv3x3::DepthwiseConv3x3(const OperationDef& definition,
+                                   bool weights_are_buffer,
+                                   bool local_mem_uploads,
+                                   const GpuInfo& gpu_info)
+    : GPUOperation(definition), local_mem_uploads_(local_mem_uploads) {
+  work_group_size_ = int3(8, 4, 1);
+  code_ = GenerateDepthwiseConvCode(gpu_info, definition_, weights_are_buffer,
+                                    local_mem_uploads_);
+
+  if (definition_.precision == CalculationsPrecision::F16 &&
+      gpu_info.IsPowerVR()) {
+    compiler_options_.push_back(CompilerOptions::kClFastRelaxedMath);
+  }
+}
+
+DepthwiseConv3x3::DepthwiseConv3x3(DepthwiseConv3x3&& operation)
+    : GPUOperation(std::move(operation)),
+      local_mem_uploads_(operation.local_mem_uploads_) {}
+
+DepthwiseConv3x3& DepthwiseConv3x3::operator=(DepthwiseConv3x3&& operation) {
+  if (this != &operation) {
+    std::swap(local_mem_uploads_, operation.local_mem_uploads_);
+    GPUOperation::operator=(std::move(operation));
+  }
+  return *this;
+}
+
+std::string DepthwiseConv3x3::GenerateDepthwiseConvCode(
+    const GpuInfo& gpu_info, const OperationDef& op_def,
+    bool weights_are_buffer, bool local_mem_uploads) {
+  auto src_desc = op_def.src_tensors[0];
+  AddSrcTensor("src_tensor", src_desc);
+  AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
+
+  std::string c;
+  if (local_mem_uploads && gpu_info.IsApiOpenCl()) {
+    c += "__attribute__((reqd_work_group_size(8, 4, 1)))\n";
+  }
+  c += "MAIN_FUNCTION($0) {\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = (linear_id / args.dst_tensor.Batch()) * 2;\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = GLOBAL_ID_0 * 2;\n";
+  }
+  c += "  int Y = GLOBAL_ID_1 * 2;\n";
+  c += "  int S = GLOBAL_ID_2;\n";
+  c += "   ACCUM_FLT4 r0 = INIT_ACCUM_FLT4(0.0f);\n";
+  c += "   ACCUM_FLT4 r1 = INIT_ACCUM_FLT4(0.0f);\n";
+  c += "   ACCUM_FLT4 r2 = INIT_ACCUM_FLT4(0.0f);\n";
+  c += "   ACCUM_FLT4 r3 = INIT_ACCUM_FLT4(0.0f);\n";
+  if (!local_mem_uploads) {
+    c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() "
+         "|| S >= args.dst_tensor.Slices()) { \n";
+    c += "    return; \n";
+    c += "  } \n";
+  }
+  if (local_mem_uploads) {
+    c += "  __local FLT4 f[10];\n";
+    if (gpu_info.IsApiOpenCl() && gpu_info.IsPowerVR()) {
+      c += "  event_t e = async_work_group_copy(f, args.weights.GetPtr() + S * "
+           "10, 10, 0);\n";
+      c += "  wait_group_events(1, &e);\n";
+    } else {
+      c += "  int local_id = LOCAL_ID_1 * 8 + LOCAL_ID_0;\n";
+      c += "  if (local_id < 10) {\n";
+      c += "    f[local_id] = args.weights.Read(S * 10 + local_id);\n";
+      c += "  }\n";
+      c += "  LOCAL_MEM_BARRIER;\n";
+    }
+  } else if (weights_are_buffer && gpu_info.SupportsPointersInKernels()) {
+    c += "  __global FLT4* f = args.weights.GetPtr() + S * 10;\n";
+  }
+  c += "  FLT4 s0;\n";
+  c += "  FLT4 s1;\n";
+  c += "  FLT4 s2;\n";
+  c += "  FLT4 s3;\n";
+  std::string W[9] = {"f0", "f1", "f2", "f3", "f4", "f5", "f6", "f7", "f8"};
+  std::string bias = "bias";
+  std::string xc[4] = {"X - 1", "X", "X + 1", "X + 2"};
+  std::string yc[4] = {"Y - 1", "Y", "Y + 1", "Y + 2"};
+  if (!weights_are_buffer) {
+    c += "   FLT4 f0 = args.weights.Read(0, S);\n";
+    c += "   FLT4 f1 = args.weights.Read(1, S);\n";
+    c += "   FLT4 f2 = args.weights.Read(2, S);\n";
+    c += "   FLT4 f3 = args.weights.Read(3, S);\n";
+    c += "   FLT4 f4 = args.weights.Read(4, S);\n";
+    c += "   FLT4 f5 = args.weights.Read(5, S);\n";
+    c += "   FLT4 f6 = args.weights.Read(6, S);\n";
+    c += "   FLT4 f7 = args.weights.Read(7, S);\n";
+    c += "   FLT4 f8 = args.weights.Read(8, S);\n";
+  }
+  if (!op_def.src_tensors[0].SupportsZeroClamp(Axis::WIDTH, gpu_info)) {
+    c += "  int x0 = X - 1;\n";
+    c += "  int x1 = X;\n";
+    c += "  int x2 = X + 1;\n";
+    c += "  int x3 = X + 2;\n";
+    c += "  bool x0_in = x0 >= 0 && x0 < args.dst_tensor.Width();\n";
+    c += "  bool x1_in = x1 >= 0 && x1 < args.dst_tensor.Width();\n";
+    c += "  bool x2_in = x2 >= 0 && x2 < args.dst_tensor.Width();\n";
+    c += "  bool x3_in = x3 >= 0 && x3 < args.dst_tensor.Width();\n";
+    c += "  x0 = clamp(x0, 0, args.dst_tensor.Width() - 1);\n";
+    c += "  x1 = clamp(x1, 0, args.dst_tensor.Width() - 1);\n";
+    c += "  x2 = clamp(x2, 0, args.dst_tensor.Width() - 1);\n";
+    c += "  x3 = clamp(x3, 0, args.dst_tensor.Width() - 1);\n";
+    xc[0] = "x0";
+    xc[1] = "x1";
+    xc[2] = "x2";
+    xc[3] = "x3";
+  }
+  if (!op_def.src_tensors[0].SupportsZeroClamp(Axis::HEIGHT, gpu_info)) {
+    c += "  int y0 = Y - 1;\n";
+    c += "  int y1 = Y;\n";
+    c += "  int y2 = Y + 1;\n";
+    c += "  int y3 = Y + 2;\n";
+    c += "  bool y0_in = y0 >= 0 && y0 < args.dst_tensor.Height();\n";
+    c += "  bool y1_in = y1 >= 0 && y1 < args.dst_tensor.Height();\n";
+    c += "  bool y2_in = y2 >= 0 && y2 < args.dst_tensor.Height();\n";
+    c += "  bool y3_in = y3 >= 0 && y3 < args.dst_tensor.Height();\n";
+    c += "  y0 = clamp(y0, 0, args.dst_tensor.Height() - 1);\n";
+    c += "  y1 = clamp(y1, 0, args.dst_tensor.Height() - 1);\n";
+    c += "  y2 = clamp(y2, 0, args.dst_tensor.Height() - 1);\n";
+    c += "  y3 = clamp(y3, 0, args.dst_tensor.Height() - 1);\n";
+    yc[0] = "y0";
+    yc[1] = "y1";
+    yc[2] = "y2";
+    yc[3] = "y3";
+  }
+  if (local_mem_uploads || weights_are_buffer) {
+    const bool use_direct_buffer =
+        !local_mem_uploads && !gpu_info.SupportsPointersInKernels();
+    const std::string fetch_start =
+        use_direct_buffer ? "args.weights.Read(S * 10 + " : "f[";
+    const std::string fetch_end = use_direct_buffer ? ")" : "]";
+    W[0] = fetch_start + "0" + fetch_end;
+    W[1] = fetch_start + "1" + fetch_end;
+    W[2] = fetch_start + "2" + fetch_end;
+    W[3] = fetch_start + "3" + fetch_end;
+    W[4] = fetch_start + "4" + fetch_end;
+    W[5] = fetch_start + "5" + fetch_end;
+    W[6] = fetch_start + "6" + fetch_end;
+    W[7] = fetch_start + "7" + fetch_end;
+    W[8] = fetch_start + "8" + fetch_end;
+    bias = fetch_start + "9" + fetch_end;
+  }
+  auto read_4x_line = [&](int y) {
+    std::string s0_check, s1_check, s2_check, s3_check;
+    if (!op_def.src_tensors[0].SupportsZeroClamp(Axis::WIDTH, gpu_info)) {
+      s0_check += "x0_in";
+      s1_check += "x1_in";
+      s2_check += "x2_in";
+      s3_check += "x3_in";
+    }
+    if (!op_def.src_tensors[0].SupportsZeroClamp(Axis::HEIGHT, gpu_info)) {
+      const std::string y_in = "y" + std::to_string(y) + "_in";
+      s0_check += s0_check.empty() ? y_in : (" && " + y_in);
+      s1_check += s1_check.empty() ? y_in : (" && " + y_in);
+      s2_check += s2_check.empty() ? y_in : (" && " + y_in);
+      s3_check += s3_check.empty() ? y_in : (" && " + y_in);
+    }
+    if (!s0_check.empty()) {
+      s0_check = " * INIT_FLT(" + s0_check + ")";
+    }
+    if (!s1_check.empty()) {
+      s1_check = " * INIT_FLT(" + s1_check + ")";
+    }
+    if (!s2_check.empty()) {
+      s2_check = " * INIT_FLT(" + s2_check + ")";
+    }
+    if (!s3_check.empty()) {
+      s3_check = " * INIT_FLT(" + s3_check + ")";
+    }
+    c += "    s0 = args.src_tensor.Read(" + xc[0] + ", " + yc[y] + ", S)" +
+         s0_check + ";\n";
+    c += "    s1 = args.src_tensor.Read(" + xc[1] + ", " + yc[y] + ", S)" +
+         s1_check + ";\n";
+    c += "    s2 = args.src_tensor.Read(" + xc[2] + ", " + yc[y] + ", S)" +
+         s2_check + ";\n";
+    c += "    s3 = args.src_tensor.Read(" + xc[3] + ", " + yc[y] + ", S)" +
+         s3_check + ";\n";
+  };
+  c += "  {\n";
+  read_4x_line(0);
+  c += "    r0 += TO_ACCUM_TYPE(" + W[0] + " * s0);\n";
+  c += "    r0 += TO_ACCUM_TYPE(" + W[1] + " * s1);\n";
+  c += "    r1 += TO_ACCUM_TYPE(" + W[0] + " * s1);\n";
+  c += "    r0 += TO_ACCUM_TYPE(" + W[2] + " * s2);\n";
+  c += "    r1 += TO_ACCUM_TYPE(" + W[1] + " * s2);\n";
+  c += "    r1 += TO_ACCUM_TYPE(" + W[2] + " * s3);\n";
+  c += "  }\n";
+  c += "  {\n";
+  read_4x_line(1);
+  c += "    r0 += TO_ACCUM_TYPE(" + W[3] + " * s0);\n";
+  c += "    r2 += TO_ACCUM_TYPE(" + W[0] + " * s0);\n";
+  c += "    r0 += TO_ACCUM_TYPE(" + W[4] + " * s1);\n";
+  c += "    r1 += TO_ACCUM_TYPE(" + W[3] + " * s1);\n";
+  c += "    r2 += TO_ACCUM_TYPE(" + W[1] + " * s1);\n";
+  c += "    r3 += TO_ACCUM_TYPE(" + W[0] + " * s1);\n";
+  c += "    r0 += TO_ACCUM_TYPE(" + W[5] + " * s2);\n";
+  c += "    r1 += TO_ACCUM_TYPE(" + W[4] + " * s2);\n";
+  c += "    r2 += TO_ACCUM_TYPE(" + W[2] + " * s2);\n";
+  c += "    r3 += TO_ACCUM_TYPE(" + W[1] + " * s2);\n";
+  c += "    r1 += TO_ACCUM_TYPE(" + W[5] + " * s3);\n";
+  c += "    r3 += TO_ACCUM_TYPE(" + W[2] + " * s3);\n";
+  c += "  }\n";
+  c += "  {\n";
+  read_4x_line(2);
+  c += "    r0 += TO_ACCUM_TYPE(" + W[6] + " * s0);\n";
+  c += "    r2 += TO_ACCUM_TYPE(" + W[3] + " * s0);\n";
+  c += "    r0 += TO_ACCUM_TYPE(" + W[7] + " * s1);\n";
+  c += "    r1 += TO_ACCUM_TYPE(" + W[6] + " * s1);\n";
+  c += "    r2 += TO_ACCUM_TYPE(" + W[4] + " * s1);\n";
+  c += "    r3 += TO_ACCUM_TYPE(" + W[3] + " * s1);\n";
+  c += "    r0 += TO_ACCUM_TYPE(" + W[8] + " * s2);\n";
+  c += "    r1 += TO_ACCUM_TYPE(" + W[7] + " * s2);\n";
+  c += "    r2 += TO_ACCUM_TYPE(" + W[5] + " * s2);\n";
+  c += "    r3 += TO_ACCUM_TYPE(" + W[4] + " * s2);\n";
+  c += "    r1 += TO_ACCUM_TYPE(" + W[8] + " * s3);\n";
+  c += "    r3 += TO_ACCUM_TYPE(" + W[5] + " * s3);\n";
+  c += "  }\n";
+  c += "  {\n";
+  read_4x_line(3);
+  c += "    r2 += TO_ACCUM_TYPE(" + W[6] + " * s0);\n";
+  c += "    r2 += TO_ACCUM_TYPE(" + W[7] + " * s1);\n";
+  c += "    r3 += TO_ACCUM_TYPE(" + W[6] + " * s1);\n";
+  c += "    r2 += TO_ACCUM_TYPE(" + W[8] + " * s2);\n";
+  c += "    r3 += TO_ACCUM_TYPE(" + W[7] + " * s2);\n";
+  c += "    r3 += TO_ACCUM_TYPE(" + W[8] + " * s3);\n";
+  c += "  }\n";
+  if (!weights_are_buffer) {
+    c += "   FLT4 bias = args.weights.Read(9, S);\n";
+  }
+  c += "  r0 += TO_ACCUM_TYPE(" + bias + ");\n";
+  c += "  r1 += TO_ACCUM_TYPE(" + bias + ");\n";
+  c += "  r2 += TO_ACCUM_TYPE(" + bias + ");\n";
+  c += "  r3 += TO_ACCUM_TYPE(" + bias + ");\n";
+  if (local_mem_uploads) {
+    c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() "
+         "|| S >= args.dst_tensor.Slices()) { \n";
+    c += "    return; \n";
+    c += "  } \n";
+  }
+  c += "  if(X + 0 < args.dst_tensor.Width() && Y + 0 < "
+       "args.dst_tensor.Height()) {\n";
+  c += "    FLT4 result = TO_FLT4(r0);\n";
+  c += "    args.dst_tensor.Write(result, X + 0, Y + 0, S);\n";
+  c += "  }\n";
+  c += "  if(X + 1 < args.dst_tensor.Width() && Y + 0 < "
+       "args.dst_tensor.Height()) {\n";
+  c += "    FLT4 result = TO_FLT4(r1);\n";
+  c += "    args.dst_tensor.Write(result, X + 1, Y + 0, S);\n";
+  c += "  }\n";
+  c += "  if(X + 0 < args.dst_tensor.Width() && Y + 1 < "
+       "args.dst_tensor.Height()) {\n";
+  c += "    FLT4 result = TO_FLT4(r2);\n";
+  c += "    args.dst_tensor.Write(result, X + 0, Y + 1, S);\n";
+  c += "  }\n";
+  c += "  if(X + 1 < args.dst_tensor.Width() && Y + 1 < "
+       "args.dst_tensor.Height()) {\n";
+  c += "    FLT4 result = TO_FLT4(r3);\n";
+  c += "    args.dst_tensor.Write(result, X + 1, Y + 1, S);\n";
+  c += "  }\n";
+  c += "}\n";
+
+  return c;
+}
+
+int3 DepthwiseConv3x3::GetGridSize() const {
+  const int grid_x = DivideRoundUp(dst_[0]->Width(), 2) * dst_[0]->Batch();
+  const int grid_y = DivideRoundUp(dst_[0]->Height(), 2);
+  const int grid_z = dst_[0]->Slices();
+  return int3(grid_x, grid_y, grid_z);
+}
+
+void DepthwiseConv3x3::GetPossibleKernelWorkGroups(
+    TuningType tuning_type, const GpuInfo& gpu_info,
+    const KernelInfo& kernel_info, std::vector<int3>* work_groups) const {
+  if (local_mem_uploads_) {
+    work_groups->push_back(work_group_size_);
+  } else {
+    GetPossibleWorkGroups(tuning_type, gpu_info, kernel_info, grid_size_,
+                          work_groups);
+  }
+}
+
+bool IsDepthwiseConv3x3Supported(const GpuInfo& gpu_info,
+                                 const DepthwiseConvolution2DAttributes& attr) {
+  if (gpu_info.IsApiOpenCl() && gpu_info.IsAdreno()) {
+    const std::string kBadDriver =
+        "OpenCL 2.0 QUALCOMM build: commit #7daed58 changeid #I7ece6fe30d "
+        "Date: 10/19/16";
+    if (absl::StrContains(gpu_info.opencl_info.platform_version, kBadDriver)) {
+      return false;
+    }
+  }
+  return attr.weights.shape.o == 1 && attr.dilations.w == 1 &&
+         attr.dilations.h == 1 && attr.weights.shape.w == 3 &&
+         attr.weights.shape.h == 3 && attr.strides.w == 1 &&
+         attr.strides.h == 1 && attr.padding.prepended.w == 1 &&
+         attr.padding.prepended.h == 1 && attr.padding.appended.w == 1 &&
+         attr.padding.appended.h == 1;
+}
+
+DepthwiseConv3x3 CreateDepthwiseConv3x3(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const DepthwiseConvolution2DAttributes& attr) {
+  bool weights_are_buffer = !gpu_info.SupportsImages() ||
+                            gpu_info.IsPowerVR() || gpu_info.IsMali() ||
+                            gpu_info.IsApple();
+  bool local_mem_uploads =
+      (weights_are_buffer && gpu_info.IsPowerVR() && gpu_info.IsApiOpenCl() &&
+       gpu_info.opencl_info.dedicated_local_memory) ||
+      (gpu_info.IsApple() &&
+       gpu_info.apple_info.IsLocalMemoryPreferredOverGlobal());
+  DepthwiseConv3x3 result(definition, weights_are_buffer, local_mem_uploads,
+                          gpu_info);
+  result.UploadWeightsAndBiases(attr.weights, attr.bias, weights_are_buffer);
+  return result;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_3x3.h

@@ -0,0 +1,160 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_3X3_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_3X3_H_
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/buffer_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tensor_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/tensor.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+class DepthwiseConv3x3 : public GPUOperation {
+ public:
+  DepthwiseConv3x3() = default;
+  void GetPossibleKernelWorkGroups(
+      TuningType tuning_type, const GpuInfo& gpu_info,
+      const KernelInfo& kernel_info,
+      std::vector<int3>* work_groups) const override;
+  int3 GetGridSize() const override;
+
+  // Move only
+  DepthwiseConv3x3(DepthwiseConv3x3&& operation);
+  DepthwiseConv3x3& operator=(DepthwiseConv3x3&& operation);
+  DepthwiseConv3x3(const DepthwiseConv3x3&) = delete;
+  DepthwiseConv3x3& operator=(const DepthwiseConv3x3&) = delete;
+
+ private:
+  explicit DepthwiseConv3x3(const OperationDef& definition,
+                            bool weights_are_buffer, bool local_mem_uploads,
+                            const GpuInfo& gpu_info);
+  template <DataType T>
+  void UploadWeightsAndBiases(const tflite::gpu::Tensor<OHWI, T>& weights,
+                              const tflite::gpu::Tensor<Linear, T>& biases,
+                              bool weights_are_buffer);
+
+  friend DepthwiseConv3x3 CreateDepthwiseConv3x3(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const DepthwiseConvolution2DAttributes& attr);
+
+  template <DataType S, typename T>
+  void RearrangeWeightsAndBiasesData(
+      const tflite::gpu::Tensor<OHWI, S>& weights,
+      const tflite::gpu::Tensor<Linear, S>& biases, absl::Span<T> dst);
+
+  std::string GenerateDepthwiseConvCode(const GpuInfo& gpu_info,
+                                        const OperationDef& op_def,
+                                        bool weights_are_buffer,
+                                        bool local_mem_uploads);
+
+  bool local_mem_uploads_;
+};
+
+template <DataType T>
+void DepthwiseConv3x3::UploadWeightsAndBiases(
+    const tflite::gpu::Tensor<OHWI, T>& weights,
+    const tflite::gpu::Tensor<Linear, T>& biases, bool weights_are_buffer) {
+  const int src_depth = DivideRoundUp(weights.shape.i, 4);
+  int texture_width = 10;  // 3x3 kernel + 1 bias
+  int texture_height = src_depth;
+  const int elements_count = texture_width * texture_height;
+  const bool fp32_weights = definition_.precision == CalculationsPrecision::F32;
+  const int float4_size = fp32_weights ? 16 : 8;
+
+  std::vector<uint8_t> data(float4_size * elements_count);
+  if (fp32_weights) {
+    float4* ptr = reinterpret_cast<float4*>(data.data());
+    RearrangeWeightsAndBiasesData(weights, biases,
+                                  absl::MakeSpan(ptr, elements_count));
+  } else {
+    half4* ptr = reinterpret_cast<half4*>(data.data());
+    RearrangeWeightsAndBiasesData(weights, biases,
+                                  absl::MakeSpan(ptr, elements_count));
+  }
+
+  if (weights_are_buffer) {
+    BufferDescriptor desc;
+    desc.element_type = fp32_weights ? DataType::FLOAT32 : DataType::FLOAT16;
+    desc.element_size = 4;
+    desc.size = float4_size * elements_count;
+    desc.data = std::move(data);
+    args_.AddObject("weights",
+                    std::make_unique<BufferDescriptor>(std::move(desc)));
+  } else {
+    TensorDescriptor desc = CreateConstantHWVec4TensorDescriptor(
+        fp32_weights ? DataType::FLOAT32 : DataType::FLOAT16,
+        TensorStorageType::TEXTURE_2D, texture_width, texture_height,
+        data.data());
+    args_.AddObject("weights", std::make_unique<TensorDescriptor>(desc));
+  }
+}
+
+template <DataType S, typename T>
+void DepthwiseConv3x3::RearrangeWeightsAndBiasesData(
+    const tflite::gpu::Tensor<OHWI, S>& weights,
+    const tflite::gpu::Tensor<Linear, S>& biases, absl::Span<T> dst) {
+  const int src_depth = DivideRoundUp(weights.shape.i, 4);
+
+  int counter = 0;
+  for (int s = 0; s < src_depth; ++s) {
+    for (int y = 0; y < 3; ++y) {
+      for (int x = 0; x < 3; ++x) {
+        T filter_val;
+        for (int i = 0; i < 4; ++i) {
+          const int s_ch = s * 4 + i;
+          if (s_ch < weights.shape.i) {
+            const int f_index = weights.shape.LinearIndex({0, y, x, s_ch});
+            filter_val[i] = weights.data[f_index];
+          } else {
+            filter_val[i] = 0.0f;
+          }
+        }
+        dst[counter++] = filter_val;
+      }
+    }
+
+    T bias_val;
+    for (int i = 0; i < 4; ++i) {
+      const int dst_ch = s * 4 + i;
+      bias_val[i] = dst_ch >= biases.shape.v ? 0.0f : biases.data[dst_ch];
+    }
+    dst[counter++] = bias_val;
+  }
+}
+
+bool IsDepthwiseConv3x3Supported(const GpuInfo& gpu_info,
+                                 const DepthwiseConvolution2DAttributes& attr);
+
+DepthwiseConv3x3 CreateDepthwiseConv3x3(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const DepthwiseConvolution2DAttributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_3X3_H_

tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_3x3_stride_h2.cc

@@ -0,0 +1,251 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_3x3_stride_h2.h"
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
+
+namespace tflite {
+namespace gpu {
+namespace {
+std::string GetKernelDepthWiseConv3x3StrideH2(const GpuInfo& gpu_info,
+                                              const OperationDef& definition,
+                                              bool weights_are_buffer,
+                                              bool local_mem_uploads) {
+  std::string c = "MAIN_FUNCTION($0) {\n";
+  if (definition.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  c += R"(
+  int Y = GLOBAL_ID_1 * 2;
+  int S = GLOBAL_ID_2;
+
+  ACCUM_FLT4 r0 = INIT_ACCUM_FLT4(0.0f);
+  ACCUM_FLT4 l0 = INIT_ACCUM_FLT4(0.0f);
+)";
+  if (local_mem_uploads) {
+    c += "  __local FLT4 f[10];\n";
+    c += "  int local_id = LOCAL_ID_1 * 8 + LOCAL_ID_0;\n";
+    c += "  if (local_id < 10) {\n";
+    c += "    f[local_id] = args.weights.Read(S * 10 + local_id);\n";
+    c += "  }\n";
+    c += "  LOCAL_MEM_BARRIER;\n";
+  } else if (weights_are_buffer && gpu_info.SupportsPointersInKernels()) {
+    c += "  __global FLT4* f = args.weights.GetPtr() + S * 10;\n";
+  }
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() "
+       "|| S >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
+  c += "  FLT4 s0, s1, s2;\n";
+  c += "  int x0 = X * args.stride_x + args.padding_x;\n";
+  c += "  int x1 = X * args.stride_x + args.padding_x + args.dilation_x;\n";
+  c += "  int x2 = X * args.stride_x + args.padding_x + 2 * args.dilation_x;\n";
+  c += "  int y0 = Y * 2 + args.padding_y;\n";
+  c += "  int y1 = Y * 2 + args.padding_y + 1;\n";
+  c += "  int y2 = Y * 2 + args.padding_y + 2;\n";
+  c += "  int y3 = Y * 2 + args.padding_y + 3;\n";
+  c += "  int y4 = Y * 2 + args.padding_y + 4;\n";
+  std::string W[9] = {"f0", "f1", "f2", "f3", "f4", "f5", "f6", "f7", "f8"};
+  std::string bias = "bias";
+  if (!weights_are_buffer) {
+    c += "   FLT4 f0 = args.weights.Read(0, S);\n";
+    c += "   FLT4 f1 = args.weights.Read(1, S);\n";
+    c += "   FLT4 f2 = args.weights.Read(2, S);\n";
+    c += "   FLT4 f3 = args.weights.Read(3, S);\n";
+    c += "   FLT4 f4 = args.weights.Read(4, S);\n";
+    c += "   FLT4 f5 = args.weights.Read(5, S);\n";
+    c += "   FLT4 f6 = args.weights.Read(6, S);\n";
+    c += "   FLT4 f7 = args.weights.Read(7, S);\n";
+    c += "   FLT4 f8 = args.weights.Read(8, S);\n";
+  }
+  if (!definition.src_tensors[0].SupportsZeroClamp(Axis::WIDTH, gpu_info)) {
+    c += "  bool x0_in = x0 >= 0 && x0 < args.src_tensor.Width();\n";
+    c += "  bool x1_in = x1 >= 0 && x1 < args.src_tensor.Width();\n";
+    c += "  bool x2_in = x2 >= 0 && x2 < args.src_tensor.Width();\n";
+    c += "  x0 = clamp(x0, 0, args.src_tensor.Width() - 1);\n";
+    c += "  x1 = clamp(x1, 0, args.src_tensor.Width() - 1);\n";
+    c += "  x2 = clamp(x2, 0, args.src_tensor.Width() - 1);\n";
+  }
+  if (!definition.src_tensors[0].SupportsZeroClamp(Axis::HEIGHT, gpu_info)) {
+    c += "  bool y0_in = y0 >= 0 && y0 < args.src_tensor.Height();\n";
+    c += "  bool y1_in = y1 >= 0 && y1 < args.src_tensor.Height();\n";
+    c += "  bool y2_in = y2 >= 0 && y2 < args.src_tensor.Height();\n";
+    c += "  bool y3_in = y3 >= 0 && y3 < args.src_tensor.Height();\n";
+    c += "  bool y4_in = y4 >= 0 && y4 < args.src_tensor.Height();\n";
+    c += "  y0 = clamp(y0, 0, args.src_tensor.Height() - 1);\n";
+    c += "  y1 = clamp(y1, 0, args.src_tensor.Height() - 1);\n";
+    c += "  y2 = clamp(y2, 0, args.src_tensor.Height() - 1);\n";
+    c += "  y3 = clamp(y3, 0, args.src_tensor.Height() - 1);\n";
+    c += "  y4 = clamp(y4, 0, args.src_tensor.Height() - 1);\n";
+  }
+
+  if (local_mem_uploads || weights_are_buffer) {
+    const bool use_direct_buffer =
+        !local_mem_uploads && !gpu_info.SupportsPointersInKernels();
+    const std::string fetch_start =
+        use_direct_buffer ? "args.weights.Read(S * 10 + " : "f[";
+    const std::string fetch_end = use_direct_buffer ? ")" : "]";
+    W[0] = fetch_start + "0" + fetch_end;
+    W[1] = fetch_start + "1" + fetch_end;
+    W[2] = fetch_start + "2" + fetch_end;
+    W[3] = fetch_start + "3" + fetch_end;
+    W[4] = fetch_start + "4" + fetch_end;
+    W[5] = fetch_start + "5" + fetch_end;
+    W[6] = fetch_start + "6" + fetch_end;
+    W[7] = fetch_start + "7" + fetch_end;
+    W[8] = fetch_start + "8" + fetch_end;
+    bias = fetch_start + "9" + fetch_end;
+  }
+  auto read_3x_line = [&](int y) {
+    std::string s0_check, s1_check, s2_check;
+    if (!definition.src_tensors[0].SupportsZeroClamp(Axis::WIDTH, gpu_info)) {
+      s0_check += "x0_in";
+      s1_check += "x1_in";
+      s2_check += "x2_in";
+    }
+    if (!definition.src_tensors[0].SupportsZeroClamp(Axis::HEIGHT, gpu_info)) {
+      const std::string y_in = "y" + std::to_string(y) + "_in";
+      s0_check += s0_check.empty() ? y_in : (" && " + y_in);
+      s1_check += s1_check.empty() ? y_in : (" && " + y_in);
+      s2_check += s2_check.empty() ? y_in : (" && " + y_in);
+    }
+    if (!s0_check.empty()) {
+      s0_check = " * INIT_FLT(" + s0_check + ")";
+    }
+    if (!s1_check.empty()) {
+      s1_check = " * INIT_FLT(" + s1_check + ")";
+    }
+    if (!s2_check.empty()) {
+      s2_check = " * INIT_FLT(" + s2_check + ")";
+    }
+
+    const std::string yc = "y" + std::to_string(y);
+    c += "    s0 = args.src_tensor.Read(x0, " + yc + ", S)" + s0_check + ";\n";
+    c += "    s1 = args.src_tensor.Read(x1, " + yc + ", S)" + s1_check + ";\n";
+    c += "    s2 = args.src_tensor.Read(x2, " + yc + ", S)" + s2_check + ";\n";
+  };
+  read_3x_line(0);
+  c += "    r0 += TO_ACCUM_TYPE(" + W[0] + " * s0);\n";
+  c += "    r0 += TO_ACCUM_TYPE(" + W[1] + " * s1);\n";
+  c += "    r0 += TO_ACCUM_TYPE(" + W[2] + " * s2);\n";
+  read_3x_line(1);
+  c += "    r0 += TO_ACCUM_TYPE(" + W[3] + " * s0);\n";
+  c += "    r0 += TO_ACCUM_TYPE(" + W[4] + " * s1);\n";
+  c += "    r0 += TO_ACCUM_TYPE(" + W[5] + " * s2);\n";
+  read_3x_line(2);
+  c += "    r0 += TO_ACCUM_TYPE(" + W[6] + " * s0);\n";
+  c += "    r0 += TO_ACCUM_TYPE(" + W[7] + " * s1);\n";
+  c += "    r0 += TO_ACCUM_TYPE(" + W[8] + " * s2);\n";
+  c += "    l0 += TO_ACCUM_TYPE(" + W[0] + " * s0);\n";
+  c += "    l0 += TO_ACCUM_TYPE(" + W[1] + " * s1);\n";
+  c += "    l0 += TO_ACCUM_TYPE(" + W[2] + " * s2);\n";
+  read_3x_line(3);
+  c += "    l0 += TO_ACCUM_TYPE(" + W[3] + " * s0);\n";
+  c += "    l0 += TO_ACCUM_TYPE(" + W[4] + " * s1);\n";
+  c += "    l0 += TO_ACCUM_TYPE(" + W[5] + " * s2);\n";
+  read_3x_line(4);
+  c += "    l0 += TO_ACCUM_TYPE(" + W[6] + " * s0);\n";
+  c += "    l0 += TO_ACCUM_TYPE(" + W[7] + " * s1);\n";
+  c += "    l0 += TO_ACCUM_TYPE(" + W[8] + " * s2);\n";
+  if (!weights_are_buffer) {
+    c += "   FLT4 bias = args.weights.Read(9, S);\n";
+  }
+  c += "  r0 += TO_ACCUM_TYPE(" + bias + ");\n";
+  c += "  l0 += TO_ACCUM_TYPE(" + bias + ");\n";
+  c += R"(
+  if (Y < args.dst_tensor.Height()) {
+    FLT4 value = TO_FLT4(r0);
+    args.dst_tensor.Write(value, X, Y, S);
+  }
+  if (Y + 1 < args.dst_tensor.Height()) {
+    FLT4 value = TO_FLT4(l0);
+    args.dst_tensor.Write(value, X, Y + 1, S);
+  }
+}
+)";
+
+  return c;
+}
+
+}  // namespace
+
+int3 DepthWiseConv3x3StrideH2::GetGridSize() const {
+  const int grid_x = dst_[0]->Width() * dst_[0]->Batch();
+  const int grid_y = DivideRoundUp(dst_[0]->Height(), 2);
+  const int grid_z = dst_[0]->Slices();
+  return int3(grid_x, grid_y, grid_z);
+}
+
+void DepthWiseConv3x3StrideH2::GetPossibleKernelWorkGroups(
+    TuningType tuning_type, const GpuInfo& gpu_info,
+    const KernelInfo& kernel_info, std::vector<int3>* work_groups) const {
+  if (local_mem_uploads_) {
+    work_groups->push_back(work_group_size_);
+  } else {
+    GetPossibleWorkGroups(tuning_type, gpu_info, kernel_info, grid_size_,
+                          work_groups);
+  }
+}
+
+DepthWiseConv3x3StrideH2 CreateDepthWiseConv3x3StrideH2(
+    const OperationDef& definition,
+    const DepthwiseConvolution2DAttributes& attr, const GpuInfo& gpu_info) {
+  bool weights_are_buffer = !gpu_info.SupportsImages() ||
+                            gpu_info.IsPowerVR() || gpu_info.IsMali() ||
+                            gpu_info.IsApple();
+
+  DepthWiseConv3x3StrideH2 desc(definition);
+  desc.local_mem_uploads_ =
+      (weights_are_buffer && gpu_info.IsPowerVR() && gpu_info.IsApiOpenCl() &&
+       gpu_info.opencl_info.dedicated_local_memory) ||
+      (gpu_info.IsApple() &&
+       gpu_info.apple_info.IsLocalMemoryPreferredOverGlobal());
+  desc.work_group_size_ = int3(8, 4, 1);
+  desc.code_ = GetKernelDepthWiseConv3x3StrideH2(
+      gpu_info, definition, weights_are_buffer, desc.local_mem_uploads_);
+  auto src_desc = definition.src_tensors[0];
+  desc.AddSrcTensor("src_tensor", src_desc);
+  desc.AddDstTensor("dst_tensor", definition.dst_tensors[0]);
+
+  desc.args_.AddInt("padding_x", -attr.padding.prepended.w);
+  desc.args_.AddInt("padding_y", -attr.padding.prepended.h);
+  desc.args_.AddInt("stride_x", attr.strides.w);
+  desc.args_.AddInt("dilation_x", attr.dilations.w);
+
+  desc.UploadWeightsAndBiases(attr.weights, attr.bias, weights_are_buffer);
+  return desc;
+}
+
+bool IsDepthWiseConv3x3StrideH2Supported(
+    const DepthwiseConvolution2DAttributes& attr) {
+  return attr.weights.shape.o == 1 && attr.weights.shape.h == 3 &&
+         attr.weights.shape.w == 3 && attr.strides.h == 2 &&
+         attr.dilations.h == 1;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_3x3_stride_h2.h

@@ -0,0 +1,153 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_3X3_STRIDE_H2_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_3X3_STRIDE_H2_H_
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/task/buffer_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+
+namespace tflite {
+namespace gpu {
+
+// Depth Wise Convolution for kernel 3x3
+// require:
+//   channels_multiplier = 1;
+//   kernel_size = 3x3;
+//   dilation.y = 1;
+//   stride.y = 2;
+class DepthWiseConv3x3StrideH2 : public GPUOperation {
+ public:
+  DepthWiseConv3x3StrideH2() = default;
+  void GetPossibleKernelWorkGroups(
+      TuningType tuning_type, const GpuInfo& gpu_info,
+      const KernelInfo& kernel_info,
+      std::vector<int3>* work_groups) const override;
+  int3 GetGridSize() const override;
+
+  // Move only
+  DepthWiseConv3x3StrideH2(DepthWiseConv3x3StrideH2&& kernel) = default;
+  DepthWiseConv3x3StrideH2& operator=(DepthWiseConv3x3StrideH2&& kernel) =
+      default;
+  DepthWiseConv3x3StrideH2(const DepthWiseConv3x3StrideH2&) = delete;
+  DepthWiseConv3x3StrideH2& operator=(const DepthWiseConv3x3StrideH2&) = delete;
+
+ private:
+  explicit DepthWiseConv3x3StrideH2(const OperationDef& definition)
+      : GPUOperation(definition) {}
+  friend DepthWiseConv3x3StrideH2 CreateDepthWiseConv3x3StrideH2(
+      const OperationDef& definition,
+      const DepthwiseConvolution2DAttributes& attr, const GpuInfo& gpu_info);
+
+  template <DataType T>
+  void UploadWeightsAndBiases(const tflite::gpu::Tensor<OHWI, T>& weights,
+                              const tflite::gpu::Tensor<Linear, T>& biases,
+                              bool weights_are_buffer);
+  template <DataType S, typename T>
+  void RearrangeWeightsAndBiasesData(
+      const tflite::gpu::Tensor<OHWI, S>& weights,
+      const tflite::gpu::Tensor<Linear, S>& biases, absl::Span<T> dst);
+
+  bool local_mem_uploads_;
+};
+
+template <DataType T>
+void DepthWiseConv3x3StrideH2::UploadWeightsAndBiases(
+    const tflite::gpu::Tensor<OHWI, T>& weights,
+    const tflite::gpu::Tensor<Linear, T>& biases, bool weights_are_buffer) {
+  const int src_depth = DivideRoundUp(weights.shape.i, 4);
+  int texture_width = 10;  // 3x3 kernel + 1 bias
+  int texture_height = src_depth;
+  const int elements_count = texture_width * texture_height;
+  const bool fp32_weights = definition_.precision == CalculationsPrecision::F32;
+  const int float4_size = fp32_weights ? 16 : 8;
+
+  std::vector<uint8_t> data(float4_size * elements_count);
+  if (fp32_weights) {
+    float4* ptr = reinterpret_cast<float4*>(data.data());
+    RearrangeWeightsAndBiasesData(weights, biases,
+                                  absl::MakeSpan(ptr, elements_count));
+  } else {
+    half4* ptr = reinterpret_cast<half4*>(data.data());
+    RearrangeWeightsAndBiasesData(weights, biases,
+                                  absl::MakeSpan(ptr, elements_count));
+  }
+
+  if (weights_are_buffer) {
+    BufferDescriptor desc;
+    desc.element_type = fp32_weights ? DataType::FLOAT32 : DataType::FLOAT16;
+    desc.element_size = 4;
+    desc.size = float4_size * elements_count;
+    desc.data = std::move(data);
+    args_.AddObject("weights",
+                    std::make_unique<BufferDescriptor>(std::move(desc)));
+  } else {
+    TensorDescriptor desc = CreateConstantHWVec4TensorDescriptor(
+        fp32_weights ? DataType::FLOAT32 : DataType::FLOAT16,
+        TensorStorageType::TEXTURE_2D, texture_width, texture_height,
+        data.data());
+    args_.AddObject("weights", std::make_unique<TensorDescriptor>(desc));
+  }
+}
+
+template <DataType S, typename T>
+void DepthWiseConv3x3StrideH2::RearrangeWeightsAndBiasesData(
+    const tflite::gpu::Tensor<OHWI, S>& weights,
+    const tflite::gpu::Tensor<Linear, S>& biases, absl::Span<T> dst) {
+  const int src_depth = DivideRoundUp(weights.shape.i, 4);
+
+  int counter = 0;
+  for (int s = 0; s < src_depth; ++s) {
+    for (int y = 0; y < 3; ++y) {
+      for (int x = 0; x < 3; ++x) {
+        T filter_val;
+        for (int i = 0; i < 4; ++i) {
+          const int s_ch = s * 4 + i;
+          if (s_ch < weights.shape.i) {
+            const int f_index = weights.shape.LinearIndex({0, y, x, s_ch});
+            filter_val[i] = weights.data[f_index];
+          } else {
+            filter_val[i] = 0.0f;
+          }
+        }
+        dst[counter++] = filter_val;
+      }
+    }
+
+    T bias_val;
+    for (int i = 0; i < 4; ++i) {
+      const int dst_ch = s * 4 + i;
+      bias_val[i] = dst_ch >= biases.shape.v ? 0.0f : biases.data[dst_ch];
+    }
+    dst[counter++] = bias_val;
+  }
+}
+
+DepthWiseConv3x3StrideH2 CreateDepthWiseConv3x3StrideH2(
+    const OperationDef& definition,
+    const DepthwiseConvolution2DAttributes& attr, const GpuInfo& gpu_info);
+
+bool IsDepthWiseConv3x3StrideH2Supported(
+    const DepthwiseConvolution2DAttributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_3X3_STRIDE_H2_H_

tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_3x3_stride_h2_test_util.cc

@@ -0,0 +1,68 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_3x3_stride_h2_test_util.h"
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_3x3_stride_h2.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status DepthWiseConv3x3StrideH2SimpleWeightsTest(
+    TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 3, 3, 1);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f};
+
+  DepthwiseConvolution2DAttributes attr;
+  attr.padding.prepended = HW(1, 1);
+  attr.padding.appended = HW(1, 1);
+  attr.strides = HW(2, 2);
+  attr.dilations = HW(1, 1);
+  attr.weights.shape = OHWI(1, 3, 3, 1);
+  attr.weights.data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f};
+  attr.bias.shape = Linear(1);
+  attr.bias.data = {0.0f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      DepthWiseConv3x3StrideH2 operation =
+          CreateDepthWiseConv3x3StrideH2(op_def, attr, env->GetGpuInfo());
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor,
+          std::make_unique<DepthWiseConv3x3StrideH2>(std::move(operation)),
+          BHWC(1, 2, 2, 1), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({8.0f, 12.0f, 20.0f, 24.0f}, dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_3x3_stride_h2_test_util.h

@@ -0,0 +1,31 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_3X3_STRIDE_H2_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_3X3_STRIDE_H2_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status DepthWiseConv3x3StrideH2SimpleWeightsTest(
+    TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_3X3_STRIDE_H2_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_3x3_test_util.cc

@@ -0,0 +1,106 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_3x3_test_util.h"
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_3x3.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status DepthwiseConv3x3SimpleWeightsTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  DepthwiseConvolution2DAttributes attr;
+  attr.padding.prepended = HW(1, 1);
+  attr.padding.appended = HW(1, 1);
+  attr.strides = HW(1, 1);
+  attr.dilations = HW(1, 1);
+  attr.weights.shape = OHWI(1, 3, 3, 2);
+  attr.weights.data = {0.0f, 1.0f, 1.0f, 1.0f, 0.0f, 1.0f, 1.0f, 1.0f, 1.0f,
+                       1.0f, 1.0f, 1.0f, 0.0f, 1.0f, 1.0f, 1.0f, 0.0f, 1.0f};
+  attr.bias.shape = Linear(2);
+  attr.bias.data = {0.0f, 0.0f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      DepthwiseConv3x3 operation =
+          CreateDepthwiseConv3x3(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<DepthwiseConv3x3>(std::move(operation)),
+          BHWC(1, 2, 2, 2), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({6.0f, 16.0f, 8.0f, 16.0f, 10.0f, 16.0f, 12.0f, 16.0f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status DepthwiseConv3x3Test(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  DepthwiseConvolution2DAttributes attr;
+  attr.padding.prepended = HW(1, 1);
+  attr.padding.appended = HW(1, 1);
+  attr.strides = HW(1, 1);
+  attr.dilations = HW(1, 1);
+  attr.weights.shape = OHWI(1, 3, 3, 2);
+  attr.weights.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 0.0f, 1.0f, 2.0f,
+                       3.0f, 4.0f, 5.0f, 0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
+  attr.bias.shape = Linear(2);
+  attr.bias.data = {0.5f, -0.5f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      DepthwiseConv3x3 operation =
+          CreateDepthwiseConv3x3(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<DepthwiseConv3x3>(std::move(operation)),
+          BHWC(1, 2, 2, 2), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear(
+          {40.5f, 67.5f, 16.5f, 35.5f, 40.5f, 67.5f, 16.5f, 35.5f},
+          dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_3x3_test_util.h

@@ -0,0 +1,31 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_3X3_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_3X3_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status DepthwiseConv3x3SimpleWeightsTest(TestExecutionEnvironment* env);
+absl::Status DepthwiseConv3x3Test(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_3X3_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_test_util.cc

@@ -0,0 +1,143 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_test_util.h"
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status DepthwiseConvSimpleWeightsTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  DepthwiseConvolution2DAttributes attr;
+  attr.padding.prepended = HW(1, 0);
+  attr.padding.appended = HW(1, 0);
+  attr.strides = HW(1, 1);
+  attr.dilations = HW(1, 1);
+  attr.weights.shape = OHWI(1, 3, 1, 2);
+  attr.weights.data = {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f};
+  attr.bias.shape = Linear(2);
+  attr.bias.data = {0.0f, 0.0f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      DepthwiseConv operation =
+          CreateDepthwiseConvolution2D(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<DepthwiseConv>(std::move(operation)),
+          BHWC(1, 2, 2, 2), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({4.0f, 6.0f, 8.0f, 10.0f, 4.0f, 6.0f, 8.0f, 10.0f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status DepthwiseConvNoMultiplierTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  DepthwiseConvolution2DAttributes attr;
+  attr.padding.prepended = HW(1, 0);
+  attr.padding.appended = HW(1, 0);
+  attr.strides = HW(1, 1);
+  attr.dilations = HW(1, 1);
+  attr.weights.shape = OHWI(1, 3, 1, 2);
+  attr.weights.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
+  attr.bias.shape = Linear(2);
+  attr.bias.data = {0.5f, -0.5f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      DepthwiseConv operation =
+          CreateDepthwiseConvolution2D(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<DepthwiseConv>(std::move(operation)),
+          BHWC(1, 2, 2, 2), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({16.5f, 27.5f, 28.5f, 43.5f, 8.5f, 15.5f, 12.5f, 23.5f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status DepthwiseConvMultiplier2Test(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  DepthwiseConvolution2DAttributes attr;
+  attr.padding.prepended = HW(1, 0);
+  attr.padding.appended = HW(1, 0);
+  attr.strides = HW(1, 1);
+  attr.dilations = HW(1, 1);
+  attr.weights.shape = OHWI(2, 3, 1, 2);
+  attr.weights.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f,  5.0f,
+                       6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f};
+  attr.bias.shape = Linear(4);
+  attr.bias.data = {0.5f, -0.5f, 1.0f, -1.0f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      DepthwiseConv operation =
+          CreateDepthwiseConvolution2D(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<DepthwiseConv>(std::move(operation)),
+          BHWC(1, 2, 2, 4), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear(
+          {16.5f, 39.5f, 29.0f, 63.0f, 28.5f, 75.5f, 45.0f, 103.0f, 8.5f, 31.5f,
+           17.0f, 51.0f, 12.5f, 59.5f, 25.0f, 83.0f},
+          dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/depthwise_conv_test_util.h

@@ -0,0 +1,32 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status DepthwiseConvSimpleWeightsTest(TestExecutionEnvironment* env);
+absl::Status DepthwiseConvNoMultiplierTest(TestExecutionEnvironment* env);
+absl::Status DepthwiseConvMultiplier2Test(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_DEPTHWISE_CONV_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/elementwise.cc

@@ -0,0 +1,551 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/elementwise.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_replace.h"
+#include "absl/strings/substitute.h"
+
+namespace tflite {
+namespace gpu {
+
+namespace {
+std::string GetOneInputCode(const GpuInfo& gpu_info,
+                            const OperationType& op_type,
+                            CalculationsPrecision precision,
+                            const std::string& input_value,
+                            const std::string& output_value) {
+  const bool use_native_opencl_functions =
+      gpu_info.IsApiOpenCl() && precision != CalculationsPrecision::F32 &&
+      gpu_info.IsAdreno();
+  std::string result;
+  switch (op_type) {
+    case OperationType::ABS:
+      result = "$0 = fabs($1);";
+      break;
+    case OperationType::CEIL:
+      result = "$0 = ceil($1);";
+      break;
+    case OperationType::COS:
+      if (use_native_opencl_functions) {
+        result = "$0 = convert_half4(native_cos(convert_float4($1)));";
+      } else {
+        result = "$0 = cos($1);";
+      }
+      break;
+    case OperationType::COPY:
+      result = "$0 = $1;";
+      break;
+    case OperationType::ELU:
+      if (gpu_info.IsApiOpenCl()) {
+        result = R"(
+$0.x = $1.x < INIT_FLT(0.0f) ? expm1($1.x) : $1.x;
+$0.y = $1.y < INIT_FLT(0.0f) ? expm1($1.y) : $1.y;
+$0.z = $1.z < INIT_FLT(0.0f) ? expm1($1.z) : $1.z;
+$0.w = $1.w < INIT_FLT(0.0f) ? expm1($1.w) : $1.w;)";
+      } else {
+        result = R"(
+$0.x = $1.x < INIT_FLT(0.0f) ? exp($1.x) - INIT_FLT(1.0f) : $1.x;
+$0.y = $1.y < INIT_FLT(0.0f) ? exp($1.y) - INIT_FLT(1.0f) : $1.y;
+$0.z = $1.z < INIT_FLT(0.0f) ? exp($1.z) - INIT_FLT(1.0f) : $1.z;
+$0.w = $1.w < INIT_FLT(0.0f) ? exp($1.w) - INIT_FLT(1.0f) : $1.w;)";
+      }
+      break;
+    case OperationType::EXP:
+      if (use_native_opencl_functions) {
+        result = "$0 = convert_half4(native_exp(convert_float4($1)));";
+      } else {
+        result = "$0 = exp($1);";
+      }
+      break;
+    case OperationType::FLOOR:
+      result = "$0 = floor($1);";
+      break;
+    case OperationType::GELU:
+      // OpenCL has erfc and so it can use the more accurate gelu calculation
+      // as compared to the OpenGL and Vulkan implementations.
+      // gelu(x) = 0.5 * x * erfc(x * -sqrt(0.5))
+      result =
+          "$0 = INIT_FLT4(0.5f) * $1 * erfc($1 * "
+          "INIT_FLT4(-0.70710678118654752440f));";
+      break;
+    case OperationType::HARD_SWISH:
+      result =
+          "$0 = $1 * clamp($1 * INIT_FLT(0.16666667f) + INIT_FLT(0.5f), "
+          "INIT_FLT4(0.0f), "
+          "INIT_FLT4(1.0f));";
+      break;
+    case OperationType::LOG:
+      if (use_native_opencl_functions) {
+        result = "$0 = convert_half4(native_log(convert_float4($1)));";
+      } else {
+        result = "$0 = log($1);";
+      }
+      break;
+    case OperationType::NEG:
+      result = "$0 = -($1);";
+      break;
+    case OperationType::RSQRT:
+      if (use_native_opencl_functions) {
+        result = "$0 = convert_half4(native_rsqrt(convert_float4($1)));";
+      } else {
+        result = "$0 = rsqrt($1);";
+      }
+      break;
+    case OperationType::SIGMOID:
+      if (use_native_opencl_functions) {
+        result =
+            "$0 = convert_half4(native_recip(1.0f + "
+            "native_exp(convert_float4(-$1))));";
+      } else {
+        result = "$0 = INIT_FLT4(1.0f) / (INIT_FLT4(1.0f) + exp(-($1)));";
+      }
+      break;
+    case OperationType::SIGN:
+      result = "$0 = sign($1);";
+      break;
+    case OperationType::SIN:
+      if (use_native_opencl_functions) {
+        result = "$0 = convert_half4(native_sin(convert_float4($1)));";
+      } else {
+        result = "$0 = sin($1);";
+      }
+      break;
+    case OperationType::SQRT:
+      if (use_native_opencl_functions) {
+        result = "$0 = convert_half4(native_sqrt(convert_float4($1)));";
+      } else {
+        result = "$0 = sqrt($1);";
+      }
+      break;
+    case OperationType::SQUARE:
+      result = "$0 = $1 * $1;";
+      break;
+    case OperationType::TANH:
+      if (use_native_opencl_functions) {
+        result =
+            "FLT4 exp_val = convert_half4(native_exp(2.0f * "
+            "convert_float4($1)));\n";
+        result +=
+            "$0 = isinf(exp_val) ? sign($1) : ((exp_val - INIT_FLT4(1.0f)) / "
+            "(exp_val + INIT_FLT4(1.0f)));\n";
+      } else {
+        result = "$0 = tanh($1);";
+      }
+      break;
+    default:
+      return "Unknown operation type;";
+  }
+  return absl::Substitute(result, output_value, input_value);
+}
+
+std::string GetTwoInputCode(const OperationType& op_type,
+                            const std::string& result_var,
+                            const std::string& input0,
+                            const std::string& input1,
+                            bool swap_inputs = false) {
+  std::string result;
+  switch (op_type) {
+    case OperationType::ADD:
+      result += "$0 = $1 + $2;";
+      break;
+    case OperationType::DIV:
+      result += "$0 = $1 / $2;";
+      break;
+    case OperationType::FLOOR_DIV:
+      result = "$0 = floor($1 / $2);";
+      break;
+    case OperationType::FLOOR_MOD:
+      result = "$0 = $1 - floor($1 / $2) * $2;";
+      break;
+    case OperationType::MAXIMUM:
+      result += "$0 = max($1, $2);";
+      break;
+    case OperationType::MINIMUM:
+      result += "$0 = min($1, $2);";
+      break;
+    case OperationType::MUL:
+      result += "$0 = $1 * $2;";
+      break;
+    case OperationType::POW:
+      result += "$0 = pow($1, $2);";
+      break;
+    case OperationType::SQUARED_DIFF:
+      result += "$0 = ($1 - $2) * ($1 - $2);";
+      break;
+    case OperationType::SUB:
+      result += "$0 = $1 - $2;";
+      break;
+    // Comparison operators
+    case OperationType::LESS:
+      result = "$0.x = $1.x < $2.x;\n";
+      result += "$0.y = $1.y < $2.y;\n";
+      result += "$0.z = $1.z < $2.z;\n";
+      result += "$0.w = $1.w < $2.w;";
+      break;
+    case OperationType::LESS_EQUAL:
+      result = "$0.x = $1.x <= $2.x;\n";
+      result += "$0.y = $1.y <= $2.y;\n";
+      result += "$0.z = $1.z <= $2.z;\n";
+      result += "$0.w = $1.w <= $2.w;";
+      break;
+    case OperationType::GREATER:
+      result = "$0.x = $1.x > $2.x;\n";
+      result += "$0.y = $1.y > $2.y;\n";
+      result += "$0.z = $1.z > $2.z;\n";
+      result += "$0.w = $1.w > $2.w;";
+      break;
+    case OperationType::GREATER_EQUAL:
+      result = "$0.x = $1.x >= $2.x;\n";
+      result += "$0.y = $1.y >= $2.y;\n";
+      result += "$0.z = $1.z >= $2.z;\n";
+      result += "$0.w = $1.w >= $2.w;";
+      break;
+    case OperationType::EQUAL:
+      result = "$0.x = $1.x == $2.x;\n";
+      result += "$0.y = $1.y == $2.y;\n";
+      result += "$0.z = $1.z == $2.z;\n";
+      result += "$0.w = $1.w == $2.w;";
+      break;
+    case OperationType::NOT_EQUAL:
+      result = "$0.x = $1.x != $2.x;\n";
+      result += "$0.y = $1.y != $2.y;\n";
+      result += "$0.z = $1.z != $2.z;\n";
+      result += "$0.w = $1.w != $2.w;";
+      break;
+    case OperationType::LOGICAL_AND:
+      result = "$0.x = ($1.x != 0) && ($2.x != 0);\n";
+      result += "$0.y = ($1.y != 0) && ($2.y != 0);\n";
+      result += "$0.z = ($1.z != 0) && ($2.z != 0);\n";
+      result += "$0.w = ($1.w != 0) && ($2.w != 0);";
+      break;
+    default:
+      return "Unknown operation type;";
+  }
+  if (swap_inputs) {
+    return absl::Substitute(result, result_var, input1, input0);
+  } else {
+    return absl::Substitute(result, result_var, input0, input1);
+  }
+}
+
+// Creates simple two input (first input is runtime tensor and second input is
+// scalar argument) operation, for example sub, div, pow, etc.
+template <typename T>
+ElementwiseDescriptor CreateElementwiseOneRuntimeOneScalar(
+    const OperationDef& definition, const OperationType& op_type,
+    T scalar_parameter, bool swap_inputs) {
+  ElementwiseDescriptor op_desc;
+  if (std::is_same<T, int32_t>::value) {
+    op_desc.args.AddInt("scalar", scalar_parameter);
+    op_desc.code =
+        "int4 second_val = CONVERT_TO_INT4(INIT_FLT4(args.scalar));\n";
+    op_desc.code += GetTwoInputCode(op_type, "out_value", "in_value",
+                                    "second_val", swap_inputs);
+    return op_desc;
+  }
+  if (definition.precision == CalculationsPrecision::F32) {
+    op_desc.args.AddFloat("scalar", scalar_parameter);
+  } else {
+    op_desc.args.AddHalf("scalar", half(scalar_parameter));
+  }
+  op_desc.code = "FLT4 second_val = INIT_FLT4(args.scalar);\n";
+  op_desc.code += GetTwoInputCode(op_type, "out_value", "in_value",
+                                  "second_val", swap_inputs);
+  return op_desc;
+}
+
+// Creates simple two input(first input is runtime tensor and second input is
+// constant linear tensor) operation, for example sub, div and etc.
+template <DataType DataTypeT>
+ElementwiseDescriptor CreateElementwiseTwoInput(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const OperationType& op_type,
+    const tflite::gpu::Tensor<Linear, DataTypeT>& constant_tensor,
+    bool swap_inputs) {
+  TensorDescriptor const_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), constant_tensor);
+  ElementwiseDescriptor op_desc;
+  op_desc.args.AddObject("second_tensor", std::make_unique<TensorDescriptor>(
+                                              std::move(const_tensor_desc)));
+  const std::string s_coord = constant_tensor.shape.v == 1 ? "0" : "S_COORD";
+  op_desc.code = absl::StrCat(
+      "args.second_tensor::type second_val = args.second_tensor.Read(", s_coord,
+      ");\n");
+  if (constant_tensor.shape.v == 1) {
+    op_desc.code += "  second_val.y = second_val.x;\n";
+    op_desc.code += "  second_val.z = second_val.x;\n";
+    op_desc.code += "  second_val.w = second_val.x;\n";
+  }
+  op_desc.code += GetTwoInputCode(op_type, "out_value", "in_value",
+                                  "second_val", swap_inputs);
+  return op_desc;
+}
+
+// Creates simple two input(first input is runtime tensor and second input is
+// constant HWC tensor) operation, for example sub, div and etc.
+template <DataType DataTypeT>
+ElementwiseDescriptor CreateElementwiseTwoInput(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const OperationType& op_type,
+    const tflite::gpu::Tensor<HWC, DataTypeT>& constant_tensor,
+    bool swap_inputs) {
+  const BHWC shape = BHWC(1, constant_tensor.shape.h, constant_tensor.shape.w,
+                          constant_tensor.shape.c);
+  TensorDescriptor const_tensor_desc =
+      TensorDescriptor(definition.src_tensors[0].GetDataType(),
+                       definition.src_tensors[0].GetStorageType(), Layout::HWC);
+  auto status = const_tensor_desc.UpdateToSupportedStorageType(gpu_info, shape);
+  const_tensor_desc.UploadData(constant_tensor);
+
+  ElementwiseDescriptor op_desc;
+  op_desc.args.AddObject("second_tensor", std::make_unique<TensorDescriptor>(
+                                              std::move(const_tensor_desc)));
+  const std::string x_coord = shape.w == 1 ? "0" : "X_COORD";
+  const std::string y_coord = shape.h == 1 ? "0" : "Y_COORD";
+  const std::string s_coord = shape.c == 1 ? "0" : "S_COORD";
+  op_desc.code = absl::StrCat(
+      "args.second_tensor::type second_val = args.second_tensor.Read(", x_coord,
+      ", ", y_coord, ", ", s_coord, ");\n");
+  if (shape.c == 1) {
+    op_desc.code += "  second_val.y = second_val.x;\n";
+    op_desc.code += "  second_val.z = second_val.x;\n";
+    op_desc.code += "  second_val.w = second_val.x;\n";
+  }
+  op_desc.code += GetTwoInputCode(op_type, "out_value", "in_value",
+                                  "second_val", swap_inputs);
+
+  return op_desc;
+}
+
+template <DataType DataTypeT, typename T>
+ElementwiseDescriptor CreateElementwiseDesc(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const OperationType& op_type,
+    const ElementwiseAttributesBase<DataTypeT, T>& attr) {
+  const T* scalar = std::get_if<T>(&attr.param);
+  const auto* linear_tensor =
+      std::get_if<tflite::gpu::Tensor<Linear, DataTypeT>>(&attr.param);
+  const auto* hwc_tensor =
+      std::get_if<tflite::gpu::Tensor<HWC, DataTypeT>>(&attr.param);
+
+  if (scalar) {
+    return CreateElementwiseOneRuntimeOneScalar(definition, op_type, *scalar,
+                                                attr.runtime_tensor_is_second);
+  } else if (linear_tensor) {
+    return CreateElementwiseTwoInput(gpu_info, definition, op_type,
+                                     *linear_tensor,
+                                     attr.runtime_tensor_is_second);
+  } else if (hwc_tensor) {
+    return CreateElementwiseTwoInput(gpu_info, definition, op_type, *hwc_tensor,
+                                     attr.runtime_tensor_is_second);
+  } else {
+    return ElementwiseDescriptor();
+  }
+}
+
+}  // namespace
+
+ElementwiseDescriptor CreateElementwiseOneInput(const GpuInfo& gpu_info,
+                                                CalculationsPrecision precision,
+                                                const OperationType& op_type) {
+  ElementwiseDescriptor op_desc;
+  op_desc.code =
+      GetOneInputCode(gpu_info, op_type, precision, "in_value", "out_value");
+  return op_desc;
+}
+
+GPUOperation CreateElementwiseOneInput(const GpuInfo& gpu_info,
+                                       const OperationDef& definition,
+                                       const OperationType& op_type) {
+  return CreateGpuOperation(
+      definition,
+      CreateElementwiseOneInput(gpu_info, definition.precision, op_type));
+}
+
+template <DataType DataTypeT, typename T>
+GPUOperation CreateElementwise(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const OperationType& op_type,
+    const ElementwiseAttributesBase<DataTypeT, T>& attr) {
+  return CreateGpuOperation(
+      definition, CreateElementwiseDesc(gpu_info, definition, op_type, attr));
+}
+
+GPUOperation CreateElementwiseTwoInput(const OperationDef& definition,
+                                       const OperationType& op_type,
+                                       const BHWC& shape) {
+  ElementwiseDescriptor op_desc;
+  op_desc.code =
+      GetTwoInputCode(op_type, "out_value", "in_value", "in2_value", false);
+  return CreateGpuOperation(definition, std::move(op_desc), shape);
+}
+
+namespace {
+std::string GetKernelBodyCode(const TensorDescriptor& dst_desc) {
+  std::string c;
+  c += "MAIN_FUNCTION($$0) {\n";
+  if (dst_desc.HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  c += "  int Y = GLOBAL_ID_1;\n";
+  c += "  int S = GLOBAL_ID_2;\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "S >= args.dst_tensor.Slices()) return; \n";
+  c += "  args.dst_tensor::type result;\n";
+  c += "  $0\n";
+  c += "  args.dst_tensor.Write(result, X, Y, S);\n";
+  c += "} \n";
+  return c;
+}
+std::string GetReadBroadcastedValueCode(const BHWC& src_shape,
+                                        const TensorDescriptor& src_desc,
+                                        const BHWC& dst_shape) {
+  const std::string x_coord = src_shape.w != dst_shape.w ? "0" : "X";
+  const std::string y_coord = src_shape.h != dst_shape.h ? "0" : "Y";
+  const std::string s_coord = src_shape.c != dst_shape.c ? "0" : "S";
+  std::string coords = absl::StrCat(x_coord, ", ", y_coord, ", ", s_coord);
+  if (src_desc.HasAxis(Axis::BATCH)) {
+    const std::string b_coord = src_shape.b != dst_shape.b ? "0" : "B";
+    coords += ", " + b_coord;
+  }
+  std::string read_value_code =
+      absl::StrCat("args.$0::type $1 = args.$0.Read(", coords, ");\n");
+  if (src_shape.c != dst_shape.c) {
+    read_value_code += "  $1.y = $1.x;\n";
+    read_value_code += "  $1.z = $1.x;\n";
+    read_value_code += "  $1.w = $1.x;\n";
+  }
+  return read_value_code;
+}
+}  // namespace
+
+GPUOperation CreateElementwiseOneInputWithBroadcast(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const OperationType& op_type, const BHWC& input_shape,
+    const BHWC& output_shape) {
+  GPUOperation op(definition);
+  op.AddSrcTensor("src_tensor", definition.src_tensors[0]);
+  op.AddDstTensor("dst_tensor", definition.dst_tensors[0]);
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+  std::string c;
+  c += "  " + absl::Substitute(
+                  GetReadBroadcastedValueCode(
+                      input_shape, definition.src_tensors[0], output_shape),
+                  "src_tensor", "first_value");
+  c += "  " + GetOneInputCode(gpu_info, op_type, definition.precision,
+                              "first_value", "result");
+  op.code_ = absl::Substitute(GetKernelBodyCode(definition.dst_tensors[0]), c);
+  return op;
+}
+
+template <DataType DataTypeT, typename T>
+GPUOperation CreateElementwiseWithBroadcast(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const OperationType& op_type,
+    const ElementwiseAttributesBase<DataTypeT, T>& attr,
+    const BHWC& input_shape, const BHWC& output_shape) {
+  ElementwiseDescriptor op_desc =
+      CreateElementwiseDesc(gpu_info, definition, op_type, attr);
+
+  GPUOperation op(definition);
+  op.args_ = std::move(op_desc.args);
+  op.AddSrcTensor("src_tensor", definition.src_tensors[0]);
+  op.AddDstTensor("dst_tensor", definition.dst_tensors[0]);
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+  std::string c;
+  c += "  " + absl::Substitute(
+                  GetReadBroadcastedValueCode(
+                      input_shape, definition.src_tensors[0], output_shape),
+                  "src_tensor", "first_value");
+  c += "  " + absl::StrReplaceAll(op_desc.code, {{"in_value", "first_value"},
+                                                 {"out_value", "result"},
+                                                 {"X_COORD", "X"},
+                                                 {"Y_COORD", "Y"},
+                                                 {"S_COORD", "S"},
+                                                 {"B_COORD", "B"}});
+  op.code_ = absl::Substitute(GetKernelBodyCode(definition.dst_tensors[0]), c);
+  return op;
+}
+
+GPUOperation CreateElementwiseTwoInputWithBroadcast(
+    const OperationDef& definition, const OperationType& op_type,
+    const BHWC& first_input_shape, const BHWC& second_input_shape,
+    const BHWC& output_shape) {
+  GPUOperation op(definition);
+  op.AddSrcTensor("src0_tensor", definition.src_tensors[0]);
+  op.AddSrcTensor("src1_tensor", definition.src_tensors[1]);
+  op.AddDstTensor("dst_tensor", definition.dst_tensors[0]);
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+  std::string c;
+  c += "  " + absl::Substitute(GetReadBroadcastedValueCode(
+                                   first_input_shape, definition.src_tensors[0],
+                                   output_shape),
+                               "src0_tensor", "first_value");
+  c += "  " + absl::Substitute(GetReadBroadcastedValueCode(
+                                   second_input_shape,
+                                   definition.src_tensors[1], output_shape),
+                               "src1_tensor", "second_value");
+  c += "  " +
+       GetTwoInputCode(op_type, "result", "first_value", "second_value", false);
+  op.code_ = absl::Substitute(GetKernelBodyCode(definition.dst_tensors[0]), c);
+  return op;
+}
+
+template GPUOperation CreateElementwise(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const OperationType& op_type,
+    const ElementwiseAttributesBase<DataType::BOOL, bool>& attr);
+
+template GPUOperation CreateElementwise(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const OperationType& op_type,
+    const ElementwiseAttributesBase<DataType::FLOAT32, float>& attr);
+
+template GPUOperation CreateElementwise(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const OperationType& op_type,
+    const ElementwiseAttributesBase<DataType::INT32, int32_t>& attr);
+
+template GPUOperation CreateElementwiseWithBroadcast(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const OperationType& op_type,
+    const ElementwiseAttributesBase<DataType::BOOL, bool>& attr,
+    const BHWC& input_shape, const BHWC& output_shape);
+
+template GPUOperation CreateElementwiseWithBroadcast(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const OperationType& op_type,
+    const ElementwiseAttributesBase<DataType::FLOAT32, float>& attr,
+    const BHWC& input_shape, const BHWC& output_shape);
+
+template GPUOperation CreateElementwiseWithBroadcast(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const OperationType& op_type,
+    const ElementwiseAttributesBase<DataType::INT32, int>& attr,
+    const BHWC& input_shape, const BHWC& output_shape);
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/elementwise.h

@@ -0,0 +1,85 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ELEMENTWISE_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ELEMENTWISE_H_
+
+#include <string>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/gpu_info.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/precision.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+
+namespace tflite {
+namespace gpu {
+
+// Creates simple one input operation without any parameters, for example
+// log, sin, cos, etc.
+ElementwiseDescriptor CreateElementwiseOneInput(const GpuInfo& gpu_info,
+                                                CalculationsPrecision precision,
+                                                const OperationType& op_type);
+
+GPUOperation CreateElementwiseOneInput(const GpuInfo& gpu_info,
+                                       const OperationDef& definition,
+                                       const OperationType& op_type);
+
+// Creates simple one input operation without any parameters, for example
+// log, sin, cos, etc.
+// Can broadcast input.
+GPUOperation CreateElementwiseOneInputWithBroadcast(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const OperationType& op_type, const BHWC& input_shape,
+    const BHWC& output_shape);
+
+// Creates simple two input(first input is runtime tensor and second input is
+// constant or linear/hwc tensor) operation, for example sub, div and etc.
+template <DataType DataTypeT, typename T>
+GPUOperation CreateElementwise(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const OperationType& op_type,
+    const ElementwiseAttributesBase<DataTypeT, T>& attr);
+
+// Creates simple two input(first input is runtime tensor and second input is
+// constant or linear/hwc tensor) operation, for example sub, div and etc.
+// Can broadcast input.
+template <DataType DataTypeT, typename T>
+GPUOperation CreateElementwiseWithBroadcast(
+    const GpuInfo& gpu_info, const OperationDef& definition,
+    const OperationType& op_type,
+    const ElementwiseAttributesBase<DataTypeT, T>& attr,
+    const BHWC& input_shape, const BHWC& output_shape);
+
+// Creates simple two input(2 runtime tensors) operation, for example
+// sub, div and etc.
+GPUOperation CreateElementwiseTwoInput(const OperationDef& definition,
+                                       const OperationType& op_type,
+                                       const BHWC& shape);
+
+// Creates simple two input(2 runtime tensors) operation, for example
+// sub, div and etc.
+// Can broadcast first and second input simultaneously.
+GPUOperation CreateElementwiseTwoInputWithBroadcast(
+    const OperationDef& definition, const OperationType& op_type,
+    const BHWC& first_input_shape, const BHWC& second_input_shape,
+    const BHWC& output_shape);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ELEMENTWISE_H_

tensorflow/lite/delegates/gpu/common/tasks/elementwise_test_util.h

@@ -0,0 +1,76 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ELEMENTWISE_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ELEMENTWISE_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status AbsTest(TestExecutionEnvironment* env);
+absl::Status CosTest(TestExecutionEnvironment* env);
+absl::Status CopyTest(TestExecutionEnvironment* env);
+absl::Status EluTest(TestExecutionEnvironment* env);
+absl::Status ExpTest(TestExecutionEnvironment* env);
+absl::Status FloorTest(TestExecutionEnvironment* env);
+absl::Status FloorDivTest(TestExecutionEnvironment* env);
+absl::Status FloorModTest(TestExecutionEnvironment* env);
+absl::Status GeluTest(TestExecutionEnvironment* env);
+absl::Status HardSwishTest(TestExecutionEnvironment* env);
+absl::Status LogTest(TestExecutionEnvironment* env);
+absl::Status NegTest(TestExecutionEnvironment* env);
+absl::Status RsqrtTest(TestExecutionEnvironment* env);
+absl::Status SigmoidTest(TestExecutionEnvironment* env);
+absl::Status SinTest(TestExecutionEnvironment* env);
+absl::Status SqrtTest(TestExecutionEnvironment* env);
+absl::Status SquareTest(TestExecutionEnvironment* env);
+absl::Status TanhTest(TestExecutionEnvironment* env);
+absl::Status SubTest(TestExecutionEnvironment* env);
+absl::Status SquaredDiffTest(TestExecutionEnvironment* env);
+absl::Status DivTest(TestExecutionEnvironment* env);
+absl::Status PowTest(TestExecutionEnvironment* env);
+absl::Status AddTest(TestExecutionEnvironment* env);
+absl::Status MaximumTest(TestExecutionEnvironment* env);
+absl::Status MaximumWithScalarTest(TestExecutionEnvironment* env);
+absl::Status MaximumWithConstantLinearTensorTest(TestExecutionEnvironment* env);
+absl::Status MaximumWithConstantHWCTensorTest(TestExecutionEnvironment* env);
+absl::Status MaximumWithConstantHWCTensorBroadcastChannelsTest(
+    TestExecutionEnvironment* env);
+absl::Status MinimumTest(TestExecutionEnvironment* env);
+absl::Status MinimumWithScalarTest(TestExecutionEnvironment* env);
+absl::Status MulTest(TestExecutionEnvironment* env);
+absl::Status MulBroadcastHWTest(TestExecutionEnvironment* env);
+absl::Status MulBroadcastChannelsTest(TestExecutionEnvironment* env);
+absl::Status SubWithScalarAtFirstPositionTest(TestExecutionEnvironment* env);
+absl::Status LessTest(TestExecutionEnvironment* env);
+absl::Status LessEqualTest(TestExecutionEnvironment* env);
+absl::Status GreaterTest(TestExecutionEnvironment* env);
+absl::Status GreaterEqualTest(TestExecutionEnvironment* env);
+absl::Status EqualTest(TestExecutionEnvironment* env);
+absl::Status NotEqualTest(TestExecutionEnvironment* env);
+absl::Status CosBroadcastTest(TestExecutionEnvironment* env);
+absl::Status MaximumScalarBroadcastInputTest(TestExecutionEnvironment* env);
+absl::Status MulLinearBroadcastInputTest(TestExecutionEnvironment* env);
+absl::Status MulBroadcastBothInputsTest(TestExecutionEnvironment* env);
+absl::Status LogicalAndTest(TestExecutionEnvironment* env);
+absl::Status LogicalAndWithConstantTest(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ELEMENTWISE_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/fully_connected.cc

@@ -0,0 +1,251 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/fully_connected.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/gpu_info.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/precision.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tensor_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/tensor.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+#include "tensorflow/lite/delegates/gpu/common/util.h"
+
+namespace tflite {
+namespace gpu {
+
+namespace {
+bool UseBufferForWeights(const GpuInfo& gpu_info) {
+  return gpu_info.IsAdreno() || gpu_info.IsAMD() || gpu_info.IsMali() ||
+         gpu_info.IsApple() ||
+         (gpu_info.IsIntel() && gpu_info.IsApiOpenCl() &&
+          gpu_info.opencl_info.IsCLVK());
+}
+
+void RearrangeFCWeightsToOIO4I4(
+    const tflite::gpu::Tensor<OHWI, DataType::INT8>& weights, uint8_t* dst) {
+  const int src_depth = DivideRoundUp(weights.shape.i, 4);
+  const int dst_depth = DivideRoundUp(weights.shape.o, 4);
+
+  int counter = 0;
+  for (int d = 0; d < dst_depth; ++d) {
+    for (int s = 0; s < src_depth; ++s) {
+      for (int i = 0; i < 4; ++i) {
+        const int src_ch = s * 4 + i;
+        for (int j = 0; j < 4; ++j) {
+          const int dst_ch = d * 4 + j;
+          if (src_ch < weights.shape.i && dst_ch < weights.shape.o) {
+            int t =
+                127 +
+                weights.data[weights.shape.LinearIndex({dst_ch, 0, 0, src_ch})];
+            if (t < 0) {
+              t = 0;
+            }
+            dst[counter++] = t;
+          } else {
+            dst[counter++] = 127;
+          }
+        }
+      }
+    }
+  }
+}
+
+}  // namespace
+
+FullyConnected::FullyConnected(const OperationDef& definition,
+                               const GpuInfo& gpu_info)
+    : GPUOperation(definition) {
+  if (gpu_info.IsAdreno()) {
+    if (gpu_info.adreno_info.IsAdreno3xx()) {
+      work_group_size_ = int3(16, 4, 1);
+    } else if (gpu_info.adreno_info.IsAdreno4xx()) {
+      work_group_size_ = int3(32, 4, 1);
+    } else {
+      work_group_size_ = int3(32, 4, 1);
+    }
+  } else if (gpu_info.IsIntel() || gpu_info.IsNvidia() ||
+             gpu_info.IsPowerVR() || gpu_info.IsApple()) {
+    work_group_size_ = int3(8, 4, 1);
+  } else {
+    work_group_size_ = int3(16, 4, 1);
+  }
+}
+
+FullyConnected::FullyConnected(FullyConnected&& kernel)
+    : GPUOperation(std::move(kernel)) {}
+
+FullyConnected& FullyConnected::operator=(FullyConnected&& kernel) {
+  if (this != &kernel) {
+    GPUOperation::operator=(std::move(kernel));
+  }
+  return *this;
+}
+
+// We split vec vec dot (every thread do vec vec dot product in basic
+// vec mat mult) on 4 parts to create more threads
+// tid.y thread process every 4-th element in vec vec dot
+// Good results for ~1024 x 1024 sizes, for other can be written more
+// optimized shaders
+
+std::string FullyConnected::GetFullyConnectedKernelCode(
+    const OperationDef& op_def, const GpuInfo& gpu_info,
+    bool weights_are_buffer, bool quantized) {
+  const int wg_total_size = work_group_size_.x * work_group_size_.y;
+  const std::string barrier =
+      wg_total_size == 32 && gpu_info.IsWaveSizeEqualTo32()
+          ? "SIMD_LOCAL_MEM_BARRIER"
+          : "LOCAL_MEM_BARRIER";
+  AddSrcTensor("src_tensor", op_def.src_tensors[0]);
+  AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
+
+  std::string c;
+
+  c += "#define WG_X " + std::to_string(work_group_size_.x) + "\n";
+  c += "#define WG_Y " + std::to_string(work_group_size_.y) + "\n";
+
+  c += R"(MAIN_FUNCTION($0) {
+  int gid = GLOBAL_ID_0;
+  int2 tid = INIT_INT2v2(LOCAL_ID_0, LOCAL_ID_1);
+  ACCUM_FLT4 s = INIT_ACCUM_FLT4(0.0f);
+  if (gid < args.dst_tensor.Slices()) {
+    for (int c = tid.y; c < args.src_tensor.Slices(); c += WG_Y) {
+      FLT4 v = args.src_tensor.Read(0, 0, c);
+)";
+  if (weights_are_buffer) {
+    c += R"(int weights_index = (c * args.dst_tensor.Slices() + gid) * 4;
+      FLT4 partial = v.x * args.weights.Read(weights_index + 0);
+      partial += v.y * args.weights.Read(weights_index + 1);
+      partial += v.z * args.weights.Read(weights_index + 2);
+      partial += v.w * args.weights.Read(weights_index + 3);
+      s += TO_ACCUM_TYPE(partial);
+)";
+  } else {
+    const std::string read_as_type =
+        op_def.precision == CalculationsPrecision::F32 ? "float" : "half";
+    c += "      FLT4 w0 = args.weights.Read<" + read_as_type +
+         ">(c * 4 + 0, gid);\n";
+    c += "      FLT4 w1 = args.weights.Read<" + read_as_type +
+         ">(c * 4 + 1, gid);\n";
+    c += "      FLT4 w2 = args.weights.Read<" + read_as_type +
+         ">(c * 4 + 2, gid);\n";
+    c += "      FLT4 w3 = args.weights.Read<" + read_as_type +
+         ">(c * 4 + 3, gid);\n";
+    if (quantized) {
+      c += R"(w0 = w0 * args.q0 + args.q1;
+      w1 = w1 * args.q0 + args.q1;
+      w2 = w2 * args.q0 + args.q1;
+      w3 = w3 * args.q0 + args.q1;
+)";
+    }
+    c += R"(FLT4 partial = v.x * w0;
+      partial += v.y * w1;
+      partial += v.z * w2;
+      partial += v.w * w3;
+      s += TO_ACCUM_TYPE(partial);
+)";
+  }
+  c += R"(    }
+  }
+  __local ACCUM_FLT4 temp[WG_X][WG_Y];
+  temp[tid.x][tid.y] = s;
+)";
+  c += "  " + barrier + ";\n";
+  c += R"(
+  if (gid >= args.dst_tensor.Slices()) {
+    return;
+  }
+  if (tid.y == 0) {
+)";
+  for (int i = 1; i < work_group_size_.y; ++i) {
+    c += "    s += temp[tid.x][" + std::to_string(i) + "];\n";
+  }
+  c += R"(    FLT4 r0 = TO_FLT4(s) + args.biases.Read(gid);
+    args.dst_tensor.Write(r0, 0, 0, gid);
+  }
+})";
+
+  return c;
+}
+
+int3 FullyConnected::GetGridSize() const {
+  return int3(dst_[0]->Slices(), 1, 1);
+}
+
+void FullyConnected::UploadQuantizedWeights(
+    const tflite::gpu::Tensor<OHWI, DataType::INT8>& weights, float scale,
+    float zero_point) {
+  const int src_depth = DivideRoundUp(weights.shape.i, 4);
+  const int dst_depth = DivideRoundUp(weights.shape.o, 4);
+
+  std::vector<uint8_t> data(static_cast<size_t>(src_depth) * 4 * dst_depth * 4);
+  RearrangeFCWeightsToOIO4I4(weights, data.data());
+  TensorDescriptor desc = CreateConstantHWVec4TensorDescriptor(
+      DataType::UINT8, TensorStorageType::TEXTURE_2D, src_depth * 4, dst_depth,
+      data.data());
+
+  if (definition_.precision == CalculationsPrecision::F32) {
+    args_.AddFloat("q0", scale);
+    args_.AddFloat("q1", -scale * (127.0 + zero_point));
+  } else {
+    args_.AddHalf("q0", half(scale));
+    args_.AddHalf("q1", half(-scale * (127.0 + zero_point)));
+  }
+  args_.AddObject("weights",
+                  std::make_unique<TensorDescriptor>(std::move(desc)));
+}
+
+FullyConnected CreateFullyConnected(const GpuInfo& gpu_info,
+                                    const OperationDef& definition,
+                                    const FullyConnectedAttributes& attr) {
+  FullyConnected result(definition, gpu_info);
+  result.UploadWeights(attr.weights, UseBufferForWeights(gpu_info));
+  result.code_ = result.GetFullyConnectedKernelCode(
+      definition, gpu_info, UseBufferForWeights(gpu_info), false);
+
+  TensorDescriptor bias_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), attr.bias);
+  result.args_.AddObject("biases", std::make_unique<TensorDescriptor>(
+                                       std::move(bias_tensor_desc)));
+
+  return result;
+}
+
+FullyConnected CreateFullyConnected(const GpuInfo& gpu_info,
+                                    const OperationDef& definition,
+                                    const FullyConnectedInt8Attributes& attr) {
+  FullyConnected result(definition, gpu_info);
+  result.UploadQuantizedWeights(attr.weights, attr.scale, attr.zero_point);
+  result.code_ =
+      result.GetFullyConnectedKernelCode(definition, gpu_info, false, true);
+
+  TensorDescriptor bias_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), attr.bias);
+  result.args_.AddObject("biases", std::make_unique<TensorDescriptor>(
+                                       std::move(bias_tensor_desc)));
+
+  return result;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/fully_connected.h

@@ -0,0 +1,217 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_FULLY_CONNECTED_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_FULLY_CONNECTED_H_
+
+#include <stdint.h>
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/memory/memory.h"
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/gpu_info.h"
+#include "tensorflow/lite/delegates/gpu/common/kernel_info.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/precision.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/task/buffer_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tensor_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tuning_type.h"
+#include "tensorflow/lite/delegates/gpu/common/tensor.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+#include "tensorflow/lite/delegates/gpu/common/util.h"
+
+namespace tflite {
+namespace gpu {
+
+template <DataType T, typename S>
+void RearrangeFCWeightsToIOO4I4(const tflite::gpu::Tensor<OHWI, T>& weights,
+                                S* dst) {
+  const int src_channels = weights.shape.i;
+  const int padded_src_channels = AlignByN(src_channels, 4);
+  const int dst_channels = weights.shape.o;
+  const int padded_dst_channels = AlignByN(dst_channels, 4);
+
+  // Change the travelsal order of the weight matrix in the following way:
+  // The matrix is segmented to blocks of 4x4. If (any) dimension of the matrix
+  // size is not divisible by 4, then pad with zeros. Each block is stored
+  // contigously. The 16 elements within a block are ordered as 4 elements of
+  // the first column, 4 elems of the second, etc. Blocks then traversed as
+  // columns first, rows last. As an example, an 8x8 matrix would be traversed
+  // as below.
+  //
+  //  |  0  4  8 12 32 36 40 44 |
+  //  |  1  5  9 13 33 37 41 45 |
+  //  |  2  6 10 14 34 38 42 46 |
+  //  |  3  7 11 15 35 39 43 47 |
+  //  | 16 20 24 28 48 52 56 60 |
+  //  | 17 21 25 29 49 53 57 61 |
+  //  | 18 22 26 30 50 54 58 62 |
+  //  | 19 23 27 31 51 55 59 63 |
+  //
+  // The benefit of doing this is that reading contigous 16 elements gives a 4x4
+  // block of the matrix, where the first 4 elements is the first row of the
+  // block, second 4 elements is the second row of the block, etc. Subsequent
+  // blocks contain elements of the same 4 columns.
+
+  for (int block_y = 0; 4 * block_y < padded_dst_channels; block_y++) {
+    for (int y_in_block = 0; y_in_block < 4; y_in_block++) {
+      for (int block_x = 0; 4 * block_x < padded_src_channels; block_x++) {
+        for (int x_in_block = 0; x_in_block < 4; x_in_block++) {
+          int y = 4 * block_y + y_in_block;
+          int x = 4 * block_x + x_in_block;
+          // Consider destination as an array with extents
+          // [padded_src_channels/4][padded_dst_channels/4][4][4]
+          int dst_index = block_x * padded_dst_channels * 4 + block_y * 16 +
+                          x_in_block * 4 + y_in_block;
+          if (x < src_channels && y < dst_channels) {
+            dst[dst_index] = weights.data[src_channels * y + x];
+          } else {
+            dst[dst_index] = 0.0f;
+          }
+        }
+      }
+    }
+  }
+}
+
+template <DataType T, typename S>
+void RearrangeFCWeightsToOIO4I4(const tflite::gpu::Tensor<OHWI, T>& weights,
+                                S* dst) {
+  const int src_channels = weights.shape.i;
+  const int src_depth = DivideRoundUp(src_channels, 4);
+  const int dst_channels = weights.shape.o;
+  const int dst_depth = DivideRoundUp(dst_channels, 4);
+
+  int counter = 0;
+  for (int d = 0; d < dst_depth; ++d) {
+    for (int s = 0; s < src_depth; ++s) {
+      for (int i = 0; i < 4; ++i) {
+        const int src_ch = s * 4 + i;
+        for (int j = 0; j < 4; ++j) {
+          const int dst_ch = d * 4 + j;
+          if (src_ch < src_channels && dst_ch < dst_channels) {
+            dst[counter++] = weights.data[dst_ch * src_channels + src_ch];
+          } else {
+            dst[counter++] = 0.0f;
+          }
+        }
+      }
+    }
+  }
+}
+
+class FullyConnected : public GPUOperation {
+ public:
+  FullyConnected() = default;
+  void GetPossibleKernelWorkGroups(
+      TuningType tuning_type, const GpuInfo& gpu_info,
+      const KernelInfo& kernel_info,
+      std::vector<int3>* work_groups) const override {
+    work_groups->push_back(work_group_size_);
+  }
+  int3 GetGridSize() const override;
+
+  // Move only
+  FullyConnected(FullyConnected&& kernel);
+  FullyConnected& operator=(FullyConnected&& kernel);
+  FullyConnected(const FullyConnected&) = delete;
+  FullyConnected& operator=(const FullyConnected&) = delete;
+
+ private:
+  FullyConnected(const OperationDef& definition, const GpuInfo& gpu_info);
+  friend FullyConnected CreateFullyConnected(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const FullyConnectedAttributes& attr);
+  friend FullyConnected CreateFullyConnected(
+      const GpuInfo& gpu_info, const OperationDef& definition,
+      const FullyConnectedInt8Attributes& attr);
+
+  void UploadQuantizedWeights(
+      const tflite::gpu::Tensor<OHWI, DataType::INT8>& weights, float scale,
+      float zero_point);
+  template <DataType T>
+  void UploadWeights(const tflite::gpu::Tensor<OHWI, T>& weights,
+                     bool weights_are_buffer);
+
+  std::string GetFullyConnectedKernelCode(const OperationDef& op_def,
+                                          const GpuInfo& gpu_info,
+                                          bool weights_are_buffer,
+                                          bool quantized);
+};
+
+template <DataType T>
+void FullyConnected::UploadWeights(const tflite::gpu::Tensor<OHWI, T>& weights,
+                                   bool weights_are_buffer) {
+  const int src_depth = DivideRoundUp(weights.shape.i, 4);
+  const int dst_depth = DivideRoundUp(weights.shape.o, 4);
+
+  const size_t elements_count = static_cast<size_t>(src_depth) * dst_depth * 4;
+  const bool f32_weights = definition_.precision == CalculationsPrecision::F32;
+
+  const int float4_size = f32_weights ? 16 : 8;
+
+  if (weights_are_buffer) {
+    BufferDescriptor desc;
+    desc.element_type = f32_weights ? DataType::FLOAT32 : DataType::FLOAT16;
+    desc.element_size = 4;
+    desc.size = float4_size * elements_count;
+    desc.data.resize(desc.size);
+
+    if (f32_weights) {
+      float* ptr = reinterpret_cast<float*>(desc.data.data());
+      RearrangeFCWeightsToIOO4I4(weights, ptr);
+    } else {
+      half* ptr = reinterpret_cast<half*>(desc.data.data());
+      RearrangeFCWeightsToIOO4I4(weights, ptr);
+    }
+
+    args_.AddObject("weights",
+                    std::make_unique<BufferDescriptor>(std::move(desc)));
+  } else {
+    std::vector<uint8_t> data(float4_size * elements_count);
+    if (f32_weights) {
+      float* ptr = reinterpret_cast<float*>(data.data());
+      RearrangeFCWeightsToOIO4I4(weights, ptr);
+    } else {
+      half* ptr = reinterpret_cast<half*>(data.data());
+      RearrangeFCWeightsToOIO4I4(weights, ptr);
+    }
+
+    TensorDescriptor desc = CreateConstantHWVec4TensorDescriptor(
+        f32_weights ? DataType::FLOAT32 : DataType::FLOAT16,
+        TensorStorageType::TEXTURE_2D, src_depth * 4, dst_depth, data.data());
+    args_.AddObject("weights",
+                    std::make_unique<TensorDescriptor>(std::move(desc)));
+  }
+}
+
+FullyConnected CreateFullyConnected(const GpuInfo& gpu_info,
+                                    const OperationDef& definition,
+                                    const FullyConnectedAttributes& attr);
+
+FullyConnected CreateFullyConnected(const GpuInfo& gpu_info,
+                                    const OperationDef& definition,
+                                    const FullyConnectedInt8Attributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_FULLY_CONNECTED_H_

tensorflow/lite/delegates/gpu/common/tasks/fully_connected_test_util.cc

@@ -0,0 +1,205 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/fully_connected_test_util.h"
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/precision.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/fully_connected.h"
+#include "tensorflow/lite/delegates/gpu/common/tensor.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status FullyConnectedTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 1, 1, 4);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f};
+
+  FullyConnectedAttributes attr;
+  attr.weights.shape = OHWI(2, 1, 1, 4);
+  attr.weights.data = {0.0f, 1.0f, 2.0f, 3.0f,  //
+                       4.0f, 5.0f, 6.0f, 7.0f};
+  attr.bias.shape = Linear(2);
+  attr.bias.data = {0.5f, -0.5f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      FullyConnected operation =
+          CreateFullyConnected(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<FullyConnected>(std::move(operation)),
+          BHWC(1, 1, 1, 2), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear({14.5f, 37.5f}, dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status FullyConnectedLargeTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 1, 1, 8);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
+
+  FullyConnectedAttributes attr;
+  attr.weights.shape = OHWI(12, 1, 1, 8);
+  attr.weights.data = {
+      0.0f,  1.0f,  2.0f,  3.0f,  4.0f,  5.0f,  6.0f,  7.0f,   //
+      8.0f,  9.0f,  10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f,  //
+      16.0f, 17.0f, 18.0f, 19.0f, 20.0f, 21.0f, 22.0f, 23.0f,  //
+      24.0f, 25.0f, 26.0f, 27.0f, 28.0f, 29.0f, 30.0f, 31.0f,  //
+      32.0f, 33.0f, 34.0f, 35.0f, 36.0f, 37.0f, 38.0f, 39.0f,  //
+      40.0f, 41.0f, 42.0f, 43.0f, 44.0f, 45.0f, 46.0f, 47.0f,  //
+      48.0f, 49.0f, 50.0f, 51.0f, 52.0f, 53.0f, 54.0f, 55.0f,  //
+      56.0f, 57.0f, 58.0f, 59.0f, 60.0f, 61.0f, 62.0f, 63.0f,  //
+      64.0f, 65.0f, 66.0f, 67.0f, 68.0f, 69.0f, 70.0f, 71.0f,  //
+      72.0f, 73.0f, 74.0f, 75.0f, 76.0f, 77.0f, 78.0f, 79.0f,  //
+      80.0f, 81.0f, 82.0f, 83.0f, 84.0f, 85.0f, 86.0f, 87.0f,  //
+      88.0f, 89.0f, 90.0f, 91.0f, 92.0f, 93.0f, 94.0f, 95.0f,  //
+  };
+  attr.bias.shape = Linear(12);
+  attr.bias.data = {-0.6f, -0.5f, -0.4f, -0.3f, -0.2f, -0.1f,
+                    0.1f,  0.2f,  0.3f,  0.4f,  0.5f,  0.6f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 0.0f : 1.0f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      FullyConnected operation =
+          CreateFullyConnected(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<FullyConnected>(std::move(operation)),
+          BHWC(1, 1, 1, 12), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({139.4f, 363.5f, 587.6f, 811.7f, 1035.8f, 1259.9f,
+                         1484.1f, 1708.2f, 1932.3f, 2156.4f, 2380.5f, 2604.6f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status FullyConnectedExtraLargeTest(TestExecutionEnvironment* env) {
+  static const int kInputSize = 1024;
+  static const int kOutputSize = 1024;
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 1, 1, kInputSize);
+  src_tensor.data.assign(kInputSize, 1.1f);
+
+  FullyConnectedAttributes attr;
+  attr.weights.shape = OHWI(1024, 1, 1, kInputSize);
+  attr.weights.data.assign(kOutputSize * kInputSize, 2.2f);
+  attr.bias.shape = Linear(kOutputSize);
+  attr.bias.data.assign(kOutputSize, 3.3f);
+
+  std::vector<float> expected(kOutputSize, 2481.38f);
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      float eps;
+      switch (precision) {
+        case CalculationsPrecision::F32:
+          eps = 2.45e-3f;
+          break;
+        case CalculationsPrecision::F32_F16:
+          eps = 1.38f;
+          break;
+        case CalculationsPrecision::F16:
+          eps = 39.0f;
+          break;
+      }
+      if (precision == CalculationsPrecision::F32_F16 &&
+          env->GetGpuInfo().IsApiMetal() && env->GetGpuInfo().IsIntel()) {
+        eps = 3.5f;
+      }
+      if (precision == CalculationsPrecision::F32_F16 &&
+          env->GetGpuInfo().IsGlsl()) {
+        eps = 3.5f;
+      }
+      if (!env->GetGpuInfo().IsRoundToNearestSupported()) {
+        eps *= 4.0f;
+      }
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      FullyConnected operation =
+          CreateFullyConnected(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<FullyConnected>(std::move(operation)),
+          BHWC(1, 1, 1, kOutputSize), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear(expected, dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status FullyConnectedInt8Test(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 1, 1, 4);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f};
+
+  FullyConnectedInt8Attributes attr;
+  attr.weights.shape = OHWI(2, 1, 1, 4);
+  attr.weights.data = {2,  4,  6,   8,  //
+                       10, 12, -14, 16};
+  attr.bias.shape = Linear(2);
+  attr.bias.data = {0.5f, -0.5f};
+  attr.scale = 0.5f;
+  attr.zero_point = 0;
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      FullyConnected operation =
+          CreateFullyConnected(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<FullyConnected>(std::move(operation)),
+          BHWC(1, 1, 1, 2), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear({20.5f, 15.5f}, dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/fully_connected_test_util.h

@@ -0,0 +1,33 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_FULLY_CONNECTED_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_FULLY_CONNECTED_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status FullyConnectedTest(TestExecutionEnvironment* env);
+absl::Status FullyConnectedLargeTest(TestExecutionEnvironment* env);
+absl::Status FullyConnectedExtraLargeTest(TestExecutionEnvironment* env);
+absl::Status FullyConnectedInt8Test(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_FULLY_CONNECTED_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/gather.cc

@@ -0,0 +1,112 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/gather.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/gpu_info.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tensor_desc.h"
+
+namespace tflite {
+namespace gpu {
+
+namespace {
+std::string GetGatherCode(const OperationDef& op_def, GatherAttributes attr) {
+  std::string c;
+  c += "MAIN_FUNCTION($0) {\n";
+  if (op_def.IsBatchSupported()) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  c += "  int Y = GLOBAL_ID_1;\n";
+  c += "  int S = GLOBAL_ID_2;\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "S >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
+  c += "  int idx;\n";
+  c += "  args.src_tensor::type result;\n";
+  switch (attr.axis) {
+    case Axis::BATCH:
+      c += "  idx = args.indices.Read<int>(0, 0, 0, B).x;\n";
+      c += "  result = args.src_tensor.Read(X, Y, "
+           "S, idx);\n";
+      break;
+    case Axis::HEIGHT:
+      c += "  idx = args.indices.Read<int>(0, 0, 0, Y).x;\n";
+      c += "  result = args.src_tensor.Read(X, idx, "
+           "S, B);\n";
+      break;
+    case Axis::WIDTH:
+      c += "  idx = args.indices.Read<int>(0, 0, 0, X).x;\n";
+      c += "  result = args.src_tensor.Read(idx, Y, "
+           ", S, B);\n";
+      break;
+    case Axis::CHANNELS:
+      c += "  idx = args.indices.Read<int>(0, 0, 0, S * 4).x;\n";
+      c += "  args.src_tensor.ReadPerChannel(result.x, X, Y, idx, B);\n";
+      c += "  idx = args.indices.Read<int>(0, 0, 0, S * 4 + 1).x;\n";
+      c += "  args.src_tensor.ReadPerChannel(result.y, X, Y, idx, B);\n";
+      c += "  idx = args.indices.Read<int>(0, 0, 0, S * 4 + 2).x;\n";
+      c += "  args.src_tensor.ReadPerChannel(result.z, X, Y, idx, B);\n";
+      c += "  idx = args.indices.Read<int>(0, 0, 0, S * 4 + 3).x;\n";
+      c += "  args.src_tensor.ReadPerChannel(result.w, X, Y, idx, B);\n";
+      break;
+    default:
+      c += "  return;\n";
+  }
+  c += "  args.dst_tensor.Write(result, X, Y, S);\n";
+  c += "}\n";
+  return c;
+}
+}  // namespace
+
+GPUOperation CreateGather(const GpuInfo& gpu_info, const OperationDef& op_def,
+                          const GatherAttributes& attr) {
+  GPUOperation op(op_def);
+  op.AddSrcTensor("src_tensor", op_def.src_tensors[0]);
+  op.AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
+  if (op_def.src_tensors.size() == 1) {  // Constant indices
+    BHWC shape = BHWC(attr.indices.shape.v, 1, 1, 1);
+    TensorStorageType storage_type = GetStorageTypeForLinearTensor(
+        gpu_info, DataType::INT32, attr.indices.shape);
+    TensorDescriptor indices =
+        CreateBhwcTensorDescriptor(DataType::INT32, storage_type, shape);
+    indices.UploadData(attr.indices);
+    op.args_.AddObject("indices",
+                       std::make_unique<TensorDescriptor>(std::move(indices)));
+  } else {  // Runtime indices
+    op.AddSrcTensor("indices", op_def.src_tensors[1]);
+  }
+  op.code_ = GetGatherCode(op_def, attr);
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+  return op;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/gather.h

@@ -0,0 +1,34 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_GATHER_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_GATHER_H_
+
+#include "tensorflow/lite/delegates/gpu/common/gpu_info.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+GPUOperation CreateGather(const GpuInfo& gpu_info, const OperationDef& op_def,
+                          const GatherAttributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_GATHER_H_

tensorflow/lite/delegates/gpu/common/tasks/gather_test_util.cc

@@ -0,0 +1,241 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/gather_test_util.h"
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/precision.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tensor_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/gather.h"
+#include "tensorflow/lite/delegates/gpu/common/tensor.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status GatherBatchTest(TestExecutionEnvironment* env, bool constant_idx) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(5, 1, 1, 1);
+  src_tensor.data = {half(1.5f), half(2.4f), half(3.3f), half(4.2f),
+                     half(5.1f)};
+  std::vector<int> src_indices_data{1, 2, 3, 0, 1, 4, 2, 3, 1};
+  GatherAttributes attr;
+  attr.axis = Axis::BATCH;
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::BHWC});
+      TensorDescriptor src_0, src_1, dst;
+      src_0 = op_def.src_tensors[0];
+      src_0.UploadData(src_tensor);
+      dst.SetBHWDCShape(BHWDC(9, 1, 1, 1, 1));
+      if (!constant_idx) {
+        op_def.src_tensors.push_back({DataType::INT32, storage, Layout::BHWC});
+        TensorDescriptor src_1;
+        src_1 = op_def.src_tensors[1];
+        tflite::gpu::Tensor<BHWC, DataType::INT32> src_indices;
+        src_indices.shape = BHWC(9, 1, 1, 1);
+        src_indices.data = src_indices_data;
+        src_1.UploadData(src_indices);
+        GPUOperation operation = CreateGather(env->GetGpuInfo(), op_def, attr);
+        RETURN_IF_ERROR(env->ExecuteGPUOperation(
+            {&src_0, &src_1}, {&dst},
+            std::make_unique<GPUOperation>(std::move(operation))));
+      } else {
+        attr.indices.shape = Linear(9);
+        attr.indices.data = src_indices_data;
+        GPUOperation operation = CreateGather(env->GetGpuInfo(), op_def, attr);
+        RETURN_IF_ERROR(env->ExecuteGPUOperation(
+            {&src_0}, {&dst},
+            std::make_unique<GPUOperation>(std::move(operation))));
+      }
+      TensorFloat32 dst_tensor;
+      dst.DownloadData(&dst_tensor);
+      RETURN_IF_ERROR(PointWiseNear(
+          {half(2.4f), half(3.3f), half(4.2f), half(1.5f), half(2.4f),
+           half(5.1f), half(3.3f), half(4.2f), half(2.4f)},
+          dst_tensor.data, 0.0f));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status GatherHeightTest(TestExecutionEnvironment* env,
+                              bool constant_idx) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 5, 1, 1);
+  src_tensor.data = {half(1.5f), half(2.4f), half(3.3f), half(4.2f),
+                     half(5.1f)};
+  std::vector<int> src_indices_data{1, 2, 3, 0, 1, 4, 2, 3, 1};
+  GatherAttributes attr;
+  attr.axis = Axis::HEIGHT;
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::BHWC});
+      TensorDescriptor src_0, src_1, dst;
+      src_0 = op_def.src_tensors[0];
+      src_0.UploadData(src_tensor);
+      dst.SetBHWDCShape(BHWDC(1, 9, 1, 1, 1));
+      if (!constant_idx) {
+        op_def.src_tensors.push_back({DataType::INT32, storage, Layout::BHWC});
+        TensorDescriptor src_1;
+        src_1 = op_def.src_tensors[1];
+        tflite::gpu::Tensor<BHWC, DataType::INT32> src_indices;
+        src_indices.shape = BHWC(9, 1, 1, 1);
+        src_indices.data = src_indices_data;
+        src_1.UploadData(src_indices);
+        GPUOperation operation = CreateGather(env->GetGpuInfo(), op_def, attr);
+        RETURN_IF_ERROR(env->ExecuteGPUOperation(
+            {&src_0, &src_1}, {&dst},
+            std::make_unique<GPUOperation>(std::move(operation))));
+      } else {
+        attr.indices.shape = Linear(9);
+        attr.indices.data = src_indices_data;
+        GPUOperation operation = CreateGather(env->GetGpuInfo(), op_def, attr);
+        RETURN_IF_ERROR(env->ExecuteGPUOperation(
+            {&src_0}, {&dst},
+            std::make_unique<GPUOperation>(std::move(operation))));
+      }
+      TensorFloat32 dst_tensor;
+      dst.DownloadData(&dst_tensor);
+      RETURN_IF_ERROR(PointWiseNear(
+          {half(2.4f), half(3.3f), half(4.2f), half(1.5f), half(2.4f),
+           half(5.1f), half(3.3f), half(4.2f), half(2.4f)},
+          dst_tensor.data, 0.0f));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status GatherWidthTest(TestExecutionEnvironment* env, bool constant_idx) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 1, 5, 1);
+  src_tensor.data = {half(1.5f), half(2.4f), half(3.3f), half(4.2f),
+                     half(5.1f)};
+  std::vector<int> src_indices_data{1, 2, 3, 0, 1, 4, 2, 3, 1};
+  GatherAttributes attr;
+  attr.axis = Axis::WIDTH;
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::BHWC});
+      TensorDescriptor src_0, src_1, dst;
+      src_0 = op_def.src_tensors[0];
+      src_0.UploadData(src_tensor);
+      dst.SetBHWDCShape(BHWDC(1, 1, 9, 1, 1));
+      GPUOperation operation = CreateGather(env->GetGpuInfo(), op_def, attr);
+      if (!constant_idx) {
+        op_def.src_tensors.push_back({DataType::INT32, storage, Layout::BHWC});
+        TensorDescriptor src_1;
+        src_1 = op_def.src_tensors[1];
+        tflite::gpu::Tensor<BHWC, DataType::INT32> src_indices;
+        src_indices.shape = BHWC(9, 1, 1, 1);
+        src_indices.data = src_indices_data;
+        src_1.UploadData(src_indices);
+        GPUOperation operation = CreateGather(env->GetGpuInfo(), op_def, attr);
+        RETURN_IF_ERROR(env->ExecuteGPUOperation(
+            {&src_0, &src_1}, {&dst},
+            std::make_unique<GPUOperation>(std::move(operation))));
+      } else {
+        attr.indices.shape = Linear(9);
+        attr.indices.data = src_indices_data;
+        GPUOperation operation = CreateGather(env->GetGpuInfo(), op_def, attr);
+        RETURN_IF_ERROR(env->ExecuteGPUOperation(
+            {&src_0}, {&dst},
+            std::make_unique<GPUOperation>(std::move(operation))));
+      }
+      TensorFloat32 dst_tensor;
+      dst.DownloadData(&dst_tensor);
+      RETURN_IF_ERROR(PointWiseNear(
+          {half(2.4f), half(3.3f), half(4.2f), half(1.5f), half(2.4f),
+           half(5.1f), half(3.3f), half(4.2f), half(2.4f)},
+          dst_tensor.data, 0.0f));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status GatherChannelsTest(TestExecutionEnvironment* env,
+                                bool constant_idx) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 1, 1, 5);
+  src_tensor.data = {half(1.5f), half(2.4f), half(3.3f), half(4.2f),
+                     half(5.1f)};
+  std::vector<int> src_indices_data{1, 2, 3, 0, 1, 4, 2, 3, 1};
+  GatherAttributes attr;
+  attr.axis = Axis::CHANNELS;
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::BHWC});
+      TensorDescriptor src_0, src_1, dst;
+      src_0 = op_def.src_tensors[0];
+      src_0.UploadData(src_tensor);
+      dst.SetBHWDCShape(BHWDC(1, 1, 1, 1, 9));
+      GPUOperation operation = CreateGather(env->GetGpuInfo(), op_def, attr);
+      if (!constant_idx) {
+        op_def.src_tensors.push_back({DataType::INT32, storage, Layout::BHWC});
+        TensorDescriptor src_1;
+        src_1 = op_def.src_tensors[1];
+        tflite::gpu::Tensor<BHWC, DataType::INT32> src_indices;
+        src_indices.shape = BHWC(9, 1, 1, 1);
+        src_indices.data = src_indices_data;
+        src_1.UploadData(src_indices);
+        GPUOperation operation = CreateGather(env->GetGpuInfo(), op_def, attr);
+        RETURN_IF_ERROR(env->ExecuteGPUOperation(
+            {&src_0, &src_1}, {&dst},
+            std::make_unique<GPUOperation>(std::move(operation))));
+      } else {
+        attr.indices.shape = Linear(9);
+        attr.indices.data = src_indices_data;
+        GPUOperation operation = CreateGather(env->GetGpuInfo(), op_def, attr);
+        RETURN_IF_ERROR(env->ExecuteGPUOperation(
+            {&src_0}, {&dst},
+            std::make_unique<GPUOperation>(std::move(operation))));
+      }
+      TensorFloat32 dst_tensor;
+      dst.DownloadData(&dst_tensor);
+      RETURN_IF_ERROR(PointWiseNear(
+          {half(2.4f), half(3.3f), half(4.2f), half(1.5f), half(2.4f),
+           half(5.1f), half(3.3f), half(4.2f), half(2.4f)},
+          dst_tensor.data, 0.0f));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/gather_test_util.h

@@ -0,0 +1,34 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_GATHER_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_GATHER_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status GatherBatchTest(TestExecutionEnvironment* env, bool constant_idx);
+absl::Status GatherHeightTest(TestExecutionEnvironment* env, bool constant_idx);
+absl::Status GatherWidthTest(TestExecutionEnvironment* env, bool constant_idx);
+absl::Status GatherChannelsTest(TestExecutionEnvironment* env,
+                                bool constant_idx);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_GATHER_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/lstm.cc

@@ -0,0 +1,101 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/lstm.h"
+
+#include <string>
+
+#include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
+
+namespace tflite {
+namespace gpu {
+
+namespace {
+std::string GetLSTMCode(const OperationDef& op_def, const GpuInfo& gpu_info) {
+  std::string c;
+  c += "MAIN_FUNCTION(\n";
+  c += "$0) {\n";
+  c += "  int B = GLOBAL_ID_0;\n";
+  c += "  int Z = GLOBAL_ID_2;\n";
+  c += "  if (Z >= args.activation.Slices() || B >= args.activation.Batch()) "
+       "return;\n";
+  c += "  FLT4 prev_st = args.prev_state.Read(0, 0, Z, B);\n";
+  c += "  FLT4 r0 = args.intermediate.Read(0, 0, Z, B);\n";
+  c += "  int state_stride = args.activation.Slices();\n";
+  c += "  FLT4 r1 = args.intermediate.Read(0, 0, Z + state_stride, B);\n";
+  c += "  FLT4 r2 = args.intermediate.Read(0, 0, Z + state_stride * 2, B);\n";
+  c += "  FLT4 r3 = args.intermediate.Read(0, 0, Z + state_stride * 3, B);\n";
+  if (gpu_info.IsApiOpenCl() &&
+      op_def.precision != CalculationsPrecision::F32 && gpu_info.IsAdreno()) {
+    c += "  FLT4 input_gate;\n";
+    c += "  FLT4 new_input;\n";
+    c += "  FLT4 forget_gate;\n";
+    c += "  FLT4 output_gate;\n";
+    c += "  input_gate.x = native_recip(1.0h + native_exp(-r0.x));\n";
+    c += "  input_gate.y = native_recip(1.0h + native_exp(-r0.y));\n";
+    c += "  input_gate.z = native_recip(1.0h + native_exp(-r0.z));\n";
+    c += "  input_gate.w = native_recip(1.0h + native_exp(-r0.w));\n";
+    c += "  new_input.x = 1.0h - 2.0h * native_recip(1.0h + native_exp(2.0h * "
+         "r1.x));\n";
+    c += "  new_input.y = 1.0h - 2.0h * native_recip(1.0h + native_exp(2.0h * "
+         "r1.y));\n";
+    c += "  new_input.z = 1.0h - 2.0h * native_recip(1.0h + native_exp(2.0h * "
+         "r1.z));\n";
+    c += "  new_input.w = 1.0h - 2.0h * native_recip(1.0h + native_exp(2.0h * "
+         "r1.w));\n";
+    c += "  forget_gate.x = native_recip(1.0h + native_exp(-r2.x));\n";
+    c += "  forget_gate.y = native_recip(1.0h + native_exp(-r2.y));\n";
+    c += "  forget_gate.z = native_recip(1.0h + native_exp(-r2.z));\n";
+    c += "  forget_gate.w = native_recip(1.0h + native_exp(-r2.w));\n";
+    c += "  output_gate.x = native_recip(1.0h + native_exp(-r3.x));\n";
+    c += "  output_gate.y = native_recip(1.0h + native_exp(-r3.y));\n";
+    c += "  output_gate.z = native_recip(1.0h + native_exp(-r3.z));\n";
+    c += "  output_gate.w = native_recip(1.0h + native_exp(-r3.w));\n";
+  } else {
+    c += "  FLT4 input_gate  = INIT_FLT4(1.0f) / (INIT_FLT4(1.0f) + "
+         "exp(INIT_FLT4(-1.0f) "
+         "* r0));\n";
+    c += "  FLT4 new_input   = tanh(r1);\n";
+    c += "  FLT4 forget_gate = INIT_FLT4(1.0f) / (INIT_FLT4(1.0f) + "
+         "exp(INIT_FLT4(-1.0f) "
+         "* r2));\n";
+    c += "  FLT4 output_gate = INIT_FLT4(1.0f) / (INIT_FLT4(1.0f) + "
+         "exp(INIT_FLT4(-1.0f) "
+         "* r3));\n";
+  }
+  c += "  FLT4 new_st = input_gate * new_input + forget_gate * prev_st;\n";
+  c += "  FLT4 act_value = output_gate * tanh(new_st);\n";
+  c += "  args.activation.Write(act_value, 0, 0, Z, B);\n";
+  c += "  args.new_state.Write(new_st, 0, 0, Z, B);\n";
+  c += "}\n";
+  return c;
+}
+
+}  // namespace
+
+GPUOperation CreateLSTM(const OperationDef& definition,
+                        const GpuInfo& gpu_info) {
+  GPUOperation op(definition);
+  op.AddSrcTensor("intermediate", definition.src_tensors[0]);
+  op.AddSrcTensor("prev_state", definition.src_tensors[1]);
+  op.AddDstTensor("new_state", definition.dst_tensors[0]);
+  op.AddDstTensor("activation", definition.dst_tensors[1]);
+  op.code_ = GetLSTMCode(definition, gpu_info);
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+  return op;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/lstm.h

@@ -0,0 +1,33 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_LSTM_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_LSTM_H_
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+GPUOperation CreateLSTM(const OperationDef& definition,
+                        const GpuInfo& gpu_info);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_LSTM_H_

tensorflow/lite/delegates/gpu/common/tasks/lstm_test_util.cc

@@ -0,0 +1,82 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/lstm_test_util.h"
+
+#include <memory>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/lstm.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status LstmTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 1, 1, 16);
+  src_tensor.data = {
+      -std::log(2.0f), -std::log(2.0f), -std::log(2.0f), -std::log(2.0f),
+      std::log(3.0f),  std::log(3.0f),  std::log(3.0f),  std::log(3.0f),
+      -std::log(4.0f), -std::log(4.0f), -std::log(4.0f), -std::log(4.0f),
+      -std::log(5.0f), -std::log(5.0f), -std::log(5.0f), -std::log(5.0f)};
+  // input_gate = 1.0 / (1.0 + exp(log(2.0f))) = 1.0 / 3.0;
+  // new_input = tanh(log(3.0f)) = (exp(2 * log(3.0f)) - 1) / exp(2 * log(3.0f))
+  // + 1 = (9 - 1) / (9 + 1) = 0.8;
+  // forget_gate = 1.0 / (1.0 + exp(log(4.0f)))
+  //  = 1.0 / 5.0;
+  // output_gate = 1.0 / (1.0 + exp(log(5.0f))) = 1.0 / 6.0;
+  // new_st = input_gate * new_input + forget_gate * prev_st
+  //   = 1.0 / 3.0 * 0.8 + 1.0 / 5.0 * prev_st
+  //   = 4.0 / 15.0 + 3.0 / 15.0 = 7.0 / 15.0
+  // activation = output_gate * tanh(new_st)
+  TensorFloat32 prev_state;
+  prev_state.shape = BHWC(1, 1, 1, 4);
+  prev_state.data = {1.0f, 2.0f, 3.0f, 4.0f};
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
+      op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::BHWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::BHWC});
+      TensorFloat32 new_state;
+      TensorFloat32 new_activ;
+      GPUOperation operation = CreateLSTM(op_def, env->GetGpuInfo());
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {src_tensor, prev_state},
+          std::make_unique<GPUOperation>(std::move(operation)),
+          {BHWC(1, 1, 1, 4), BHWC(1, 1, 1, 4)}, {&new_state, &new_activ}));
+      RETURN_IF_ERROR(
+          PointWiseNear({7.0 / 15.0, 10.0 / 15.0, 13.0 / 15.0, 16.0 / 15.0},
+                        new_state.data, eps));
+      RETURN_IF_ERROR(PointWiseNear(
+          {static_cast<float>((1.0 / 6.0) * std::tanh(7.0 / 15.0)),
+           static_cast<float>((1.0 / 6.0) * std::tanh(10.0 / 15.0)),
+           static_cast<float>((1.0 / 6.0) * std::tanh(13.0 / 15.0)),
+           static_cast<float>((1.0 / 6.0) * std::tanh(16.0 / 15.0))},
+          new_activ.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/lstm_test_util.h

@@ -0,0 +1,30 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_LSTM_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_LSTM_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status LstmTest(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_LSTM_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/max_unpooling.cc

@@ -0,0 +1,154 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/max_unpooling.h"
+
+#include <string>
+
+namespace tflite {
+namespace gpu {
+namespace {
+void AppendConditionally(const std::string& value, const std::string& delimeter,
+                         std::string* result) {
+  if (!result->empty()) {
+    *result += delimeter;
+  }
+  *result += value;
+}
+
+std::string GetMaxUnpoolingKernelCode(const GpuInfo& gpu_info,
+                                      const OperationDef& op_def,
+                                      GPUOperation* op) {
+  op->AddSrcTensor("src_tensor", op_def.src_tensors[0]);
+  op->AddSrcTensor("src_indices", op_def.src_tensors[1]);
+  op->AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
+
+  std::string c;
+  c += "MAIN_FUNCTION($0) {\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+    c += "  args.src_indices.SetBatchRef(B);\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  int linear_id_1 = GLOBAL_ID_1;\n";
+    c += "  int Y = linear_id_1 / args.dst_tensor.Depth();\n";
+    c += "  int Z = linear_id_1 % args.dst_tensor.Depth();\n";
+  } else {
+    c += "  int Y = GLOBAL_ID_1;\n";
+  }
+  c += "  int S = GLOBAL_ID_2;\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "S >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
+  c += "  int src_x = (X + args.padding_x) / args.stride_x;\n";
+  c += "  int t_x = X - (src_x * args.stride_x - args.padding_x);\n";
+  c += "  int src_y = (Y + args.padding_y) / args.stride_y;\n";
+  c += "  int t_y = Y - (src_y * args.stride_y - args.padding_y);\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  int src_z = (Z + args.padding_z) / args.stride_z;\n";
+    c += "  int t_z = Z - (src_z * args.stride_z - args.padding_z);\n";
+    c += "  int t_index = (t_y * args.kernel_size_x + t_x) * "
+         "args.kernel_size_z + t_z;\n";
+  } else {
+    c += "  int t_index = t_y * args.kernel_size_x + t_x;\n";
+  }
+  std::string inbounds_check;
+  if (!op_def.src_tensors[0].SupportsZeroClamp(Axis::WIDTH, gpu_info) ||
+      !op_def.src_tensors[1].SupportsZeroClamp(Axis::WIDTH, gpu_info)) {
+    c += "  bool inside_x = src_x >= 0 && src_x < args.src_tensor.Width();\n";
+    c += "  src_x = clamp(src_x, 0, args.src_tensor.Width() - 1);\n";
+    AppendConditionally("inside_x", " && ", &inbounds_check);
+  }
+  if (!op_def.src_tensors[0].SupportsZeroClamp(Axis::HEIGHT, gpu_info) ||
+      !op_def.src_tensors[1].SupportsZeroClamp(Axis::HEIGHT, gpu_info)) {
+    c += "  bool inside_y = src_y >= 0 && src_y < args.src_tensor.Height();\n";
+    c += "  src_y = clamp(src_y, 0, args.src_tensor.Height() - 1);\n";
+    AppendConditionally("inside_y", " && ", &inbounds_check);
+  }
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    if (!op_def.src_tensors[0].SupportsZeroClamp(Axis::DEPTH, gpu_info) ||
+        !op_def.src_tensors[1].SupportsZeroClamp(Axis::DEPTH, gpu_info)) {
+      c += "  bool inside_z = src_z >= 0 && src_z < args.src_tensor.Depth();\n";
+      c += "  src_z = clamp(src_z, 0, args.src_tensor.Depth() - 1);\n";
+      AppendConditionally("inside_z", " && ", &inbounds_check);
+    }
+  }
+  std::string src_args = op_def.dst_tensors[0].HasAxis(Axis::DEPTH)
+                             ? "src_x, src_y, src_z, S"
+                             : "src_x, src_y, S";
+  c +=
+      "  args.src_tensor::type src = args.src_tensor.Read(" + src_args + ");\n";
+  c += "  int4 ind = args.src_indices.Read<int>(" + src_args + ");\n";
+  if (!inbounds_check.empty()) {
+    c += "  src *= INIT_FLT(" + inbounds_check + ");\n";
+    c += "  ind *= INIT_INT(" + inbounds_check + ");\n";
+  }
+  c += "  args.src_tensor::type result;\n";
+  c += "  result.x = t_index == ind.x ? src.x : INIT_FLT(0.0f);\n";
+  c += "  result.y = t_index == ind.y ? src.y : INIT_FLT(0.0f);\n";
+  c += "  result.z = t_index == ind.z ? src.z : INIT_FLT(0.0f);\n";
+  c += "  result.w = t_index == ind.w ? src.w : INIT_FLT(0.0f);\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  args.dst_tensor.Write(result, X, Y, Z, S);\n";
+  } else {
+    c += "  args.dst_tensor.Write(result, X, Y, S);\n";
+  }
+  c += "}\n";
+  return c;
+}
+}  // namespace
+
+GPUOperation CreateMaxUnpooling(const GpuInfo& gpu_info,
+                                const OperationDef& definition,
+                                const MaxUnpooling2DAttributes& attr) {
+  GPUOperation op(definition);
+  op.args_.AddInt("kernel_size_x", attr.kernel.w);
+  op.args_.AddInt("padding_x", attr.padding.appended.w);
+  op.args_.AddInt("stride_x", attr.strides.w);
+  op.args_.AddInt("kernel_size_y", attr.kernel.h);
+  op.args_.AddInt("padding_y", attr.padding.appended.h);
+  op.args_.AddInt("stride_y", attr.strides.h);
+  op.code_ = GetMaxUnpoolingKernelCode(gpu_info, definition, &op);
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+  return op;
+}
+
+GPUOperation CreateMaxUnpooling(const GpuInfo& gpu_info,
+                                const OperationDef& definition,
+                                const MaxUnpooling3DAttributes& attr) {
+  GPUOperation op(definition);
+  op.args_.AddInt("kernel_size_x", attr.kernel.w);
+  op.args_.AddInt("padding_x", attr.padding.appended.w);
+  op.args_.AddInt("stride_x", attr.strides.w);
+  op.args_.AddInt("kernel_size_y", attr.kernel.h);
+  op.args_.AddInt("padding_y", attr.padding.appended.h);
+  op.args_.AddInt("stride_y", attr.strides.h);
+  op.args_.AddInt("kernel_size_z", attr.kernel.d);
+  op.args_.AddInt("padding_z", attr.padding.appended.d);
+  op.args_.AddInt("stride_z", attr.strides.d);
+  op.code_ = GetMaxUnpoolingKernelCode(gpu_info, definition, &op);
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+  return op;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/max_unpooling.h

@@ -0,0 +1,38 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_MAX_UNPOOLING_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_MAX_UNPOOLING_H_
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+GPUOperation CreateMaxUnpooling(const GpuInfo& gpu_info,
+                                const OperationDef& definition,
+                                const MaxUnpooling2DAttributes& attr);
+
+GPUOperation CreateMaxUnpooling(const GpuInfo& gpu_info,
+                                const OperationDef& definition,
+                                const MaxUnpooling3DAttributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_MAX_UNPOOLING_H_

tensorflow/lite/delegates/gpu/common/tasks/max_unpooling_test_util.cc

@@ -0,0 +1,69 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/max_unpooling_test_util.h"
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/max_unpooling.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status MaxUnpoolingTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 2, 1);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f};
+  TensorFloat32 src_ind_tensor;
+  src_ind_tensor.shape = BHWC(1, 2, 2, 1);
+  src_ind_tensor.data = {0.1f, 1.1f, 2.1f, 3.1f};
+
+  MaxUnpooling2DAttributes attr;
+  attr.padding.prepended = HW(0, 0);
+  attr.padding.appended = HW(0, 0);
+  attr.strides = HW(2, 2);
+  attr.kernel = HW(2, 2);
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation =
+          CreateMaxUnpooling(env->GetGpuInfo(), op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {src_tensor, src_ind_tensor},
+          std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 4, 4, 1), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f,
+                         0.0f, 0.0f, 0.0f, 2.0f, 0.0f, 0.0f, 3.0f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/max_unpooling_test_util.h

@@ -0,0 +1,30 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_MAX_UNPOOLING_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_MAX_UNPOOLING_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status MaxUnpoolingTest(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_MAX_UNPOOLING_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/mean_stddev_normalization.cc

@@ -0,0 +1,554 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/mean_stddev_normalization.h"
+
+#include <algorithm>
+#include <map>
+#include <memory>
+#include <set>
+#include <string>
+#include <utility>
+
+#include "absl/strings/substitute.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/util.h"
+
+namespace tflite {
+namespace gpu {
+
+namespace {
+
+absl::Status CheckIfValidNodeOfType(const Node* node,
+                                    OperationType required_type) {
+  if (node == nullptr) {
+    return absl::NotFoundError("Invalid node.");
+  }
+  if (OperationTypeFromString(node->operation.type) != required_type) {
+    return absl::NotFoundError("Type mismatch.");
+  }
+  return absl::OkStatus();
+}
+
+absl::Status GetElementwiseScalarValue(const Node* node, float* result) {
+  auto attr = absl::any_cast<ElementwiseAttributes>(node->operation.attributes);
+  const float* value = std::get_if<float>(&attr.param);
+  if (!value) {
+    return absl::NotFoundError("Not a scalar value inside attributes.");
+  }
+  *result = *value;
+  return absl::OkStatus();
+}
+
+absl::Status GetNextSingleNode(const GraphFloat32& graph, const Node& node,
+                               OperationType next_type, Node** next_node) {
+  auto consumers = graph.FindConsumers(graph.FindOutputs(node.id)[0]->id);
+  if (consumers.size() != 1) {
+    return absl::NotFoundError("Not a single consumer.");
+  }
+  RETURN_IF_ERROR(CheckIfValidNodeOfType(consumers[0], next_type));
+  *next_node = consumers[0];
+  return absl::OkStatus();
+}
+
+std::string GetReduceCode(const std::string& value, int3 work_group_size,
+                          bool two_step) {
+  int reduction_size = work_group_size.z;
+  std::string mem_name = work_group_size.x * work_group_size.y != 1
+                             ? "shared_mem[LOCAL_ID_1][LOCAL_ID_0]"
+                             : "shared_mem";
+  if (reduction_size <= 8) {
+    std::string result;
+    result += "  {  // reduction\n";
+    result += "    " + mem_name + "[local_id] = " + value + ";\n";
+    result += "    LOCAL_MEM_BARRIER;\n";
+    result += "    if (LOCAL_ID_2 == 0) {\n";
+    result += "      " + value + " = " + mem_name + "[0];\n";
+    for (int i = 1; i < reduction_size; ++i) {
+      result += "      " + value + " += " + mem_name + "[" + std::to_string(i) +
+                "];\n";
+    }
+    result += "      " + mem_name + "[0] = " + value + ";\n";
+    result += "    }\n";
+    result += "    LOCAL_MEM_BARRIER;\n";
+    result += "    " + value + " = " + mem_name + "[0];\n";
+    if (two_step) {
+      result += "    LOCAL_MEM_BARRIER;\n";
+    }
+    result += "  }\n";
+    return result;
+  } else {
+    // In the reduction step add upper half of the still-to-be-summed vector to
+    // the lower half, while taking care of odd sizes and rounding. E.g.:
+    // Number of items still to be summed before: 5
+    // Local memory before: [a, b, c, d, e];
+    // Local memory after: [a+d, b+e, c, d, e];
+    // Threads doing work: id < 2 = floor(5/2)
+    // Offset to the added items: 3 = ceil(5/2)
+    // Number of items still to be summed after: 3 = ceil(5/2)
+    return absl::Substitute(R"(
+  {  // reduction, all threads inside workgroup must execute this code
+    $2[local_id] = $1;
+    LOCAL_MEM_BARRIER;
+    // The number of items still need to be summed
+    int reduction_size = $0;
+    while (reduction_size > 1) {
+      int active_thread_limit = reduction_size / 2;
+      int offset = (reduction_size + 1) / 2;
+      if (local_id < active_thread_limit) {
+        $1 += $2[local_id + offset];
+        $2[local_id] = $1;
+      }
+      LOCAL_MEM_BARRIER;
+      reduction_size = offset;
+    }
+    $1 = $2[0];
+  }
+)",
+                            reduction_size, value, mem_name);
+  }
+}
+
+std::string ZeroClampVec4Code(const std::string& slice_name,
+                              const std::string& channels_name,
+                              const std::string& value_name) {
+  return absl::Substitute(R"(
+    // no need to check first element, always valid
+    if ($0 * 4 + 1 >= $1) { $2.y = 0.0f; }
+    if ($0 * 4 + 2 >= $1) { $2.z = 0.0f; }
+    if ($0 * 4 + 3 >= $1) { $2.w = 0.0f; }
+)",
+                          slice_name, channels_name, value_name);
+}
+
+bool UseWorkGroupReduction(const GpuInfo& gpu_info, const BHWC& shape) {
+  const int tensor_slices = DivideRoundUp(shape.c, 4);
+  if (gpu_info.IsAdreno() && tensor_slices <= 32 &&
+      shape.w * shape.h * shape.b >= 128) {
+    return false;
+  } else {
+    return true;
+  }
+}
+
+int3 GetRecommendedWorkGroupSize(const GpuInfo& gpu_info, const BHWC& shape) {
+  const int tensor_slices = DivideRoundUp(shape.c, 4);
+  int desired_work_group_size = gpu_info.GetMaxWorkGroupSizeForZ();
+  if (gpu_info.IsMali()) {
+    // Don't use more than 64 work items per work group on ARM Mali. They
+    // implement local memory using the global memory, larger workgroups have
+    // severe performance penalty.
+    desired_work_group_size = 64;
+  }
+  if (gpu_info.IsAdreno()) {
+    AdrenoInfo info = gpu_info.adreno_info;
+    desired_work_group_size = 256;
+    if (info.IsAdreno3xx()) {
+      if (info.adreno_gpu == AdrenoGpu::kAdreno320 ||
+          info.adreno_gpu == AdrenoGpu::kAdreno330) {
+        desired_work_group_size = 128;
+      } else {
+        desired_work_group_size = 64;
+      }
+    } else if (info.IsAdreno4xx()) {
+      if (info.adreno_gpu == AdrenoGpu::kAdreno430) {
+        desired_work_group_size = 256;
+      } else {
+        desired_work_group_size = 128;
+      }
+    } else if (info.IsAdreno5xx()) {
+      if (info.adreno_gpu == AdrenoGpu::kAdreno530 ||
+          info.adreno_gpu == AdrenoGpu::kAdreno540) {
+        desired_work_group_size = 256;
+      } else {
+        desired_work_group_size = 128;
+      }
+    }
+  }
+  if (gpu_info.IsPowerVR()) {
+    desired_work_group_size = 64;
+  }
+  if (gpu_info.IsApple()) {
+    desired_work_group_size = 64;
+  }
+  if (gpu_info.IsAMD()) {
+    desired_work_group_size = 512;
+  }
+  int3 work_group_size(1, 1, 1);
+  if (shape.w * shape.h == 1) {
+    desired_work_group_size =
+        std::min(desired_work_group_size, gpu_info.GetMaxWorkGroupSizeForZ());
+    while (desired_work_group_size >= tensor_slices * 2) {
+      desired_work_group_size /= 2;
+    }
+    work_group_size.x = 1;
+    work_group_size.y = 1;
+    work_group_size.z = desired_work_group_size;
+  } else {
+    if (tensor_slices >= 16) {
+      work_group_size.z = 8;
+    } else if (tensor_slices >= 10) {
+      work_group_size.z = 4;
+    } else {
+      std::map<int, int> slices_to_group_size = {
+          {1, 1}, {2, 2}, {3, 3}, {4, 4}, {5, 3},
+          {6, 3}, {7, 4}, {8, 4}, {9, 3},
+      };
+      work_group_size.z = slices_to_group_size[tensor_slices];
+    }
+    desired_work_group_size =
+        std::min(desired_work_group_size, gpu_info.GetMaxWorkGroupTotalSize());
+    work_group_size.x = 1;
+    work_group_size.y =
+        desired_work_group_size / AlignByN(work_group_size.z, 4);
+    while (work_group_size.y > work_group_size.x) {
+      work_group_size.y /= 2;
+      work_group_size.x *= 2;
+    }
+  }
+  return work_group_size;
+}
+
+std::string GetVarianceCalculationCode(const GpuInfo& gpu_info,
+                                       bool work_group_reduction,
+                                       const int3& work_group_size,
+                                       bool has_batch, bool channels_x4,
+                                       bool two_step) {
+  std::string c;
+  if (work_group_reduction && gpu_info.IsApiOpenCl()) {
+    c += "__attribute__((reqd_work_group_size(" +
+         std::to_string(work_group_size.x) + ", " +
+         std::to_string(work_group_size.y) + ", " +
+         std::to_string(work_group_size.z) + ")))\n";
+  }
+  c += "MAIN_FUNCTION($0) {\n";
+  if (work_group_reduction) {
+    std::string accum_type = two_step ? "float" : "float2";
+    if (work_group_size.x * work_group_size.y == 1) {
+      c += "__local " + accum_type + " shared_mem[" +
+           std::to_string(work_group_size.z) + "];\n";
+    } else {
+      c += "__local " + accum_type + " shared_mem[" +
+           std::to_string(work_group_size.y) + "][" +
+           std::to_string(work_group_size.x) + "][" +
+           std::to_string(work_group_size.z) + "];\n";
+    }
+  }
+  if (has_batch) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  c += "  int Y = GLOBAL_ID_1;\n";
+  if (!work_group_reduction) {
+    c += "  if (X >= args.dst_tensor.Width()) { return; }\n";
+    c += "  if (Y >= args.dst_tensor.Height()) { return; }\n";
+  }
+  if (!two_step) {
+    c += "  float4 private_sum4_sq = INIT_FLOAT4(0.0f);\n";
+  }
+  if (work_group_reduction) {
+    c += "  int local_id = LOCAL_ID_2;\n";
+    c += "  int reduction_group_size = GROUP_SIZE_2;\n";
+  } else {
+    c += "  int local_id = 0;\n";
+    c += "  int reduction_group_size = 1;\n";
+  }
+  c += R"(
+  float4 private_sum4 = INIT_FLOAT4(0.0f);
+  for (int S = local_id; S < args.src_tensor.Slices(); S += reduction_group_size) {
+    int x_clamped = min(X, args.src_tensor.Width() - 1);
+    int y_clamped = min(Y, args.src_tensor.Height() - 1);
+    float4 t = args.src_tensor.Read<float>(x_clamped, y_clamped, S);
+)";
+  if (!channels_x4) {
+    c += ZeroClampVec4Code("S", "args.src_tensor.Channels()", "t");
+  }
+  if (two_step) {
+    c += "    private_sum4 += t;\n";
+    c += "  }\n";
+    c += "  float sum = dot(private_sum4, INIT_FLOAT4(1.0f));\n";
+  } else {
+    c += "    private_sum4 += t;\n";
+    c += "    private_sum4_sq += t * t;\n";
+    c += "  }\n";
+    c += "  float2 sum;\n";
+    c += "  sum.x = dot(private_sum4, INIT_FLOAT4(1.0f));\n";
+    c += "  sum.y = dot(private_sum4_sq, INIT_FLOAT4(1.0f));\n";
+  }
+  if (work_group_reduction) {
+    c += GetReduceCode("sum", work_group_size, two_step);
+  }
+  if (two_step) {
+    c += R"(
+  // Calculate the mean
+  float mean = sum * args.inv_ch_count;
+  // Calculate the squared sum of the difference from the mean.
+  float4 private_sum_diff_sq4 = INIT_FLOAT4(0.0f);
+  for (int S = local_id; S < args.src_tensor.Slices(); S += reduction_group_size) {
+    int x_clamped = min(X, args.src_tensor.Width() - 1);
+    int y_clamped = min(Y, args.src_tensor.Height() - 1);
+    float4 t = args.src_tensor.Read<float>(x_clamped, y_clamped, S);
+    float4 diff = t - mean;)";
+    if (!channels_x4) {
+      c += ZeroClampVec4Code("S", "args.src_tensor.Channels()", "diff");
+    }
+    c += R"(
+    private_sum_diff_sq4 += diff * diff;
+  }
+  // Reduce
+  float sum_diff_sq = dot(private_sum_diff_sq4, INIT_FLOAT4(1.0f));
+)";
+    if (work_group_reduction) {
+      c += GetReduceCode("sum_diff_sq", work_group_size, two_step);
+    }
+    c += "  float variance = sum_diff_sq * args.inv_ch_count;\n";
+  } else {
+    c += "  float mean = sum.x * args.inv_ch_count;\n";
+    c += "  float mean_sq = sum.y * args.inv_ch_count;\n";
+    c += "  float variance = mean_sq - mean * mean;\n";
+  }
+  if (work_group_reduction) {
+    c += "  // no more shared memory usage, 'useless' threads can exit now\n";
+    c += "  if (X >= args.dst_tensor.Width()) { return; }\n";
+    c += "  if (Y >= args.dst_tensor.Height()) { return; }\n";
+  }
+  return c;
+}
+}  // namespace
+
+MeanStdDevNormalization::MeanStdDevNormalization(const OperationDef& definition,
+                                                 const GpuInfo& gpu_info,
+                                                 const BHWC& shape,
+                                                 float variance_bias,
+                                                 bool two_step)
+    : GPUOperation(definition) {
+  work_group_reduction_ = UseWorkGroupReduction(gpu_info, shape);
+  if (work_group_reduction_) {
+    work_group_size_ = GetRecommendedWorkGroupSize(gpu_info, shape);
+  } else {
+    work_group_size_ = int3(8, 8, 1);
+  }
+  args_.AddFloat("variance_bias", variance_bias);
+  args_.AddFloat("inv_ch_count", 1.0f / shape.c);
+  AddSrcTensor("src_tensor", definition_.src_tensors[0]);
+  AddDstTensor("dst_tensor", definition_.dst_tensors[0]);
+  code_ = GetNormalizationCode(gpu_info, shape.c % 4 == 0, two_step);
+}
+
+std::string MeanStdDevNormalization::GetNormalizationCode(
+    const GpuInfo& gpu_info, bool channels_x4, bool two_step) {
+  std::string c = GetVarianceCalculationCode(
+      gpu_info, work_group_reduction_, work_group_size_,
+      definition_.dst_tensors[0].HasAxis(Axis::BATCH), channels_x4, two_step);
+  c += R"(
+  float stddev_inv = rsqrt(variance + args.variance_bias);
+  // Calculate (t-mean)/stddev for each element
+  for (int S = local_id; S < args.src_tensor.Slices(); S += reduction_group_size) {
+    float4 t = args.src_tensor.Read<float>(X, Y, S);
+    FLT4 result = TO_FLT4((t - mean) * stddev_inv);
+    args.dst_tensor.Write(result, X, Y, S);
+  }
+})";
+  return c;
+}
+
+int3 MeanStdDevNormalization::GetGridSize() const {
+  const int grid_x = dst_[0]->Width() * dst_[0]->Batch();
+  const int grid_y = dst_[0]->Height();
+  const int grid_z = work_group_reduction_ ? work_group_size_.z : 1;
+  return int3(grid_x, grid_y, grid_z);
+}
+
+MeanStdDevNormalization CreateMeanStdDevNormalization(
+    const OperationDef& definition, const GpuInfo& gpu_info, const BHWC& shape,
+    float variance_bias, bool two_step) {
+  return MeanStdDevNormalization(definition, gpu_info, shape, variance_bias,
+                                 two_step);
+}
+
+absl::Status TryMeanStdDevNormalization(
+    const GpuInfo& gpu_info, CalculationsPrecision precision,
+    const GraphFloat32& graph, NodeId first_node_id,
+    const std::map<ValueId, TensorDescriptor>& tensor_descriptors,
+    std::set<NodeId>* consumed_nodes, GPUOperationsSubgraph* gpu_subgraph) {
+  Node* first_mean_node = graph.GetNode(first_node_id);
+  RETURN_IF_ERROR(CheckIfValidNodeOfType(first_mean_node, OperationType::MEAN));
+  auto first_mean_attr =
+      absl::any_cast<MeanAttributes>(first_mean_node->operation.attributes);
+  if (first_mean_attr.dims != std::set<Axis>{Axis::CHANNELS}) {
+    return absl::NotFoundError("MeanStdDevNormalization not suitable.");
+  }
+  Node* sub_node;
+  RETURN_IF_ERROR(GetNextSingleNode(graph, *first_mean_node, OperationType::SUB,
+                                    &sub_node));
+  auto sub_inputs = graph.FindInputs(sub_node->id);
+  if (sub_inputs.size() != 2) {
+    return absl::NotFoundError("MeanStdDevNormalization not suitable.");
+  } else {
+    // checking structure
+    //       input
+    //       /    \
+    //      |    mean
+    //       \    /
+    //     substraction
+    Node* sub_first_parent = graph.FindProducer(sub_inputs[0]->id);
+    Node* sub_second_parent = graph.FindProducer(sub_inputs[1]->id);
+    if (sub_second_parent != first_mean_node) {
+      return absl::NotFoundError("MeanStdDevNormalization not suitable.");
+    }
+    auto mean_inputs = graph.FindInputs(first_mean_node->id);
+    Node* mean_parent = graph.FindProducer(mean_inputs[0]->id);
+    if (mean_parent != sub_first_parent) {
+      return absl::NotFoundError("MeanStdDevNormalization not suitable.");
+    }
+  }
+  auto sub_output = graph.FindOutputs(sub_node->id)[0]->id;
+  auto consumers = graph.FindConsumers(sub_output);
+  if (consumers.size() != 2) {
+    return absl::NotFoundError("MeanStdDevNormalization not suitable.");
+  }
+  Node* square_node = consumers[0];
+  Node* sub_child_mul_node = consumers[1];
+  if (!CheckIfValidNodeOfType(square_node, OperationType::SQUARE).ok()) {
+    square_node = consumers[1];
+    sub_child_mul_node = consumers[0];
+  }
+  RETURN_IF_ERROR(CheckIfValidNodeOfType(square_node, OperationType::SQUARE));
+  RETURN_IF_ERROR(
+      CheckIfValidNodeOfType(sub_child_mul_node, OperationType::MUL));
+  Node* second_mean_node;
+  RETURN_IF_ERROR(GetNextSingleNode(graph, *square_node, OperationType::MEAN,
+                                    &second_mean_node));
+  auto second_mean_attr =
+      absl::any_cast<MeanAttributes>(second_mean_node->operation.attributes);
+  if (second_mean_attr.dims != std::set<Axis>{Axis::CHANNELS}) {
+    return absl::NotFoundError("MeanStdDevNormalization not suitable.");
+  }
+  Node* add_node;
+  RETURN_IF_ERROR(GetNextSingleNode(graph, *second_mean_node,
+                                    OperationType::ADD, &add_node));
+  float add_value;
+  RETURN_IF_ERROR(GetElementwiseScalarValue(add_node, &add_value));
+  Node* rsqrt_node;
+  RETURN_IF_ERROR(
+      GetNextSingleNode(graph, *add_node, OperationType::RSQRT, &rsqrt_node));
+  Node* mul_node;
+  RETURN_IF_ERROR(
+      GetNextSingleNode(graph, *rsqrt_node, OperationType::MUL, &mul_node));
+  if (sub_child_mul_node != mul_node) {
+    return absl::NotFoundError("MeanStdDevNormalization not suitable.");
+  }
+
+  OperationDef op_def;
+  op_def.precision = precision;
+  auto input_id = graph.FindInputs(first_mean_node->id)[0]->id;
+  auto it = tensor_descriptors.find(input_id);
+  if (it != tensor_descriptors.end()) {
+    op_def.src_tensors.push_back(it->second);
+  }
+  auto output_id = graph.FindOutputs(mul_node->id)[0]->id;
+  it = tensor_descriptors.find(output_id);
+  if (it != tensor_descriptors.end()) {
+    op_def.dst_tensors.push_back(it->second);
+  }
+
+  auto subgraph_inputs = graph.FindInputs(first_mean_node->id);
+  auto subgraph_outputs = graph.FindOutputs(mul_node->id);
+  std::unique_ptr<GPUOperation>* gpu_op =
+      InitSingleOpSubgraph(subgraph_inputs, subgraph_outputs, gpu_subgraph);
+  *gpu_op =
+      std::make_unique<MeanStdDevNormalization>(CreateMeanStdDevNormalization(
+          op_def, gpu_info, subgraph_inputs[0]->tensor.shape, add_value,
+          /*two_step*/ false));
+
+  consumed_nodes->insert(first_mean_node->id);
+  consumed_nodes->insert(sub_node->id);
+  consumed_nodes->insert(square_node->id);
+  consumed_nodes->insert(second_mean_node->id);
+  consumed_nodes->insert(add_node->id);
+  consumed_nodes->insert(rsqrt_node->id);
+  consumed_nodes->insert(mul_node->id);
+
+  return absl::OkStatus();
+}
+
+LayerNormalization::LayerNormalization(
+    const OperationDef& definition, const GpuInfo& gpu_info, const BHWC& shape,
+    float variance_bias, const Tensor<Linear, DataType::FLOAT32>& mul_linear,
+    const Tensor<Linear, DataType::FLOAT32>& sub_linear, bool two_step)
+    : GPUOperation(definition) {
+  work_group_reduction_ = UseWorkGroupReduction(gpu_info, shape);
+  if (work_group_reduction_) {
+    work_group_size_ = GetRecommendedWorkGroupSize(gpu_info, shape);
+  } else {
+    work_group_size_ = int3(8, 8, 1);
+  }
+  args_.AddFloat("variance_bias", variance_bias);
+  args_.AddFloat("inv_ch_count", 1.0f / shape.c);
+  AddSrcTensor("src_tensor", definition_.src_tensors[0]);
+  AddDstTensor("dst_tensor", definition_.dst_tensors[0]);
+  TensorDescriptor mul_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), mul_linear);
+  args_.AddObject("mul_linear", std::make_unique<TensorDescriptor>(
+                                    std::move(mul_tensor_desc)));
+  TensorDescriptor sub_tensor_desc = CreateConstantLinearTensorDescriptor(
+      gpu_info, definition.src_tensors[0].GetDataType(), sub_linear);
+  args_.AddObject("sub_linear", std::make_unique<TensorDescriptor>(
+                                    std::move(sub_tensor_desc)));
+  code_ = GetNormalizationCode(gpu_info, shape.c % 4 == 0, two_step);
+}
+
+std::string LayerNormalization::GetNormalizationCode(const GpuInfo& gpu_info,
+                                                     bool channels_x4,
+                                                     bool two_step) {
+  std::string c = GetVarianceCalculationCode(
+      gpu_info, work_group_reduction_, work_group_size_,
+      definition_.dst_tensors[0].HasAxis(Axis::BATCH), channels_x4, two_step);
+  c += R"(
+  float stddev_inv = rsqrt(variance + args.variance_bias);
+  for (int S = local_id; S < args.src_tensor.Slices(); S += reduction_group_size) {
+    float4 t = args.src_tensor.Read<float>(X, Y, S);
+    float4 mul0_res = stddev_inv * args.mul_linear.Read<float>(S);
+    float4 mul1_res = mul0_res * t;
+    float4 mul2_res = mul0_res * mean;
+    float4 sub_res = args.sub_linear.Read<float>(S) - mul2_res;
+    FLT4 result = TO_FLT4(mul1_res + sub_res);
+    args.dst_tensor.Write(result, X, Y, S);
+  }
+})";
+  return c;
+}
+
+int3 LayerNormalization::GetGridSize() const {
+  const int grid_x = dst_[0]->Width() * dst_[0]->Batch();
+  const int grid_y = dst_[0]->Height();
+  const int grid_z = work_group_reduction_ ? work_group_size_.z : 1;
+  return int3(grid_x, grid_y, grid_z);
+}
+
+LayerNormalization CreateLayerNormalization(
+    const OperationDef& definition, const GpuInfo& gpu_info, const BHWC& shape,
+    float variance_bias, const Tensor<Linear, DataType::FLOAT32>& mul_linear,
+    const Tensor<Linear, DataType::FLOAT32>& sub_linear, bool two_step) {
+  return LayerNormalization(definition, gpu_info, shape, variance_bias,
+                            mul_linear, sub_linear, two_step);
+}
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/mean_stddev_normalization.h

@@ -0,0 +1,141 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_MEAN_STDDEV_NORMALIZATION_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_MEAN_STDDEV_NORMALIZATION_H_
+
+#include <map>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/model.h"
+#include "tensorflow/lite/delegates/gpu/common/selectors/subgraph.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+// Implements tensor_utils::MeanStddevNormalization
+class MeanStdDevNormalization : public GPUOperation {
+ public:
+  explicit MeanStdDevNormalization(const OperationDef& definition,
+                                   const GpuInfo& gpu_info, const BHWC& shape,
+                                   float variance_bias, bool two_step);
+
+  void GetPossibleKernelWorkGroups(
+      TuningType tuning_type, const GpuInfo& gpu_info,
+      const KernelInfo& kernel_info,
+      std::vector<int3>* work_groups) const override {
+    if (!work_group_reduction_) {
+      GetPossibleWorkGroups(tuning_type, gpu_info, kernel_info, grid_size_,
+                            work_groups);
+      return;
+    }
+    work_groups->push_back(work_group_size_);
+  }
+  int3 GetGridSize() const override;
+
+  // Move only
+  MeanStdDevNormalization(MeanStdDevNormalization&& kernel) = default;
+  MeanStdDevNormalization& operator=(MeanStdDevNormalization&& kernel) =
+      default;
+  MeanStdDevNormalization(const MeanStdDevNormalization&) = delete;
+  MeanStdDevNormalization& operator=(const MeanStdDevNormalization&) = delete;
+
+ private:
+  std::string GetNormalizationCode(const GpuInfo& gpu_info, bool channels_x4,
+                                   bool two_step);
+  bool work_group_reduction_ = true;
+};
+
+// std dev can be calculated in single step, but two step algorithm can
+// provide more stable and robust results
+MeanStdDevNormalization CreateMeanStdDevNormalization(
+    const OperationDef& definition, const GpuInfo& gpu_info, const BHWC& shape,
+    float variance_bias = 1.0e-8f, bool two_step = true);
+
+// MeanStdDevNormalization fusion works with this subgraph
+//       input
+//       /    \
+//      |    mean
+//       \    /
+//     substraction
+//       /    \
+//      |      |
+//      |    square
+//      |      |
+//      |     mean
+//      |      |
+//      |     add
+//      |      |
+//      |    rsqrt
+//      |      |
+//       \    /
+//    multiplication
+//          |
+//        output
+absl::Status TryMeanStdDevNormalization(
+    const GpuInfo& gpu_info, CalculationsPrecision precision,
+    const GraphFloat32& graph, NodeId first_node_id,
+    const std::map<ValueId, TensorDescriptor>& tensor_descriptors,
+    std::set<NodeId>* consumed_nodes, GPUOperationsSubgraph* gpu_subgraph);
+
+class LayerNormalization : public GPUOperation {
+ public:
+  LayerNormalization(const OperationDef& definition, const GpuInfo& gpu_info,
+                     const BHWC& shape, float variance_bias,
+                     const Tensor<Linear, DataType::FLOAT32>& mul_linear,
+                     const Tensor<Linear, DataType::FLOAT32>& sub_linear,
+                     bool two_step);
+
+  void GetPossibleKernelWorkGroups(
+      TuningType tuning_type, const GpuInfo& gpu_info,
+      const KernelInfo& kernel_info,
+      std::vector<int3>* work_groups) const override {
+    if (!work_group_reduction_) {
+      GetPossibleWorkGroups(tuning_type, gpu_info, kernel_info, grid_size_,
+                            work_groups);
+      return;
+    }
+    work_groups->push_back(work_group_size_);
+  }
+  int3 GetGridSize() const override;
+
+  // Move only
+  LayerNormalization(LayerNormalization&& kernel) = default;
+  LayerNormalization& operator=(LayerNormalization&& kernel) = default;
+  LayerNormalization(const LayerNormalization&) = delete;
+  LayerNormalization& operator=(const LayerNormalization&) = delete;
+
+ private:
+  std::string GetNormalizationCode(const GpuInfo& gpu_info, bool channels_x4,
+                                   bool two_step);
+  bool work_group_reduction_ = true;
+};
+
+// std dev can be calculated in single step, but two step algorithm can
+// provide more stable and robust results
+LayerNormalization CreateLayerNormalization(
+    const OperationDef& definition, const GpuInfo& gpu_info, const BHWC& shape,
+    float variance_bias, const Tensor<Linear, DataType::FLOAT32>& mul_linear,
+    const Tensor<Linear, DataType::FLOAT32>& sub_linear, bool two_step);
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_MEAN_STDDEV_NORMALIZATION_H_

tensorflow/lite/delegates/gpu/common/tasks/mean_stddev_normalization_test_util.cc

@@ -0,0 +1,220 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/mean_stddev_normalization_test_util.h"
+
+#include <cmath>
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/mean_stddev_normalization.h"
+
+namespace tflite {
+namespace gpu {
+
+// Parameterized test: mean, difference, tolerance.
+// Input is constructed as [mean-2*diff, mean-diff, mean+diff, mean+2*diff]
+absl::Status MeanStddevNormSeparateBatchesTest(float mean, float diff,
+                                               float tolerance,
+                                               TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 1, 2, 4);
+  src_tensor.data = {mean - 2 * diff, mean - diff,     mean + diff,
+                     mean + 2 * diff, mean - 2 * diff, mean - diff,
+                     mean + diff,     mean + 2 * diff};
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      auto operation = CreateMeanStdDevNormalization(op_def, env->GetGpuInfo(),
+                                                     src_tensor.shape);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {src_tensor},
+          std::make_unique<MeanStdDevNormalization>(std::move(operation)),
+          BHWC(1, 1, 2, 4), &dst_tensor));
+
+      std::vector<float> expected_output;
+      if (diff == 0.0f) {
+        expected_output.assign(
+            {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f});
+      } else {
+        const float ksqrt16 = std::sqrt(1.6f);
+        const float ksqrt04 = std::sqrt(0.4f);
+        expected_output.assign({-ksqrt16, -ksqrt04, ksqrt04, ksqrt16, -ksqrt16,
+                                -ksqrt04, ksqrt04, ksqrt16});
+      }
+      RETURN_IF_ERROR(
+          PointWiseNear(expected_output, dst_tensor.data, tolerance));
+
+      TensorFloat32 dst_tensor_single_step;
+      auto operation_single_step = CreateMeanStdDevNormalization(
+          op_def, env->GetGpuInfo(), src_tensor.shape,
+          /*variance_bias*/ 1.0e-8f, /*two_step*/ false);
+      RETURN_IF_ERROR(
+          env->ExecuteGPUOperation({src_tensor},
+                                   std::make_unique<MeanStdDevNormalization>(
+                                       std::move(operation_single_step)),
+                                   BHWC(1, 1, 2, 4), &dst_tensor_single_step));
+      RETURN_IF_ERROR(PointWiseNear(expected_output,
+                                    dst_tensor_single_step.data, tolerance));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status MeanStddevNormalizationAllBatchesTest(
+    TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(9, 1, 1, 4);
+  src_tensor.data = {
+      0.0f,    0.0f,    0.0f,   0.0f,    // zero mean, zero variance
+      -0.02f,  -0.01f,  0.01f,  0.02f,   // zero mean, small variance
+      -200.0f, -100.0f, 100.0f, 200.0f,  // zero mean, large variance
+      0.01f,   0.01f,   0.01f,  0.01f,   // small mean, zero variance
+      -0.01f,  0.0f,    0.02f,  0.03f,   // small mean, small variance
+      -199.0f, -99.0f,  101.0f, 201.0f,  // small mean, large variance
+      100.0f,  100.0f,  100.0f, 100.0f,  // large mean, zero variance
+      98.0f,   99.0f,   101.0f, 102.0f,  // large mean, small variance
+      -100.0f, 0.0f,    200.0f, 300.0f,  // large mean, large variance
+  };
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps =
+          precision == CalculationsPrecision::F32 ? 2.53e-05f : 3.57e-4f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::BHWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::BHWC});
+      TensorFloat32 dst_tensor;
+      auto operation = CreateMeanStdDevNormalization(op_def, env->GetGpuInfo(),
+                                                     src_tensor.shape);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {src_tensor},
+          std::make_unique<MeanStdDevNormalization>(std::move(operation)),
+          BHWC(9, 1, 1, 4), &dst_tensor));
+
+      const float ksqrt16 = std::sqrt(1.6f);
+      const float ksqrt04 = std::sqrt(0.4f);
+      const std::vector<float> expected_output = {
+          0.0f,     0.0f,     0.0f,    0.0f,     // zero mean, zero variance
+          -ksqrt16, -ksqrt04, ksqrt04, ksqrt16,  // zero mean, small variance
+          -ksqrt16, -ksqrt04, ksqrt04, ksqrt16,  // zero mean, large variance
+          0.0f,     0.0f,     0.0f,    0.0f,     // small mean, zero variance
+          -ksqrt16, -ksqrt04, ksqrt04, ksqrt16,  // small mean, small variance
+          -ksqrt16, -ksqrt04, ksqrt04, ksqrt16,  // small mean, large variance
+          0.0f,     0.0f,     0.0f,    0.0f,     // large mean, zero variance
+          -ksqrt16, -ksqrt04, ksqrt04, ksqrt16,  // large mean, small variance
+          -ksqrt16, -ksqrt04, ksqrt04, ksqrt16,  // large mean, large variance
+      };
+      RETURN_IF_ERROR(PointWiseNear(expected_output, dst_tensor.data, eps));
+
+      TensorFloat32 dst_tensor_single_step;
+      auto operation_single_step = CreateMeanStdDevNormalization(
+          op_def, env->GetGpuInfo(), src_tensor.shape,
+          /*variance_bias*/ 1.0e-8f, /*two_step*/ false);
+      RETURN_IF_ERROR(
+          env->ExecuteGPUOperation({src_tensor},
+                                   std::make_unique<MeanStdDevNormalization>(
+                                       std::move(operation_single_step)),
+                                   BHWC(9, 1, 1, 4), &dst_tensor_single_step));
+      RETURN_IF_ERROR(
+          PointWiseNear(expected_output, dst_tensor_single_step.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status MeanStddevNormalizationLargeVectorTest(
+    TestExecutionEnvironment* env) {
+  const float mean = 100.0f;
+  const float diff = 1.0f;
+  // Some large vector that is not a round multiple of any SIMD vector sizes.
+  constexpr int kVectorSize = 16 * 16 + 16 + 1;
+
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 1, 2, kVectorSize);
+  src_tensor.data.resize(kVectorSize * 2);
+  // First input is mean.
+  src_tensor.data[0] = mean;
+  src_tensor.data[kVectorSize] = mean;
+  // Rest is alternating between mean + diff and mean - diff.
+  for (int i = 1; i < kVectorSize - 1; i += 2) {
+    src_tensor.data[i + 0] = mean + diff;
+    src_tensor.data[i + 1] = mean - diff;
+    src_tensor.data[kVectorSize + i + 0] = mean + diff;
+    src_tensor.data[kVectorSize + i + 1] = mean - diff;
+  }
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps =
+          precision == CalculationsPrecision::F32 ? 5.0e-7f : 8.60e-4f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      auto operation = CreateMeanStdDevNormalization(op_def, env->GetGpuInfo(),
+                                                     src_tensor.shape);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {src_tensor},
+          std::make_unique<MeanStdDevNormalization>(std::move(operation)),
+          BHWC(1, 1, 2, kVectorSize), &dst_tensor));
+
+      std::vector<float> expected_output(kVectorSize * 2);
+      // First output should be 0.
+      expected_output[0] = 0.0;
+      expected_output[kVectorSize] = 0.0;
+      // Rest should be alternating between ±√(N/(N-1)).
+      const float expected_elem =
+          std::sqrt(static_cast<double>(kVectorSize) /
+                    static_cast<double>(kVectorSize - 1));
+      for (int i = 1; i < kVectorSize - 1; i += 2) {
+        expected_output[i + 0] = +expected_elem;
+        expected_output[i + 1] = -expected_elem;
+        expected_output[kVectorSize + i + 0] = +expected_elem;
+        expected_output[kVectorSize + i + 1] = -expected_elem;
+      }
+      RETURN_IF_ERROR(PointWiseNear(expected_output, dst_tensor.data, eps));
+
+      if (precision != CalculationsPrecision::F32) {
+        TensorFloat32 dst_tensor_single_step;
+        auto operation_single_step = CreateMeanStdDevNormalization(
+            op_def, env->GetGpuInfo(), src_tensor.shape,
+            /*variance_bias*/ 1.0e-8f, /*two_step*/ false);
+        RETURN_IF_ERROR(env->ExecuteGPUOperation(
+            {src_tensor},
+            std::make_unique<MeanStdDevNormalization>(
+                std::move(operation_single_step)),
+            BHWC(1, 1, 2, kVectorSize), &dst_tensor_single_step));
+        RETURN_IF_ERROR(
+            PointWiseNear(expected_output, dst_tensor_single_step.data, eps));
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/mean_stddev_normalization_test_util.h

@@ -0,0 +1,38 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_MEAN_STDDEV_NORMALIZATION_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_MEAN_STDDEV_NORMALIZATION_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status MeanStddevNormSeparateBatchesTest(float mean, float diff,
+                                               float tolerance,
+                                               TestExecutionEnvironment* env);
+
+absl::Status MeanStddevNormalizationAllBatchesTest(
+    TestExecutionEnvironment* env);
+
+absl::Status MeanStddevNormalizationLargeVectorTest(
+    TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_MEAN_STDDEV_NORMALIZATION_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/one_hot.cc

@@ -0,0 +1,76 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/delegates/gpu/common/tasks/one_hot.h"
+
+#include <string>
+#include <utility>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+
+namespace tflite {
+namespace gpu {
+
+std::string GetOneHotCode(const OperationDef& op_def,
+                          const OneHotAttributes& attr, GPUOperation* op) {
+  op->AddSrcTensor("src_tensor", op_def.src_tensors[0]);
+  op->AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
+  std::string c;
+  c += "MAIN_FUNCTION($0) {\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  c += "  int Y = GLOBAL_ID_1;\n";
+  c += "  int Z = GLOBAL_ID_2;\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "Z >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
+  c += "  int idx = Z * 4;\n";
+  c += "  int hot_idx = args.src_tensor.Read(0, 0, 0).x;\n";
+  c += "  FLT4 res = INIT_FLT4(args.off_value);\n";
+  c += "  if ((hot_idx >= idx) && (hot_idx < (idx + 4))) {\n";
+  c += "    res.x = (idx + 0) == hot_idx ? args.on_value : args.off_value;\n";
+  c += "    res.y = (idx + 1) == hot_idx ? args.on_value : args.off_value;\n";
+  c += "    res.z = (idx + 2) == hot_idx ? args.on_value : args.off_value;\n";
+  c += "    res.w = (idx + 3) == hot_idx ? args.on_value : args.off_value;\n";
+  c += "  }\n";
+  c += "  args.dst_tensor.Write(res, X, Y, Z);\n";
+  c += "}\n";
+  return c;
+}
+
+GPUOperation CreateOneHot(const OperationDef& definition,
+                          const OneHotAttributes& attr) {
+  GPUOperation op(definition);
+  op.code_ = GetOneHotCode(definition, attr, &op);
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+  if (definition.precision == CalculationsPrecision::F32) {
+    op.args_.AddFloat("on_value", attr.on_value);
+    op.args_.AddFloat("off_value", attr.off_value);
+  } else {
+    op.args_.AddHalf("on_value", half(attr.on_value));
+    op.args_.AddHalf("off_value", half(attr.off_value));
+  }
+  return op;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/one_hot.h

@@ -0,0 +1,35 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ONE_HOT_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ONE_HOT_H_
+
+#include <string>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+GPUOperation CreateOneHot(const OperationDef& definition,
+                          const OneHotAttributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ONE_HOT_H_

tensorflow/lite/delegates/gpu/common/tasks/one_hot_test_util.cc

@@ -0,0 +1,87 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/one_hot_test_util.h"
+
+#include <memory>
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/one_hot.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status OneHotTest(TestExecutionEnvironment* env) {
+  tflite::gpu::Tensor<BHWC, DataType::INT32> src_tensor;
+  src_tensor.shape = BHWC(1, 1, 1, 1);
+  src_tensor.data = {3};
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      OperationDef op_def;
+      op_def.src_tensors.push_back({DataType::INT32, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.precision = precision;
+      TensorDescriptor src = op_def.src_tensors[0];
+      TensorDescriptor dst = op_def.dst_tensors[0];
+      OneHotAttributes attr;
+      GPUOperation operation = CreateOneHot(op_def, attr);
+      src.UploadData(src_tensor);
+      dst.SetBHWCShape(BHWC(1, 1, 1, 8));
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {&src}, {&dst},
+          std::make_unique<GPUOperation>(std::move(operation))));
+      TensorFloat32 dst_tensor;
+      dst.DownloadData(&dst_tensor);
+      RETURN_IF_ERROR(
+          PointWiseNear({0, 0, 0, 1.0, 0, 0, 0, 0}, dst_tensor.data, 0.0f));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status OneHotBatchTest(TestExecutionEnvironment* env) {
+  tflite::gpu::Tensor<BHWC, DataType::INT32> src_tensor;
+  src_tensor.shape = BHWC(10, 1, 1, 1);
+  src_tensor.data = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      OperationDef op_def;
+      op_def.src_tensors.push_back({DataType::INT32, storage, Layout::BHWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::BHWC});
+      op_def.precision = precision;
+      TensorDescriptor src = op_def.src_tensors[0];
+      TensorDescriptor dst = op_def.dst_tensors[0];
+      OneHotAttributes attr = {/*on_value=*/2.0, /*off_value=*/-2.0};
+      GPUOperation operation = CreateOneHot(op_def, attr);
+      src.UploadData(src_tensor);
+      dst.SetBHWCShape(BHWC(10, 1, 1, 10));
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          {&src}, {&dst},
+          std::make_unique<GPUOperation>(std::move(operation))));
+      TensorFloat32 dst_tensor;
+      dst.DownloadData(&dst_tensor);
+      std::vector<float> expected(100, attr.off_value);
+      for (int i = 0; i < 10; i++) {
+        expected[i * 10 + i] = attr.on_value;
+      }
+      RETURN_IF_ERROR(PointWiseNear(expected, dst_tensor.data, 0.0f));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/one_hot_test_util.h

@@ -0,0 +1,32 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ONE_HOT_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ONE_HOT_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status OneHotTest(TestExecutionEnvironment* env);
+
+absl::Status OneHotBatchTest(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_ONE_HOT_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/padding.cc

@@ -0,0 +1,152 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/padding.h"
+
+#include <string>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
+
+namespace tflite {
+namespace gpu {
+
+namespace {
+std::string GetPaddingCode(const OperationDef& op_def,
+                           const PadAttributes& attr, GPUOperation* op) {
+  op->AddSrcTensor("src_tensor", op_def.src_tensors[0]);
+  op->AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
+  op->args_.AddInt("prepended_x", attr.prepended.w);
+  op->args_.AddInt("prepended_y", attr.prepended.h);
+  op->args_.AddInt("prepended_z", attr.prepended.c);
+  op->args_.AddInt("prepended_w", attr.prepended.b);
+
+  const std::string dst_batch =
+      op_def.dst_tensors[0].HasAxis(Axis::BATCH) ? "B" : "0";
+  std::string c;
+  const std::string channels[] = {".x", ".y", ".z", ".w"};
+
+  if (attr.type == PaddingContentType::REFLECT) {
+    c += "int reflect_coord(int x, int size) {\n";
+    c += "  int t = abs(x) - size + 1;\n";
+    c += "  return size - 1 - abs(t);\n";
+    c += "}\n\n";
+  }
+
+  c += "MAIN_FUNCTION($0) {\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  c += "  int Y = GLOBAL_ID_1;\n";
+  c += "  int Z = GLOBAL_ID_2;\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "Z >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
+  c += "  args.src_tensor::type result = (" +
+       ToCLDataType(op_def.src_tensors[0].GetDataType(), /*vec_size*/ 4) +
+       ")(" + std::to_string(attr.constant_values) + ");\n";
+  c += "  int s_x = X - args.prepended_x;\n";
+  c += "  int s_y = Y - args.prepended_y;\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int s_b = " + dst_batch + " - args.prepended_w;\n";
+    c += "  args.src_tensor.SetBatchRef(s_b);\n";
+  }
+  if (attr.type == PaddingContentType::REFLECT) {
+    c += "  s_x = reflect_coord(s_x, args.src_tensor.Width());\n";
+    c += "  s_y = reflect_coord(s_y, args.src_tensor.Height());\n";
+    if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+      c += "  s_b = reflect_coord(s_b, args.src_tensor.Batch());\n";
+    }
+    if (attr.prepended.c == 0 && attr.appended.c == 0) {
+      // optimized case
+      c += "  result = args.src_tensor.Read(s_x, s_y, Z);\n";
+    } else {
+      c += "  int start_channel = Z * 4;\n";
+      for (int i = 0; i < 4; ++i) {
+        const auto& s = channels[i];
+        c += "  {\n";
+        c += "    int channel = start_channel + " + std::to_string(i) + ";\n";
+        c += "    int s_z = channel - args.prepended_z;\n";
+        // We need additional clamp for z, so that we use alignment for channels
+        // and can proceed extra channels that can lead to reading out of
+        // resource.
+        c += "    s_z = clamp(reflect_coord(s_z, args.src_tensor.Channels()), "
+             "0, "
+             "args.src_tensor.Channels() - "
+             "1);\n";
+        c += "    args.src_tensor.ReadPerChannel(result" + s +
+             ", s_x, s_y, s_z);\n";
+        c += "  }\n";
+      }
+    }
+  } else {
+    c += "  bool inside_x = s_x >= 0 && s_x < args.src_tensor.Width();\n";
+    c += "  bool inside_y = s_y >= 0 && s_y < args.src_tensor.Height();\n";
+    if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+      c += "  inside_y = inside_y && (s_b >= 0 && s_b < "
+           "args.src_tensor.Batch());\n";
+    }
+    c += "  if (inside_x && inside_y) {\n";
+    if (attr.prepended.c == 0 && attr.appended.c == 0) {
+      // optimized case
+      c += "    result = args.src_tensor.Read(s_x, s_y, Z);\n";
+    } else if (attr.prepended.c % 4 == 0 &&
+               attr.prepended.b == attr.appended.b) {
+      c += "    int s_z = Z - args.prepended_z / 4;\n";
+      c += "    if (s_z >= 0 && s_z < args.src_tensor.Slices()) {\n";
+      c += "      result = args.src_tensor.Read(s_x, s_y, s_z);\n";
+      c += "    }\n";
+    } else {
+      c += "    int start_channel = Z * 4;\n";
+      for (int i = 0; i < 4; ++i) {
+        const auto& s = channels[i];
+        c += "    {\n";
+        c += "    int channel = start_channel + " + std::to_string(i) + ";\n";
+        c += "    int s_z = channel - args.prepended_z;\n";
+        c += "    if (s_z >= 0 && s_z < args.src_tensor.Channels()) {\n";
+        c += "      args.src_tensor.ReadPerChannel(result" + s +
+             ", s_x, s_y, s_z);\n";
+        c += "    }\n";
+        c += "    }\n";
+      }
+    }
+    c += "  }\n";
+  }
+  c += "  args.dst_tensor.Write(result, X, Y, Z);\n";
+  c += "}\n";
+
+  return c;
+}
+
+}  // namespace
+
+GPUOperation CreatePadding(const OperationDef& definition,
+                           const PadAttributes& attr) {
+  GPUOperation op(definition);
+  op.code_ = GetPaddingCode(definition, attr, &op);
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+  return op;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/padding.h

@@ -0,0 +1,33 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_PADDING_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_PADDING_H_
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+GPUOperation CreatePadding(const OperationDef& definition,
+                           const PadAttributes& attr);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_PADDING_H_

tensorflow/lite/delegates/gpu/common/tasks/padding_test_util.cc

@@ -0,0 +1,358 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/padding_test_util.h"
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/operations.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+#include "tensorflow/lite/delegates/gpu/common/tasks/padding.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status PaddingAppendWidthTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 1, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f};
+
+  PadAttributes attr;
+  attr.prepended = BHWC(0, 0, 0, 0);
+  attr.appended = BHWC(0, 0, 1, 0);
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreatePadding(op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 2, 2, 2), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({0.0f, 1.0f, 0.0f, 0.0f, 2.0f, 3.0f, 0.0f, 0.0f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status PaddingAppendWidthConstValuesTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 1, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f};
+
+  PadAttributes attr;
+  attr.prepended = BHWC(0, 0, 0, 0);
+  attr.appended = BHWC(0, 0, 1, 0);
+  attr.constant_values = 5;
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreatePadding(op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 2, 2, 2), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({0.0f, 1.0f, 5.0f, 5.0f, 2.0f, 3.0f, 5.0f, 5.0f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status PaddingPrependWidthTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 1, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f};
+
+  PadAttributes attr;
+  attr.prepended = BHWC(0, 0, 1, 0);
+  attr.appended = BHWC(0, 0, 0, 0);
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreatePadding(op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 2, 2, 2), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f, 2.0f, 3.0f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status PaddingAppendHeightTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 1, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f};
+
+  PadAttributes attr;
+  attr.prepended = BHWC(0, 0, 0, 0);
+  attr.appended = BHWC(0, 1, 0, 0);
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreatePadding(op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 3, 1, 2), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear({0.0f, 1.0f, 2.0f, 3.0f, 0.0f, 0.0f},
+                                    dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status PaddingPrependHeightTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 1, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f};
+
+  PadAttributes attr;
+  attr.prepended = BHWC(0, 1, 0, 0);
+  attr.appended = BHWC(0, 0, 0, 0);
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreatePadding(op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 3, 1, 2), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear({0.0f, 0.0f, 0.0f, 1.0f, 2.0f, 3.0f},
+                                    dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status PaddingAppendChannelsTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 1, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f};
+
+  PadAttributes attr;
+  attr.prepended = BHWC(0, 0, 0, 0);
+  attr.appended = BHWC(0, 0, 0, 1);
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreatePadding(op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 2, 1, 3), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear({0.0f, 1.0f, 0.0f, 2.0f, 3.0f, 0.0f},
+                                    dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status PaddingPrependChannelsTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 1, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f};
+
+  PadAttributes attr;
+  attr.prepended = BHWC(0, 0, 0, 1);
+  attr.appended = BHWC(0, 0, 0, 0);
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreatePadding(op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 2, 1, 3), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear({0.0f, 0.0f, 1.0f, 0.0f, 2.0f, 3.0f},
+                                    dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status PaddingPrependChannelsX4Test(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 1, 1, 2);
+  src_tensor.data = {1.0f, 2.0f};
+
+  PadAttributes attr;
+  attr.prepended = BHWC(0, 0, 0, 4);
+  attr.appended = BHWC(0, 0, 0, 0);
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreatePadding(op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 1, 1, 6), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear({0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 2.0f},
+                                    dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status PaddingComplexTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 2, 1, 2);
+  src_tensor.data = {0.0f, 1.0f, 2.0f, 3.0f};
+
+  PadAttributes attr;
+  attr.prepended = BHWC(0, 0, 1, 1);
+  attr.appended = BHWC(0, 1, 1, 0);
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreatePadding(op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 3, 3, 3), &dst_tensor));
+      RETURN_IF_ERROR(
+          PointWiseNear({0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f, 0.0f,
+                         0.0f, 0.0f, 0.0f, 0.0f, 2.0f, 3.0f, 0.0f, 0.0f, 0.0f,
+                         0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f},
+                        dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status PaddingReflectWidthTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 1, 3, 1);
+  src_tensor.data = {1.0f, 2.0f, 3.0f};
+
+  PadAttributes attr;
+  attr.prepended = BHWC(0, 0, 2, 0);
+  attr.appended = BHWC(0, 0, 2, 0);
+  attr.type = PaddingContentType::REFLECT;
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreatePadding(op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 1, 7, 1), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear({3.0f, 2.0f, 1.0f, 2.0f, 3.0f, 2.0f, 1.0f},
+                                    dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status PaddingReflectChannelsTest(TestExecutionEnvironment* env) {
+  TensorFloat32 src_tensor;
+  src_tensor.shape = BHWC(1, 1, 1, 3);
+  src_tensor.data = {1.0f, 2.0f, 3.0f};
+
+  PadAttributes attr;
+  attr.prepended = BHWC(0, 0, 0, 2);
+  attr.appended = BHWC(0, 0, 0, 2);
+  attr.type = PaddingContentType::REFLECT;
+
+  for (auto precision : env->GetSupportedPrecisions()) {
+    auto data_type = DeduceDataTypeFromPrecision(precision);
+    for (auto storage : env->GetSupportedStorages(data_type)) {
+      const float eps = precision == CalculationsPrecision::F32 ? 1e-6f : 1e-3f;
+      OperationDef op_def;
+      op_def.precision = precision;
+      op_def.src_tensors.push_back({data_type, storage, Layout::HWC});
+      op_def.dst_tensors.push_back({data_type, storage, Layout::HWC});
+      TensorFloat32 dst_tensor;
+      GPUOperation operation = CreatePadding(op_def, attr);
+      RETURN_IF_ERROR(env->ExecuteGPUOperation(
+          src_tensor, std::make_unique<GPUOperation>(std::move(operation)),
+          BHWC(1, 1, 1, 7), &dst_tensor));
+      RETURN_IF_ERROR(PointWiseNear({3.0f, 2.0f, 1.0f, 2.0f, 3.0f, 2.0f, 1.0f},
+                                    dst_tensor.data, eps));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/tasks/padding_test_util.h

@@ -0,0 +1,40 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_PADDING_TEST_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_PADDING_TEST_UTIL_H_
+
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/testing_util.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status PaddingAppendWidthTest(TestExecutionEnvironment* env);
+absl::Status PaddingAppendWidthConstValuesTest(TestExecutionEnvironment* env);
+absl::Status PaddingPrependWidthTest(TestExecutionEnvironment* env);
+absl::Status PaddingAppendHeightTest(TestExecutionEnvironment* env);
+absl::Status PaddingPrependHeightTest(TestExecutionEnvironment* env);
+absl::Status PaddingAppendChannelsTest(TestExecutionEnvironment* env);
+absl::Status PaddingPrependChannelsTest(TestExecutionEnvironment* env);
+absl::Status PaddingPrependChannelsX4Test(TestExecutionEnvironment* env);
+absl::Status PaddingComplexTest(TestExecutionEnvironment* env);
+absl::Status PaddingReflectWidthTest(TestExecutionEnvironment* env);
+absl::Status PaddingReflectChannelsTest(TestExecutionEnvironment* env);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASKS_PADDING_TEST_UTIL_H_

tensorflow/lite/delegates/gpu/common/tasks/pooling.cc

@@ -0,0 +1,294 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/delegates/gpu/common/tasks/pooling.h"
+
+#include <map>
+#include <string>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/gpu_info.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tensor_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/util.h"
+
+namespace tflite {
+namespace gpu {
+
+namespace {
+std::string GetAveragePoolingKernelCode(const OperationDef& op_def,
+                                        const GpuInfo& gpu_info,
+                                        GPUOperation* op) {
+  op->AddSrcTensor("src_tensor", op_def.src_tensors[0]);
+  op->AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
+
+  std::map<Axis, std::string> axis_to_src_coord = {
+      {Axis::WIDTH, "x_c"},  {Axis::HEIGHT, "y_c"}, {Axis::DEPTH, "d_c"},
+      {Axis::CHANNELS, "Z"}, {Axis::BATCH, "B"},
+  };
+
+  std::map<Axis, std::string> axis_to_dst_coord = {
+      {Axis::WIDTH, "X"},    {Axis::HEIGHT, "Y"}, {Axis::DEPTH, "D"},
+      {Axis::CHANNELS, "Z"}, {Axis::BATCH, "B"},
+  };
+
+  std::vector<std::string> src_coords;
+  std::vector<std::string> dst_coords;
+  for (auto axis : {Axis::WIDTH, Axis::HEIGHT, Axis::DEPTH, Axis::CHANNELS}) {
+    if (op_def.dst_tensors[0].HasAxis(axis)) {
+      dst_coords.push_back(axis_to_dst_coord[axis]);
+    }
+    if (op_def.src_tensors[0].HasAxis(axis)) {
+      src_coords.push_back(axis_to_src_coord[axis]);
+    }
+  }
+  std::string src_coord = src_coords[0];
+  for (int i = 1; i < src_coords.size(); ++i) {
+    src_coord += ", " + src_coords[i];
+  }
+  std::string dst_coord = dst_coords[0];
+  for (int i = 1; i < dst_coords.size(); ++i) {
+    dst_coord += ", " + dst_coords[i];
+  }
+
+  std::string c;
+  c += "MAIN_FUNCTION($0) {\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  int linear_id_1 = GLOBAL_ID_1;\n";
+    c += "  int Y = linear_id_1 / args.dst_tensor.Depth();\n";
+    c += "  int D = linear_id_1 % args.dst_tensor.Depth();\n";
+  } else {
+    c += "  int Y = GLOBAL_ID_1;\n";
+  }
+  c += "  int Z = GLOBAL_ID_2;\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "Z >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
+  c += "  float4 r = INIT_FLOAT4(0.0f);\n";
+  c += "  float window_size = 0.0;\n";
+  c += "  int xs = X * args.stride_x + args.padding_x;\n";
+  c += "  int ys = Y * args.stride_y + args.padding_y;\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  int ds = D * args.stride_z + args.padding_z;\n";
+    c += "  for (int kz = 0; kz < args.kernel_size_z; ++kz) {\n";
+    c += "    int d_c = ds + kz;\n";
+    c += "    if (d_c < 0 || d_c >= args.src_tensor.Depth()) continue;\n";
+  }
+  c += "  for (int ky = 0; ky < args.kernel_size_y; ++ky) {\n";
+  c += "    int y_c = ys + ky;\n";
+  c += "    bool outside_y = y_c < 0 || y_c >= args.src_tensor.Height();\n";
+  c += "    for (int kx = 0; kx < args.kernel_size_x; ++kx) {\n";
+  c += "      int x_c = xs + kx;\n";
+  c += "      bool outside = outside_y || x_c < 0 || x_c >= "
+       "args.src_tensor.Width();\n";
+  if (op_def.src_tensors[0].SupportsZeroClamp(Axis::WIDTH, gpu_info) &&
+      op_def.src_tensors[0].SupportsZeroClamp(Axis::HEIGHT, gpu_info)) {
+    c += "      r += args.src_tensor.Read<float>(" + src_coord + ");\n";
+  } else {
+    c += "     r += !outside ? args.src_tensor.Read<float>(" + src_coord +
+         ") : "
+         "INIT_FLOAT4(0.0f);\n";
+  }
+  c += "        window_size += !outside ? 1.0 : 0.0;\n";
+  c += "    }\n";
+  c += "  }\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  }  // Depth\n";
+  }
+  // If window_size==0, window covered nothing. This situation is a sign of
+  // incorrectly constructed operation. NaNs are expected as output.
+  c += "  FLT4 result = TO_FLT4(r / window_size);\n";
+  c += "  args.dst_tensor.Write(result, " + dst_coord + ");\n";
+  c += "}\n";
+
+  return c;
+}
+
+std::string GetMaxPoolingKernelCode(const OperationDef& op_def,
+                                    bool output_indices, GPUOperation* op) {
+  op->AddSrcTensor("src_tensor", op_def.src_tensors[0]);
+  op->AddDstTensor("dst_tensor", op_def.dst_tensors[0]);
+  if (output_indices) {
+    op->AddDstTensor("dst_indices", op_def.dst_tensors[1]);
+  }
+
+  std::map<Axis, std::string> axis_to_src_coord = {
+      {Axis::WIDTH, "x_c"},  {Axis::HEIGHT, "y_c"}, {Axis::DEPTH, "d_c"},
+      {Axis::CHANNELS, "Z"}, {Axis::BATCH, "B"},
+  };
+
+  std::map<Axis, std::string> axis_to_dst_coord = {
+      {Axis::WIDTH, "X"},    {Axis::HEIGHT, "Y"}, {Axis::DEPTH, "D"},
+      {Axis::CHANNELS, "Z"}, {Axis::BATCH, "B"},
+  };
+
+  std::vector<std::string> src_coords;
+  std::vector<std::string> dst_coords;
+  for (auto axis : {Axis::WIDTH, Axis::HEIGHT, Axis::DEPTH, Axis::CHANNELS}) {
+    if (op_def.dst_tensors[0].HasAxis(axis)) {
+      dst_coords.push_back(axis_to_dst_coord[axis]);
+    }
+    if (op_def.src_tensors[0].HasAxis(axis)) {
+      src_coords.push_back(axis_to_src_coord[axis]);
+    }
+  }
+  std::string src_coord = src_coords[0];
+  for (int i = 1; i < src_coords.size(); ++i) {
+    src_coord += ", " + src_coords[i];
+  }
+  std::string dst_coord = dst_coords[0];
+  for (int i = 1; i < dst_coords.size(); ++i) {
+    dst_coord += ", " + dst_coords[i];
+  }
+
+  std::string c;
+  c += "MAIN_FUNCTION($0) {\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::BATCH)) {
+    c += "  int linear_id = GLOBAL_ID_0;\n";
+    c += "  int X = linear_id / args.dst_tensor.Batch();\n";
+    c += "  int B = linear_id % args.dst_tensor.Batch();\n";
+    c += "  args.src_tensor.SetBatchRef(B);\n";
+    c += "  args.dst_tensor.SetBatchRef(B);\n";
+    if (output_indices) {
+      c += "  args.dst_indices.SetBatchRef(B);\n";
+    }
+  } else {
+    c += "  int X = GLOBAL_ID_0;\n";
+  }
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "  int linear_id_1 = GLOBAL_ID_1;\n";
+    c += "  int Y = linear_id_1 / args.dst_tensor.Depth();\n";
+    c += "  int D = linear_id_1 % args.dst_tensor.Depth();\n";
+  } else {
+    c += "  int Y = GLOBAL_ID_1;\n";
+  }
+  c += "  int Z = GLOBAL_ID_2;\n";
+  c += "  if (X >= args.dst_tensor.Width() || Y >= args.dst_tensor.Height() || "
+       "Z >= args.dst_tensor.Slices()) { \n";
+  c += "    return; \n";
+  c += "  } \n";
+  c += "  FLT4 maximum = INIT_FLT4(-10000.0f);\n";
+  if (output_indices) {
+    c += "  args.dst_indices::type indexes = args.dst_indices::zero_value;\n";
+  }
+  c += "  int xs = X * args.stride_x + args.padding_x;\n";
+  c += "  int ys = Y * args.stride_y + args.padding_y;\n";
+  c += "  for (int ky = 0; ky < args.kernel_size_y; ++ky) {\n";
+  c += "    int y_c = ys + ky;\n";
+  c += "    if (y_c < 0 || y_c >= args.src_tensor.Height()) continue;\n";
+  c += "    for (int kx = 0; kx < args.kernel_size_x; ++kx) {\n";
+  c += "      int x_c = xs + kx;\n";
+  c += "      if (x_c < 0 || x_c >= args.src_tensor.Width()) continue;\n";
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "    int ds = D * args.stride_z + args.padding_z;\n";
+    c += "    for (int kz = 0; kz < args.kernel_size_z; ++kz) {\n";
+    c += "    int d_c = ds + kz;\n";
+    c += "      if (d_c < 0 || d_c >= args.src_tensor.Depth()) continue;\n";
+  }
+  c += "      FLT4 src = args.src_tensor.Read(" + src_coord + ");\n";
+  if (output_indices) {
+    if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+      c += "      int index_counter = (ky * args.kernel_size_x + kx) * "
+           "args.kernel_size_z + kz;\n";
+    } else {
+      c += "      int index_counter = ky * args.kernel_size_x + kx;\n";
+    }
+    c += "      if (src.x > maximum.x) {\n";
+    c += "        indexes.x = index_counter;\n";
+    c += "        maximum.x = src.x;\n";
+    c += "      }\n";
+    c += "      if (src.y > maximum.y) {\n";
+    c += "        indexes.y = index_counter;\n";
+    c += "        maximum.y = src.y;\n";
+    c += "      }\n";
+    c += "      if (src.z > maximum.z) {\n";
+    c += "        indexes.z = index_counter;\n";
+    c += "        maximum.z = src.z;\n";
+    c += "      }\n";
+    c += "      if (src.w > maximum.w) {\n";
+    c += "        indexes.w = index_counter;\n";
+    c += "        maximum.w = src.w;\n";
+    c += "      }\n";
+  } else {
+    c += "      maximum = max(src, maximum);\n";
+  }
+  if (op_def.dst_tensors[0].HasAxis(Axis::DEPTH)) {
+    c += "    }  // Depth\n";
+  }
+  c += "    }\n";
+  c += "  }\n";
+  c += "  args.dst_tensor.Write(maximum, " + dst_coord + ");\n";
+  if (output_indices) {
+    c += "  args.dst_indices.Write(indexes, " + dst_coord + ");\n";
+  }
+  c += "}\n";
+
+  return c;
+}
+}  // namespace
+
+GPUOperation CreatePooling(const OperationDef& definition,
+                           const GpuInfo& gpu_info,
+                           const Pooling2DAttributes& attr) {
+  GPUOperation op(definition);
+  op.args_.AddInt("kernel_size_x", attr.kernel.w);
+  op.args_.AddInt("padding_x", -attr.padding.prepended.w);
+  op.args_.AddInt("stride_x", attr.strides.w);
+  op.args_.AddInt("kernel_size_y", attr.kernel.h);
+  op.args_.AddInt("padding_y", -attr.padding.prepended.h);
+  op.args_.AddInt("stride_y", attr.strides.h);
+  if (attr.type == PoolingType::AVERAGE) {
+    op.code_ = GetAveragePoolingKernelCode(definition, gpu_info, &op);
+  } else if (attr.type == PoolingType::MAX) {
+    op.code_ = GetMaxPoolingKernelCode(definition, attr.output_indices, &op);
+  }
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+  return op;
+}
+
+GPUOperation CreatePooling(const OperationDef& definition,
+                           const GpuInfo& gpu_info,
+                           const Pooling3DAttributes& attr) {
+  GPUOperation op(definition);
+  op.args_.AddInt("kernel_size_x", attr.kernel.w);
+  op.args_.AddInt("padding_x", -attr.padding.prepended.w);
+  op.args_.AddInt("stride_x", attr.strides.w);
+  op.args_.AddInt("kernel_size_y", attr.kernel.h);
+  op.args_.AddInt("padding_y", -attr.padding.prepended.h);
+  op.args_.AddInt("stride_y", attr.strides.h);
+  op.args_.AddInt("kernel_size_z", attr.kernel.d);
+  op.args_.AddInt("padding_z", -attr.padding.prepended.d);
+  op.args_.AddInt("stride_z", attr.strides.d);
+  if (attr.type == PoolingType::AVERAGE) {
+    op.code_ = GetAveragePoolingKernelCode(definition, gpu_info, &op);
+  } else if (attr.type == PoolingType::MAX) {
+    op.code_ = GetMaxPoolingKernelCode(definition, attr.output_indices, &op);
+  }
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+  return op;
+}
+
+}  // namespace gpu
+}  // namespace tflite