chore: import upstream snapshot with attribution

tensorflow/lite/delegates/gpu/common/task/BUILD

@@ -0,0 +1,255 @@
+load("@flatbuffers//:build_defs.bzl", "flatbuffer_cc_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "arguments",
+    srcs = ["arguments.cc"],
+    hdrs = ["arguments.h"],
+    deps = [
+        ":buffer_desc",
+        ":gpu_object_desc",
+        ":serialization_base_cc_fbs",
+        ":util",
+        "//tensorflow/lite/delegates/gpu/common:access_type",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common:types",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+    ],
+)
+
+cc_library(
+    name = "buffer_desc",
+    srcs = ["buffer_desc.cc"],
+    hdrs = ["buffer_desc.h"],
+    deps = [
+        ":gpu_object_desc",
+        ":util",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "compiler_options",
+    hdrs = ["compiler_options.h"],
+)
+
+cc_library(
+    name = "gpu_object_desc",
+    hdrs = ["gpu_object_desc.h"],
+    deps = [
+        ":serialization_base_cc_fbs",
+        "//tensorflow/lite/delegates/gpu/common:access_type",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:gpu_info",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "gpu_operation",
+    srcs = ["gpu_operation.cc"],
+    hdrs = ["gpu_operation.h"],
+    deps = [
+        ":arguments",
+        ":buffer_desc",
+        ":compiler_options",
+        ":gpu_tensor",
+        ":serialization_base_cc_fbs",
+        ":tensor_desc",
+        ":tuning_type",
+        ":util",
+        ":work_group_picking",
+        "//tensorflow/lite/delegates/gpu/common:access_type",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:gpu_info",
+        "//tensorflow/lite/delegates/gpu/common:kernel_info",
+        "//tensorflow/lite/delegates/gpu/common:precision",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common:types",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
+cc_library(
+    name = "gpu_tensor",
+    hdrs = ["gpu_tensor.h"],
+    deps = [
+        "//tensorflow/lite/delegates/gpu/common/task:tensor_desc",
+    ],
+)
+
+cc_library(
+    name = "profiling_info",
+    srcs = ["profiling_info.cc"],
+    hdrs = ["profiling_info.h"],
+    deps = [
+        "@com_google_absl//absl/time",
+    ],
+)
+
+cc_library(
+    name = "qcom_thin_filter_desc",
+    srcs = ["qcom_thin_filter_desc.cc"],
+    hdrs = ["qcom_thin_filter_desc.h"],
+    deps = [
+        ":gpu_object_desc",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "serialization_base",
+    srcs = ["serialization_base.cc"],
+    hdrs = ["serialization_base.h"],
+    deps = [
+        ":arguments",
+        ":buffer_desc",
+        ":gpu_object_desc",
+        ":gpu_operation",
+        ":serialization_base_cc_fbs",
+        ":tensor_desc",
+        "//tensorflow/lite/delegates/gpu/common:precision",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common:types",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+    ],
+)
+
+flatbuffer_cc_library(
+    name = "serialization_base_cc_fbs",
+    srcs = ["tflite_serialization_base.fbs"],
+    flatc_args = [
+        "--scoped-enums",
+    ],
+)
+
+cc_library(
+    name = "tensor_desc",
+    srcs = ["tensor_desc.cc"],
+    hdrs = ["tensor_desc.h"],
+    deps = [
+        ":gpu_object_desc",
+        ":util",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:shape",
+        "//tensorflow/lite/delegates/gpu/common:tensor",
+        "//tensorflow/lite/delegates/gpu/common:types",
+        "//tensorflow/lite/delegates/gpu/common:util",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "testing_util",
+    testonly = 1,
+    srcs = ["testing_util.cc"],
+    hdrs = ["testing_util.h"],
+    deps = [
+        ":gpu_operation",
+        ":tensor_desc",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:gpu_info",
+        "//tensorflow/lite/delegates/gpu/common:gpu_model",
+        "//tensorflow/lite/delegates/gpu/common:precision",
+        "//tensorflow/lite/delegates/gpu/common:shape",
+        "//tensorflow/lite/delegates/gpu/common:tensor",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "tuning_type",
+    hdrs = ["tuning_type.h"],
+)
+
+cc_library(
+    name = "util",
+    srcs = ["util.cc"],
+    hdrs = ["util.h"],
+    deps = [
+        ":gpu_object_desc",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:gpu_info",
+        "//tensorflow/lite/delegates/gpu/common:precision",
+        "//tensorflow/lite/delegates/gpu/common:types",
+        "//tensorflow/lite/delegates/gpu/common:util",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
+cc_library(
+    name = "weights_conversion",
+    srcs = ["weights_conversion.cc"],
+    hdrs = ["weights_conversion.h"],
+    deps = [
+        ":weights_layout",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:shape",
+        "//tensorflow/lite/delegates/gpu/common:status",
+        "//tensorflow/lite/delegates/gpu/common:tensor",
+        "//tensorflow/lite/delegates/gpu/common:types",
+        "//tensorflow/lite/delegates/gpu/common:util",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_library(
+    name = "weights_layout",
+    srcs = ["weights_layout.cc"],
+    hdrs = ["weights_layout.h"],
+    deps = [
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+    ],
+)
+
+cc_library(
+    name = "work_group_picking",
+    srcs = ["work_group_picking.cc"],
+    hdrs = ["work_group_picking.h"],
+    deps = [
+        ":tuning_type",
+        "//tensorflow/lite/delegates/gpu/common:gpu_info",
+        "//tensorflow/lite/delegates/gpu/common:kernel_info",
+        "//tensorflow/lite/delegates/gpu/common:types",
+        "//tensorflow/lite/delegates/gpu/common:util",
+        "//tensorflow/lite/delegates/gpu/common:workgroup_selection",
+    ],
+)
+
+cc_test(
+    name = "weights_conversion_test",
+    srcs = ["weights_conversion_test.cc"],
+    deps = [
+        ":weights_conversion",
+        ":weights_layout",
+        "//tensorflow/lite/delegates/gpu/common:data_type",
+        "//tensorflow/lite/delegates/gpu/common:shape",
+        "//tensorflow/lite/delegates/gpu/common:tensor",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tensorflow/lite/delegates/gpu/common/task/arguments.cc

@@ -0,0 +1,388 @@
+/* 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/task/arguments.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <cstring>
+#include <map>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/strings/ascii.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
+#include "tensorflow/lite/delegates/gpu/common/data_type.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_object_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/util.h"
+
+namespace tflite {
+namespace gpu {
+namespace {
+bool IsWordSymbol(char symbol) {
+  return absl::ascii_isalnum(symbol) || symbol == '_';
+}
+
+bool HasWord(const std::string& word, const std::string& text) {
+  size_t pos = text.find(word);
+  while (pos != std::string::npos) {
+    char prev = pos == 0 ? '.' : text[pos - 1];
+    char next = pos + word.size() < text.size() ? text[pos + word.size()] : '.';
+    if (!IsWordSymbol(prev) && !IsWordSymbol(next)) {
+      return true;
+    }
+    pos = text.find(word, pos + 1);
+  }
+  return false;
+}
+
+std::string RenameArg(const std::vector<std::string>& object_names,
+                      const std::string& postfix, const std::string& arg_name) {
+  for (const auto& object_name : object_names) {
+    if (absl::StartsWith(arg_name, object_name) &&
+        arg_name.size() > object_name.size() &&
+        arg_name[object_name.size()] == '_') {
+      return object_name + postfix +
+             arg_name.substr(object_name.size(),
+                             arg_name.size() - object_name.size());
+    }
+  }
+  return arg_name + postfix;
+}
+
+absl::Status BufferToKernelLanguage(const GpuInfo& gpu_info,
+                                    const std::string& buffer_name,
+                                    const BufferDescriptor* buffer_desc,
+                                    std::string* result) {
+  if (buffer_desc->element_size != 1) {
+    return absl::UnimplementedError("No support of vector types.");
+  }
+  const int elements_count =
+      buffer_desc->size /
+      (buffer_desc->element_size * SizeOf(buffer_desc->element_type));
+  if (gpu_info.IsGlsl()) {
+    const std::string glsl_type = ToGlslShaderDataType(
+        buffer_desc->element_type, buffer_desc->element_size,
+        /*add_precision*/ false, gpu_info.IsGlslSupportsExplicitFp16());
+    const std::string glsl_type_with_precision = ToGlslShaderDataType(
+        buffer_desc->element_type, buffer_desc->element_size,
+        /*add_precision*/ true, gpu_info.IsGlslSupportsExplicitFp16());
+    *result = "const " + glsl_type_with_precision + " " + buffer_name +
+              "_buffer[] = " + glsl_type + "[](\n";
+  } else if (gpu_info.IsApiMetal()) {
+    const std::string metal_type =
+        ToMetalDataType(buffer_desc->element_type, buffer_desc->element_size);
+    *result = "constant " + metal_type + " " + buffer_name + "_buffer[" +
+              std::to_string(elements_count) + "] = {\n";
+  } else if (gpu_info.IsApiOpenCl()) {
+    const std::string cl_type =
+        ToCLDataType(buffer_desc->element_type, buffer_desc->element_size);
+    *result = "__constant " + cl_type + " " + buffer_name + "_buffer[" +
+              std::to_string(elements_count) + "] = {\n";
+  } else {
+    return absl::UnimplementedError("Not supported API.");
+  }
+  if (buffer_desc->element_type == DataType::FLOAT16) {
+    std::string postfix = "f";
+    if (gpu_info.IsGlsl() && gpu_info.IsGlslSupportsExplicitFp16()) {
+      postfix = "hf";
+    }
+    const half* data_ptr =
+        reinterpret_cast<const half*>(buffer_desc->data.data());
+    for (int i = 0; i < elements_count; ++i) {
+      *result += "  " +
+                 absl::StrFormat("%.10f", static_cast<float>(data_ptr[i])) +
+                 postfix;
+      if (i != elements_count - 1) {
+        *result += ",\n";
+      }
+    }
+  } else if (buffer_desc->element_type == DataType::FLOAT32) {
+    const float* data_ptr =
+        reinterpret_cast<const float*>(buffer_desc->data.data());
+    for (int i = 0; i < elements_count; ++i) {
+      *result += "  " + absl::StrFormat("%.10f", data_ptr[i]) + "f";
+      if (i != elements_count - 1) {
+        *result += ",\n";
+      }
+    }
+  } else {
+    return absl::UnimplementedError("Not supported type.");
+  }
+  if (gpu_info.IsGlsl()) {
+    *result += ");\n";
+  } else {
+    *result += "};\n";
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+// Static
+constexpr char Arguments::kArgsPrefix[];
+
+void Arguments::AddFloat(const std::string& name, float value) {
+  float_values_[name].value = value;
+}
+void Arguments::AddHalf(const std::string& name, half value) {
+  half_values_[name].value = value;
+}
+void Arguments::AddInt(const std::string& name, int value) {
+  int_values_[name].value = value;
+}
+
+absl::Status Arguments::SetInt(const std::string& name, int value) {
+  auto it = int_values_.find(name);
+  if (it == int_values_.end()) {
+    return absl::NotFoundError(
+        absl::StrCat("No int argument with name - ", name));
+  }
+  it->second.value = value;
+  return absl::OkStatus();
+}
+absl::Status Arguments::SetFloat(const std::string& name, float value) {
+  auto it = float_values_.find(name);
+  if (it == float_values_.end()) {
+    return absl::NotFoundError(
+        absl::StrCat("No float argument with name - ", name));
+  }
+  it->second.value = value;
+  return absl::OkStatus();
+}
+
+absl::Status Arguments::SetHalf(const std::string& name, half value) {
+  auto it = half_values_.find(name);
+  if (it == half_values_.end()) {
+    return absl::NotFoundError(
+        absl::StrCat("No half argument with name - ", name));
+  }
+  it->second.value = value;
+  return absl::OkStatus();
+}
+
+void Arguments::AddObjectRef(const std::string& name, AccessType access_type,
+                             GPUObjectDescriptorPtr&& descriptor_ptr) {
+  descriptor_ptr->SetAccess(access_type);
+  object_refs_[name] = {std::move(descriptor_ptr)};
+}
+
+void Arguments::AddObject(const std::string& name,
+                          GPUObjectDescriptorPtr&& descriptor_ptr) {
+  descriptor_ptr->SetAccess(AccessType::READ);
+  objects_[name] = {std::move(descriptor_ptr)};
+}
+
+void Arguments::RenameArgs(const std::string& postfix,
+                           std::string* code) const {
+  size_t next_position = code->find(kArgsPrefix);
+  while (next_position != std::string::npos) {
+    size_t arg_pos = next_position + strlen(kArgsPrefix);
+    std::string arg_name = GetNextWord(*code, arg_pos);
+    code->replace(arg_pos, arg_name.size(), arg_name + postfix);
+    next_position = code->find(kArgsPrefix, arg_pos + arg_name.size());
+  }
+}
+
+absl::Status Arguments::Merge(Arguments&& args, const std::string& postfix,
+                              const std::vector<std::string>& exception_names) {
+  std::vector<std::string> object_names;
+  object_names.reserve(args.object_refs_.size() + args.objects_.size());
+  for (auto& v : args.object_refs_) {
+    if (std::find(exception_names.begin(), exception_names.end(), v.first) !=
+        exception_names.end()) {
+      continue;
+    }
+    object_names.push_back(v.first);
+    const std::string name = v.first + postfix;
+    if (object_refs_.find(name) != object_refs_.end()) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Object reference name collision. Name - ", name));
+    }
+    object_refs_[name] = {std::move(v.second)};
+  }
+  for (auto& v : args.objects_) {
+    if (std::find(exception_names.begin(), exception_names.end(), v.first) !=
+        exception_names.end()) {
+      continue;
+    }
+    object_names.push_back(v.first);
+    const std::string name = v.first + postfix;
+    if (objects_.find(name) != objects_.end()) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Object name collision. Name - ", name));
+    }
+    objects_[name] = {std::move(v.second)};
+  }
+  for (const auto& v : args.int_values_) {
+    AddInt(RenameArg(object_names, postfix, v.first), v.second.value);
+  }
+  for (const auto& v : args.float_values_) {
+    AddFloat(RenameArg(object_names, postfix, v.first), v.second.value);
+  }
+  for (const auto& v : args.half_values_) {
+    AddHalf(RenameArg(object_names, postfix, v.first), v.second.value);
+  }
+  return absl::OkStatus();
+}
+
+absl::Status Arguments::GetDescriptor(const std::string& name,
+                                      GPUObjectDescriptor** descriptor) const {
+  auto it_ref = object_refs_.find(name);
+  if (it_ref != object_refs_.end()) {
+    *descriptor = it_ref->second.get();
+    return absl::OkStatus();
+  }
+  auto it = objects_.find(name);
+  if (it != objects_.end()) {
+    *descriptor = it->second.get();
+    return absl::OkStatus();
+  }
+  return absl::NotFoundError(absl::StrCat("No GPU object with name - ", name));
+}
+
+void Arguments::ReleaseCPURepresentation() {
+  for (auto& t : objects_) {
+    t.second->Release();
+  }
+}
+
+void Arguments::GetActiveArguments(const std::string& code) {
+  for (auto& float_val : float_values_) {
+    float_val.second.active = HasWord(kArgsPrefix + float_val.first, code);
+  }
+  for (auto& int_val : int_values_) {
+    int_val.second.active = HasWord(kArgsPrefix + int_val.first, code);
+  }
+  for (auto& half_val : half_values_) {
+    half_val.second.active = HasWord(kArgsPrefix + half_val.first, code);
+  }
+}
+
+int Arguments::GetReadTexturesCount(const GpuInfo& gpu_info) const {
+  int counter = 0;
+  for (auto& t : objects_) {
+    counter += t.second->GetGPUResources(gpu_info).GetReadImagesCount();
+  }
+  for (auto& t : object_refs_) {
+    counter += t.second->GetGPUResources(gpu_info).GetReadImagesCount();
+  }
+  return counter;
+}
+
+int Arguments::GetWriteTexturesCount(const GpuInfo& gpu_info) const {
+  int counter = 0;
+  for (auto& t : objects_) {
+    counter += t.second->GetGPUResources(gpu_info).GetWriteImagesCount();
+  }
+  for (auto& t : object_refs_) {
+    counter += t.second->GetGPUResources(gpu_info).GetWriteImagesCount();
+  }
+  return counter;
+}
+
+void Arguments::SetStateValueForAllObjects(const std::string& key,
+                                           const std::string& value) {
+  for (auto& obj : object_refs_) {
+    obj.second->SetStateVar(key, value);
+  }
+  for (auto& obj : objects_) {
+    obj.second->SetStateVar(key, value);
+  }
+}
+
+absl::Status Arguments::Compile(const GpuInfo& gpu_info, std::string* code) {
+  RETURN_IF_ERROR(AddObjectsScalarArgs(gpu_info));
+  GetActiveArguments(*code);
+  RETURN_IF_ERROR(ResolveKernelGlobalSpaceBuffers(gpu_info, code));
+  return absl::OkStatus();
+}
+
+absl::Status Arguments::AddObjectsScalarArgs(const GpuInfo& gpu_info) {
+  for (auto& t : objects_) {
+    const auto resources = t.second->GetGPUResources(gpu_info);
+    for (const auto& r : resources.ints) {
+      AddInt(absl::StrCat(t.first, "_", r));
+    }
+    for (const auto& r : resources.floats) {
+      AddFloat(absl::StrCat(t.first, "_", r));
+    }
+  }
+  for (auto& t : object_refs_) {
+    const auto resources = t.second->GetGPUResources(gpu_info);
+    for (const auto& r : resources.ints) {
+      AddInt(absl::StrCat(t.first, "_", r));
+    }
+    for (const auto& r : resources.floats) {
+      AddFloat(absl::StrCat(t.first, "_", r));
+    }
+  }
+  return absl::OkStatus();
+}
+
+void Arguments::ResolveArgsPass(std::string* code) const {
+  size_t position = 0;
+  size_t next_position = code->find(kArgsPrefix);
+  while (next_position != std::string::npos) {
+    size_t arg_pos = next_position;
+    next_position += strlen(kArgsPrefix);
+    std::string object_name = GetNextWord(*code, next_position);
+    std::string new_name = object_name;
+    code->replace(arg_pos, object_name.size() + strlen(kArgsPrefix), new_name);
+    position = arg_pos + new_name.size();
+    next_position = code->find(kArgsPrefix, position);
+  }
+}
+
+absl::Status Arguments::ResolveKernelGlobalSpaceBuffers(const GpuInfo& gpu_info,
+                                                        std::string* code) {
+  for (auto it = objects_.begin(); it != objects_.end();) {
+    const auto* buffer_desc =
+        dynamic_cast<const BufferDescriptor*>(it->second.get());
+    if (!buffer_desc || buffer_desc->memory_type != MemoryType::CONSTANT) {
+      ++it;
+      continue;
+    }
+    bool is_kernel_global_space = false;
+    for (const auto& attribute : buffer_desc->attributes) {
+      if (attribute == "kernel_global_space") {
+        is_kernel_global_space = true;
+        break;
+      }
+    }
+    if (!is_kernel_global_space) {
+      ++it;
+      continue;
+    }
+    std::string declaration;
+    if (!BufferToKernelLanguage(gpu_info, it->first, buffer_desc, &declaration)
+             .ok()) {
+      ++it;
+      continue;
+    }
+    *code = declaration + *code;
+    objects_.erase(it++);
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/task/arguments.h

@@ -0,0 +1,153 @@
+/* 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_TASK_ARGUMENTS_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_ARGUMENTS_H_
+
+#include <map>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/access_type.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_object_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tflite_serialization_base_generated.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+class ArgumentsBinder {
+ public:
+  virtual absl::Status SetInt(const std::string& name, int value) = 0;
+  virtual absl::Status SetFloat(const std::string& name, float value) = 0;
+  virtual absl::Status SetHalf(const std::string& name, half value) = 0;
+  virtual ~ArgumentsBinder() = default;
+};
+
+class Arguments : public ArgumentsBinder {
+ public:
+  Arguments() = default;
+  ~Arguments() override = default;
+
+  // Move only
+  Arguments(Arguments&& args) = default;
+  Arguments& operator=(Arguments&& args) = default;
+  Arguments(const Arguments&) = delete;
+  Arguments& operator=(const Arguments&) = delete;
+
+  void AddFloat(const std::string& name, float value = 0.0f);
+  void AddHalf(const std::string& name, half value = half(0.0f));
+  void AddInt(const std::string& name, int value = 0);
+  absl::Status SetInt(const std::string& name, int value) override;
+  absl::Status SetFloat(const std::string& name, float value) override;
+  absl::Status SetHalf(const std::string& name, half value) override;
+  void AddObjectRef(const std::string& name, AccessType access_type,
+                    GPUObjectDescriptorPtr&& descriptor_ptr);
+  void AddObject(const std::string& name,
+                 GPUObjectDescriptorPtr&& descriptor_ptr);
+
+  void RenameArgs(const std::string& postfix, std::string* code) const;
+  absl::Status Merge(Arguments&& args, const std::string& postfix,
+                     const std::vector<std::string>& exception_names = {});
+
+  absl::Status GetDescriptor(const std::string& name,
+                             GPUObjectDescriptor** descriptor) const;
+
+  int GetReadTexturesCount(const GpuInfo& gpu_info) const;
+  int GetWriteTexturesCount(const GpuInfo& gpu_info) const;
+
+  void ReleaseCPURepresentation();
+
+  void GetActiveArguments(const std::string& code);
+
+  void SetStateValueForAllObjects(const std::string& key,
+                                  const std::string& value);
+
+  struct IntValue {
+    int value;
+
+    // many uniforms generated automatically and not used
+    // to reduce amount of data transferred we adding this optimization
+    bool active = false;
+  };
+  struct FloatValue {
+    float value;
+
+    // many uniforms generated automatically and not used
+    // to reduce amount of data transferred we adding this optimization
+    bool active = false;
+  };
+  struct HalfValue {
+    half value;
+
+    // many uniforms generated automatically and not used
+    // to reduce amount of data transferred we adding this optimization
+    bool active = false;
+  };
+
+  const std::map<std::string, IntValue>& GetIntValues() const {
+    return int_values_;
+  }
+  const std::map<std::string, FloatValue>& GetFloatValues() const {
+    return float_values_;
+  }
+  const std::map<std::string, HalfValue>& GetHalfValues() const {
+    return half_values_;
+  }
+
+  const std::map<std::string, GPUObjectDescriptorPtr>& GetObjectRefs() const {
+    return object_refs_;
+  }
+  const std::map<std::string, GPUObjectDescriptorPtr>& GetObjects() const {
+    return objects_;
+  }
+  void MoveObjectRefs(std::map<std::string, GPUObjectDescriptorPtr>* result) {
+    *result = std::move(object_refs_);
+  }
+
+  absl::Status Compile(const GpuInfo& gpu_info,
+                       std::string* code);
+
+  void ResolveObjectNames(const std::string& object_name,
+                          const std::vector<std::string>& member_names,
+                          std::string* code) const;
+  absl::Status AddObjectsScalarArgs(const GpuInfo& gpu_info);
+  void ResolveArgsPass(std::string* code) const;
+
+ private:
+  friend flatbuffers::Offset<tflite::gpu::data::Arguments> Encode(
+      const Arguments& args, flatbuffers::FlatBufferBuilder* builder);
+  friend absl::Status Decode(const tflite::gpu::data::Arguments* fb_args,
+                             Arguments* args);
+
+  absl::Status ResolveKernelGlobalSpaceBuffers(const GpuInfo& gpu_info,
+                                               std::string* code);
+
+  static constexpr char kArgsPrefix[] = "args.";
+
+  std::map<std::string, IntValue> int_values_;
+  std::map<std::string, FloatValue> float_values_;
+  std::map<std::string, HalfValue> half_values_;
+
+  std::map<std::string, GPUObjectDescriptorPtr> object_refs_;
+  std::map<std::string, GPUObjectDescriptorPtr> objects_;
+};
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_ARGUMENTS_H_

tensorflow/lite/delegates/gpu/common/task/buffer_desc.cc

@@ -0,0 +1,152 @@
+/* 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/task/buffer_desc.h"
+
+#include <string>
+#include <vector>
+
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/util.h"
+
+namespace tflite {
+namespace gpu {
+
+void BufferDescriptor::Release() { data.clear(); }
+
+GPUResources BufferDescriptor::GetGPUResources(const GpuInfo& gpu_info) const {
+  GPUResources resources;
+  GPUBufferDescriptor desc;
+  desc.data_type = element_type;
+  desc.access_type = access_type_;
+  desc.element_size = element_size;
+  desc.memory_type = memory_type;
+  desc.attributes = attributes;
+  if (gpu_info.IsGlsl() && memory_type == tflite::gpu::MemoryType::CONSTANT) {
+    desc.attributes.push_back(
+        std::to_string(size / (element_size * SizeOf(element_type))));
+  }
+  resources.buffers.push_back({"buffer", desc});
+  return resources;
+}
+
+absl::Status BufferDescriptor::PerformSelector(
+    const GpuInfo& gpu_info, absl::string_view selector,
+    const std::vector<std::string>& args,
+    const std::vector<std::string>& template_args, std::string* result) const {
+  if (selector == "Read") {
+    return PerformReadSelector(gpu_info, args, result);
+  } else if (selector == "Write") {
+    return PerformWriteSelector(gpu_info, args, result);
+  } else if (selector == "GetPtr") {
+    return PerformGetPtrSelector(args, template_args, result);
+  } else {
+    return absl::NotFoundError(absl::StrCat(
+        "BufferDescriptor don't have selector with name - ", selector));
+  }
+}
+
+absl::Status BufferDescriptor::PerformReadSelector(
+    const GpuInfo& gpu_info, const std::vector<std::string>& args,
+    std::string* result) const {
+  if (args.size() != 1) {
+    return absl::NotFoundError(
+        absl::StrCat("BufferDescriptor Read require one argument, but ",
+                     args.size(), " was passed"));
+  }
+  if (gpu_info.IsGlsl()) {
+    if (element_type == DataType::FLOAT16 &&
+        !gpu_info.IsGlslSupportsExplicitFp16()) {
+      if (memory_type == MemoryType::CONSTANT) {
+        bool is_kernel_global_space = false;
+        for (const auto& attribute : attributes) {
+          if (attribute == "kernel_global_space") {
+            is_kernel_global_space = true;
+            break;
+          }
+        }
+        if (is_kernel_global_space) {
+          *result = absl::StrCat("buffer[", args[0], "]");
+          return absl::OkStatus();
+        }
+        const std::string arg0 = "(" + args[0] + ")";
+        *result =
+            absl::StrCat("vec4(unpackHalf2x16(buffer[", arg0, " / 2][", arg0,
+                         " % 2 == 0 ? 0 : 2]), unpackHalf2x16(buffer[", arg0,
+                         " / 2][", arg0, " % 2 == 0 ? 1 : 3]))");
+      } else {
+        *result =
+            absl::StrCat("vec4(unpackHalf2x16(buffer[", args[0],
+                         "].x), unpackHalf2x16(buffer[", args[0], "].y))");
+      }
+    } else {
+      *result = absl::StrCat("buffer[", args[0], "]");
+    }
+    return absl::OkStatus();
+  } else {
+    *result = absl::StrCat("buffer[", args[0], "]");
+    return absl::OkStatus();
+  }
+}
+
+absl::Status BufferDescriptor::PerformWriteSelector(
+    const GpuInfo& gpu_info, const std::vector<std::string>& args,
+    std::string* result) const {
+  if (args.size() != 2) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "BufferDescriptor Write require two arguments(value, index), but ",
+        args.size(), " was passed"));
+  }
+  *result = absl::StrCat("buffer[", args[1], "] = ", args[0]);
+  return absl::OkStatus();
+}
+
+absl::Status BufferDescriptor::PerformGetPtrSelector(
+    const std::vector<std::string>& args,
+    const std::vector<std::string>& template_args, std::string* result) const {
+  if (args.size() > 1) {
+    return absl::NotFoundError(absl::StrCat(
+        "BufferDescriptor GetPtr require one or zero arguments, but ",
+        args.size(), " was passed"));
+  }
+  if (template_args.size() > 1) {
+    return absl::NotFoundError(
+        absl::StrCat("BufferDescriptor GetPtr require one or zero teemplate "
+                     "arguments, but ",
+                     template_args.size(), " was passed"));
+  }
+  std::string conversion;
+  if (template_args.size() == 1) {
+    const std::string type_name = ToCLDataType(element_type, element_size);
+    if (type_name != template_args[0]) {
+      conversion = absl::StrCat("(", MemoryTypeToCLType(memory_type), " ",
+                                template_args[0], "*)&");
+    }
+  }
+  if (args.empty()) {
+    *result = absl::StrCat(conversion, "buffer");
+  } else if (conversion.empty()) {
+    *result = absl::StrCat("(buffer + ", args[0], ")");
+  } else {
+    *result = absl::StrCat(conversion, "buffer[", args[0], "]");
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/task/buffer_desc.h

@@ -0,0 +1,71 @@
+/* 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_TASK_BUFFER_DESC_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_BUFFER_DESC_H_
+
+#include <string>
+#include <vector>
+
+#include "absl/strings/string_view.h"
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_object_desc.h"
+
+namespace tflite {
+namespace gpu {
+
+struct BufferDescriptor : public GPUObjectDescriptor {
+  DataType element_type;
+  int element_size;
+  MemoryType memory_type = MemoryType::GLOBAL;
+  std::vector<std::string> attributes;
+
+  // optional
+  int size = 0;
+  std::vector<uint8_t> data;
+
+  BufferDescriptor() = default;
+  BufferDescriptor(const BufferDescriptor&) = default;
+  BufferDescriptor& operator=(const BufferDescriptor&) = default;
+  BufferDescriptor(BufferDescriptor&& desc) = default;
+  BufferDescriptor& operator=(BufferDescriptor&& desc) = default;
+
+  absl::Status PerformSelector(const GpuInfo& gpu_info,
+                               absl::string_view selector,
+                               const std::vector<std::string>& args,
+                               const std::vector<std::string>& template_args,
+                               std::string* result) const override;
+
+  GPUResources GetGPUResources(const GpuInfo& gpu_info) const override;
+  absl::Status PerformReadSelector(const GpuInfo& gpu_info,
+                                   const std::vector<std::string>& args,
+                                   std::string* result) const;
+  absl::Status PerformWriteSelector(const GpuInfo& gpu_info,
+                                    const std::vector<std::string>& args,
+                                    std::string* result) const;
+  absl::Status PerformGetPtrSelector(
+      const std::vector<std::string>& args,
+      const std::vector<std::string>& template_args, std::string* result) const;
+
+  void Release() override;
+
+  uint64_t GetSizeInBytes() const override { return data.size(); };
+};
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_BUFFER_DESC_H_

tensorflow/lite/delegates/gpu/common/task/compiler_options.h

@@ -0,0 +1,35 @@
+/* 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_TASK_COMPILER_OPTIONS_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_COMPILER_OPTIONS_H_
+
+namespace tflite {
+namespace gpu {
+
+enum class CompilerOptions {
+  kAdrenoFullSimd,
+  kAdrenoMoreWaves,
+  kClFastRelaxedMath,
+  kClDisableOptimizations,
+  kCl20,
+  kCl30,
+  kClRegisterAllocation64,
+};
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_COMPILER_OPTIONS_H_

tensorflow/lite/delegates/gpu/common/task/gpu_object_desc.h

@@ -0,0 +1,230 @@
+/* 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_TASK_GPU_OBJECT_DESC_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_GPU_OBJECT_DESC_H_
+
+#include <functional>
+#include <map>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/strings/string_view.h"
+#include "tensorflow/lite/delegates/gpu/common/access_type.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/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tflite_serialization_base_generated.h"
+
+namespace tflite {
+namespace gpu {
+
+struct GPUImage2DDescriptor {
+  DataType data_type;
+  bool normalized = false;   // used with INT data types, if normalized, we read
+                             // in kernel float data.
+  DataType normalized_type;  // can be FLOAT32 or FLOAT16, using with normalized
+                             // = true
+  AccessType access_type;
+};
+
+struct GPUImage3DDescriptor {
+  DataType data_type;
+  AccessType access_type;
+};
+
+struct GPUImage2DArrayDescriptor {
+  DataType data_type;
+  AccessType access_type;
+};
+
+struct GPUImageBufferDescriptor {
+  DataType data_type;
+  AccessType access_type;
+};
+
+struct GPUCustomMemoryDescriptor {
+  std::string type_name;
+};
+
+enum class MemoryType { GLOBAL, CONSTANT, LOCAL };
+
+struct GPUBufferDescriptor {
+  DataType data_type;
+  AccessType access_type;
+  int element_size;
+  MemoryType memory_type = MemoryType::GLOBAL;
+  std::vector<std::string> attributes;
+};
+
+struct GPUResources {
+  std::vector<std::string> ints;
+  std::vector<std::string> floats;
+  std::vector<std::pair<std::string, GPUBufferDescriptor>> buffers;
+  std::vector<std::pair<std::string, GPUImage2DDescriptor>> images2d;
+  std::vector<std::pair<std::string, GPUImage2DArrayDescriptor>> image2d_arrays;
+  std::vector<std::pair<std::string, GPUImage3DDescriptor>> images3d;
+  std::vector<std::pair<std::string, GPUImageBufferDescriptor>> image_buffers;
+  std::vector<std::pair<std::string, GPUCustomMemoryDescriptor>>
+      custom_memories;
+
+  std::vector<std::string> GetNames() const {
+    std::vector<std::string> names = ints;
+    names.insert(names.end(), floats.begin(), floats.end());
+    for (const auto& obj : buffers) {
+      names.push_back(obj.first);
+    }
+    for (const auto& obj : images2d) {
+      names.push_back(obj.first);
+    }
+    for (const auto& obj : image2d_arrays) {
+      names.push_back(obj.first);
+    }
+    for (const auto& obj : images3d) {
+      names.push_back(obj.first);
+    }
+    for (const auto& obj : image_buffers) {
+      names.push_back(obj.first);
+    }
+    for (const auto& obj : custom_memories) {
+      names.push_back(obj.first);
+    }
+    return names;
+  }
+
+  int GetReadImagesCount() const {
+    int counter = 0;
+    for (const auto& t : images2d) {
+      if (t.second.access_type == tflite::gpu::AccessType::READ) {
+        counter++;
+      }
+    }
+    for (const auto& t : image2d_arrays) {
+      if (t.second.access_type == tflite::gpu::AccessType::READ) {
+        counter++;
+      }
+    }
+    for (const auto& t : images3d) {
+      if (t.second.access_type == tflite::gpu::AccessType::READ) {
+        counter++;
+      }
+    }
+    for (const auto& t : image_buffers) {
+      if (t.second.access_type == tflite::gpu::AccessType::READ) {
+        counter++;
+      }
+    }
+    return counter;
+  }
+
+  int GetWriteImagesCount() const {
+    int counter = 0;
+    for (const auto& t : images2d) {
+      if (t.second.access_type == tflite::gpu::AccessType::WRITE) {
+        counter++;
+      }
+    }
+    for (const auto& t : image2d_arrays) {
+      if (t.second.access_type == tflite::gpu::AccessType::WRITE) {
+        counter++;
+      }
+    }
+    for (const auto& t : images3d) {
+      if (t.second.access_type == tflite::gpu::AccessType::WRITE) {
+        counter++;
+      }
+    }
+    for (const auto& t : image_buffers) {
+      if (t.second.access_type == tflite::gpu::AccessType::WRITE) {
+        counter++;
+      }
+    }
+    return counter;
+  }
+};
+
+struct GenericGPUResourcesWithValue {
+  std::vector<std::pair<std::string, int>> ints;
+  std::vector<std::pair<std::string, float>> floats;
+
+  void AddFloat(absl::string_view name, float value) {
+    floats.emplace_back(name, value);
+  }
+  void AddInt(absl::string_view name, int value) {
+    ints.emplace_back(name, value);
+  }
+};
+
+class GPUObjectDescriptor {
+ public:
+  GPUObjectDescriptor() = default;
+  GPUObjectDescriptor(const GPUObjectDescriptor&) = default;
+  GPUObjectDescriptor& operator=(const GPUObjectDescriptor&) = default;
+  GPUObjectDescriptor(GPUObjectDescriptor&& obj_desc) = default;
+  GPUObjectDescriptor& operator=(GPUObjectDescriptor&& obj_desc) = default;
+  virtual ~GPUObjectDescriptor() = default;
+
+  void SetStateVar(absl::string_view key, absl::string_view value) const {
+    auto it = state_vars_.find(key);
+    if (it == state_vars_.end()) {
+      state_vars_[std::string(key)] = std::string(value);
+    } else {
+      it->second = std::string(value);
+    }
+  }
+
+  virtual absl::Status PerformConstExpr(const tflite::gpu::GpuInfo& gpu_info,
+                                        absl::string_view const_expr,
+                                        std::string* result) const {
+    return absl::UnimplementedError(
+        "No implementation of perform const expression");
+  }
+
+  virtual absl::Status PerformSelector(
+      const GpuInfo& gpu_info, absl::string_view selector,
+      const std::vector<std::string>& args,
+      const std::vector<std::string>& template_args,
+      std::string* result) const {
+    return absl::UnimplementedError("No implementation of perform selector");
+  }
+  virtual GPUResources GetGPUResources(const GpuInfo& gpu_info) const {
+    return GPUResources();
+  }
+
+  virtual void Release() {}
+
+  // For internal use, will work correct only for const objects and before
+  // Release() call.
+  virtual uint64_t GetSizeInBytes() const { return 0; }
+
+  void SetAccess(AccessType access_type) { access_type_ = access_type; }
+  AccessType GetAccess() const { return access_type_; }
+
+ protected:
+  friend flatbuffers::Offset<tflite::gpu::data::GPUObjectDescriptor> Encode(
+      const GPUObjectDescriptor& desc, flatbuffers::FlatBufferBuilder* builder);
+  friend void Decode(const tflite::gpu::data::GPUObjectDescriptor* fb_obj,
+                     GPUObjectDescriptor* obj);
+  mutable std::map<std::string, std::string, std::less<>> state_vars_;
+  AccessType access_type_;
+};
+
+using GPUObjectDescriptorPtr = std::unique_ptr<GPUObjectDescriptor>;
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_GPU_OBJECT_DESC_H_

tensorflow/lite/delegates/gpu/common/task/gpu_operation.cc

@@ -0,0 +1,766 @@
+/* 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/task/gpu_operation.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <cstring>
+#include <map>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
+#include "absl/strings/ascii.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_replace.h"
+#include "absl/strings/string_view.h"
+#include "absl/strings/substitute.h"
+#include "tensorflow/lite/delegates/gpu/common/access_type.h"
+#include "tensorflow/lite/delegates/gpu/common/task/util.h"
+#include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
+
+namespace tflite {
+namespace gpu {
+namespace {
+int3 GetWorkGroupsCountInternal(int grid_dimension, const int3& grid_size,
+                                const int3& work_group_size,
+                                const int3& work_group_launch_order) {
+  int3 work_groups_count;
+  if (grid_dimension == 1) {
+    work_groups_count.x = DivideRoundUp(grid_size.x, work_group_size.x);
+    work_groups_count.y = 1;
+    work_groups_count.z = 1;
+  } else if (grid_dimension == 2) {
+    int3 wgs;
+    wgs.x = DivideRoundUp(grid_size.x, work_group_size.x);
+    wgs.y = DivideRoundUp(grid_size.y, work_group_size.y);
+    work_groups_count.x = wgs[work_group_launch_order[0]];
+    work_groups_count.y = wgs[work_group_launch_order[1]];
+    work_groups_count.z = 1;
+  } else {  // grid_dimension == 3
+    int3 wgs;
+    wgs.x = DivideRoundUp(grid_size.x, work_group_size.x);
+    wgs.y = DivideRoundUp(grid_size.y, work_group_size.y);
+    wgs.z = DivideRoundUp(grid_size.z, work_group_size.z);
+    work_groups_count.x = wgs[work_group_launch_order[0]];
+    work_groups_count.y = wgs[work_group_launch_order[1]];
+    work_groups_count.z = wgs[work_group_launch_order[2]];
+  }
+  return work_groups_count;
+}
+
+std::string GetElementWiseCode(const OperationDef& op_def) {
+  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()) return; \n";
+  c += "  args.src_tensor::type src = args.src_tensor.Read(X, Y, Z);\n";
+  c += "  args.dst_tensor.Write(src, X, Y, Z);\n";
+  c += "} \n";
+  return c;
+}
+
+bool NeedsBroadcast(const TensorDescriptor& desc, const BHWC& shape) {
+  bool needs_broadcast = shape.w == 1 || shape.h == 1 || shape.c == 1;
+  if (desc.HasAxis(Axis::BATCH)) {
+    needs_broadcast = needs_broadcast || shape.b == 1;
+  }
+  return needs_broadcast;
+}
+
+std::string GetStringWithoutComments(const std::string& code) {
+  std::string result;
+  result.reserve(code.size());
+  for (size_t i = 0; i < code.size(); ++i) {
+    // skip long comment /*...*/
+    if (code[i] == '/' && i + 1 < code.size() && code[i + 1] == '*') {
+      i = i + 3;
+      for (; i < code.size() && (code[i - 1] != '*' && code[i] != '/'); ++i) {
+      }
+      continue;
+    }
+    // skip short comment //...
+    if (code[i] == '/' && i + 1 < code.size() && code[i + 1] == '/') {
+      i = i + 2;
+      for (; i < code.size() && code[i] != '\n'; ++i) {
+      }
+      if (i != code.size()) {
+        // i is '\n'
+        result.push_back(code[i]);
+      }
+      continue;
+    }
+    result.push_back(code[i]);
+  }
+  return result;
+}
+
+bool IsWordSymbol(char symbol) {
+  return absl::ascii_isalnum(symbol) || symbol == '_';
+}
+
+// Word is a token consisting of ascii_isalnum symbols or '_'
+// (see above IsWordSymbol(char symbol))
+// ReplaceAllWords replace all specified word-tokens(old_word) with new_word
+void ReplaceAllWords(const std::string& old_word, const std::string& new_word,
+                     std::string* str) {
+  for (size_t position = str->find(old_word); position != std::string::npos;
+       position = str->find(old_word, position)) {
+    const char prev = position == 0 ? '.' : (*str)[position - 1];
+    const char next = position + old_word.size() < str->size()
+                          ? (*str)[position + old_word.size()]
+                          : '.';
+    if (IsWordSymbol(prev) || IsWordSymbol(next)) {
+      position += 1;
+      continue;
+    }
+    str->replace(position, old_word.size(), new_word);
+    position += new_word.size();
+  }
+}
+
+struct LinkableContext {
+  std::string code;
+  TensorDescriptor* tensor_desc;
+};
+
+absl::Status ResolveLinking(const GpuInfo& gpu_info,
+                            const LinkableContext& linkable_context,
+                            std::vector<std::string>* function_args,
+                            std::string* result) {
+  std::string value_name, x_coord, y_coord, z_coord, s_coord, b_coord;
+  RETURN_IF_ERROR(
+      linkable_context.tensor_desc->GetLinkingContextFromWriteSelector(
+          *function_args, &value_name, &x_coord, &y_coord, &z_coord, &s_coord,
+          &b_coord));
+  const std::string new_value_name = value_name + "_final";
+  const std::string out_var_declaration =
+      "\n" +
+      GetTypeDeclaration(gpu_info, linkable_context.tensor_desc->GetDataType(),
+                         4) +
+      " " + new_value_name + ";\n";
+  *result = out_var_declaration +
+            "{  // elementwise code with input:" + value_name +
+            " output:" + new_value_name + "\n" +
+            absl::Substitute(linkable_context.code, "") + "\n}\n";
+  *result = absl::StrReplaceAll(*result, {{"\n", "\n  "},
+                                          {"in_value", value_name},
+                                          {"out_value", new_value_name},
+                                          {"X_COORD", x_coord},
+                                          {"Y_COORD", y_coord},
+                                          {"Z_COORD", z_coord},
+                                          {"S_COORD", s_coord},
+                                          {"B_COORD", b_coord}});
+
+  (*function_args)[0] = new_value_name;
+  return absl::OkStatus();
+}
+
+// resolve constructions of type: args.object_name::const_expr_name
+// Example: 'args.dst_tensor::type' can be replaced with 'float4'
+absl::Status ResolveConstExprPass(const GpuInfo& gpu_info,
+                                  const Arguments& args, std::string* code) {
+  std::string result;
+  size_t position = 0;
+  constexpr char kArgsPrefix[] = "args.";
+  size_t next_position = code->find(kArgsPrefix);
+  while (next_position != std::string::npos) {
+    size_t arg_pos = next_position;
+    next_position += strlen(kArgsPrefix);
+    std::string object_name = GetNextWord(*code, next_position);
+    if (next_position + object_name.size() > code->size() - 2) {
+      next_position = code->find(kArgsPrefix, next_position);
+      continue;
+    }
+    char next0 = (*code)[next_position + object_name.size()];
+    char next1 = (*code)[next_position + object_name.size() + 1];
+    if (next0 == ':' && next1 == ':') {
+      next_position += object_name.size() + 2;
+      std::string const_expr_name = GetNextWord(*code, next_position);
+      next_position += const_expr_name.size();
+      std::string patch;
+      tflite::gpu::GPUObjectDescriptor* desc_ptr;
+      RETURN_IF_ERROR(args.GetDescriptor(object_name, &desc_ptr));
+      RETURN_IF_ERROR(
+          desc_ptr->PerformConstExpr(gpu_info, const_expr_name, &patch));
+      code->replace(arg_pos, next_position - arg_pos, patch);
+      position = arg_pos + patch.size();
+    } else {
+      position = arg_pos + strlen(kArgsPrefix);
+    }
+    next_position = code->find(kArgsPrefix, position);
+  }
+  return absl::OkStatus();
+}
+
+// resolve constructions of type: args.object_name.method_name(list of args)
+// Example: 'args.bias.Read(S)' can be replaced with 'args.bias_buffer[S]'
+absl::Status ResolveSelectorsPass(
+    const GpuInfo& gpu_info,
+    const std::map<std::string, LinkableContext>& linkables,
+    const Arguments& args, std::string* code) {
+  std::string result;
+  size_t position = 0;
+  constexpr char kArgsPrefix[] = "args.";
+  size_t next_position = code->find(kArgsPrefix);
+  while (next_position != std::string::npos) {
+    size_t arg_pos = next_position;
+    next_position += strlen(kArgsPrefix);
+    std::string object_name = GetNextWord(*code, next_position);
+    char next = (*code)[next_position + object_name.size()];
+    if (next == '.') {
+      next_position += object_name.size() + 1;
+      std::string selector_name = GetNextWord(*code, next_position);
+      next_position += selector_name.size();
+      next = (*code)[next_position];
+      std::vector<std::string> template_args;
+      if (next == '<') {
+        size_t close_bracket_pos;
+        RETURN_IF_ERROR(ParseArgsInsideBrackets(
+            *code, next_position, &close_bracket_pos, &template_args));
+        next_position = close_bracket_pos;
+        next = (*code)[next_position];
+      }
+      if (next != '(') {
+        return absl::NotFoundError(absl::StrCat(
+            "Expected ( after ", object_name, ".", selector_name, " call"));
+      }
+      std::vector<std::string> function_args;
+      size_t close_bracket_pos;
+      RETURN_IF_ERROR(ParseArgsInsideBrackets(
+          *code, next_position, &close_bracket_pos, &function_args));
+      for (auto& arg : function_args) {
+        RETURN_IF_ERROR(ResolveSelectorsPass(gpu_info, {}, args, &arg));
+      }
+      std::string patch;
+      std::string linkable_patch;
+      GPUObjectDescriptor* desc_ptr;
+      RETURN_IF_ERROR(args.GetDescriptor(object_name, &desc_ptr));
+      auto names = desc_ptr->GetGPUResources(gpu_info).GetNames();
+      if (desc_ptr && !linkables.empty() && selector_name == "Write") {
+        auto it = linkables.find(object_name);
+        if (it != linkables.end()) {
+          RETURN_IF_ERROR(ResolveLinking(gpu_info, it->second, &function_args,
+                                         &linkable_patch));
+          RETURN_IF_ERROR(
+              ResolveConstExprPass(gpu_info, args, &linkable_patch));
+          RETURN_IF_ERROR(
+              ResolveSelectorsPass(gpu_info, {}, args, &linkable_patch));
+        }
+      }
+      RETURN_IF_ERROR(desc_ptr->PerformSelector(
+          gpu_info, selector_name, function_args, template_args, &patch));
+      for (const auto& member_name : names) {
+        const std::string new_name =
+            kArgsPrefix + object_name + "_" + member_name;
+        ReplaceAllWords(member_name, new_name, &patch);
+      }
+      if (!linkable_patch.empty()) {
+        patch = "{\n" + linkable_patch + patch + ";\n}";
+      }
+      code->replace(arg_pos, close_bracket_pos - arg_pos, patch);
+      position = arg_pos + patch.size();
+    } else {
+      position = arg_pos + strlen(kArgsPrefix);
+    }
+    next_position = code->find(kArgsPrefix, position);
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+DataType OperationDef::GetDataType() const {
+  return DeduceDataTypeFromPrecision(precision);
+}
+
+DataType OperationDef::GetPrimaryDataType() const {
+  return src_tensors[0].GetDataType();
+}
+TensorStorageType OperationDef::GetPrimaryStorageType() const {
+  return src_tensors[0].GetStorageType();
+}
+
+bool OperationDef::IsBatchSupported() const {
+  for (const auto& src : src_tensors) {
+    if (src.HasAxis(Axis::BATCH)) {
+      return true;
+    }
+  }
+  for (const auto& dst : dst_tensors) {
+    if (dst.HasAxis(Axis::BATCH)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+GPUOperation::GPUOperation(const OperationDef& definition)
+    : definition_(definition) {}
+
+void GPUOperation::SetSrc(GpuSpatialTensor* ptr, int index) {
+  if (index >= src_.size()) {
+    src_.resize(index + 1, nullptr);
+  }
+  src_[index] = ptr;
+}
+
+void GPUOperation::SetDst(GpuSpatialTensor* ptr, int index) {
+  if (index >= dst_.size()) {
+    dst_.resize(index + 1, nullptr);
+  }
+  dst_[index] = ptr;
+}
+
+GPUOperation::GPUOperation(GPUOperation&& operation)
+    : args_(std::move(operation.args_)),
+      code_(std::move(operation.code_)),
+      work_group_size_(operation.work_group_size_),
+      compiler_options_(std::move(operation.compiler_options_)),
+      tensor_to_grid_(operation.tensor_to_grid_),
+      flops_(operation.flops_),
+      const_args_size_(operation.const_args_size_),
+      definition_(std::move(operation.definition_)),
+      src_(std::move(operation.src_)),
+      dst_(std::move(operation.dst_)),
+      grid_dimension_(operation.grid_dimension_),
+      work_group_launch_order_(operation.work_group_launch_order_),
+      grid_size_(operation.grid_size_),
+      src_tensors_names_(std::move(operation.src_tensors_names_)),
+      dst_tensors_names_(std::move(operation.dst_tensors_names_)),
+      work_groups_count_(operation.work_groups_count_),
+      elementwise_(operation.elementwise_),
+      elementwise_inputs_(operation.elementwise_inputs_),
+      second_elementwise_tensor_name_(
+          operation.second_elementwise_tensor_name_),
+      linkable_count_(operation.linkable_count_),
+      elementwise_code_(std::move(operation.elementwise_code_)) {}
+
+GPUOperation& GPUOperation::operator=(GPUOperation&& operation) {
+  if (this != &operation) {
+    args_ = std::move(operation.args_);
+    code_ = std::move(operation.code_);
+    std::swap(work_group_size_, operation.work_group_size_);
+    compiler_options_ = std::move(operation.compiler_options_);
+    tensor_to_grid_ = operation.tensor_to_grid_;
+    flops_ = operation.flops_;
+    const_args_size_ = operation.const_args_size_;
+    definition_ = std::move(operation.definition_);
+    src_ = std::move(operation.src_);
+    dst_ = std::move(operation.dst_);
+    std::swap(grid_dimension_, operation.grid_dimension_);
+    std::swap(work_group_launch_order_, operation.work_group_launch_order_);
+    std::swap(grid_size_, operation.grid_size_);
+    src_tensors_names_ = std::move(operation.src_tensors_names_);
+    dst_tensors_names_ = std::move(operation.dst_tensors_names_);
+    std::swap(work_groups_count_, operation.work_groups_count_);
+    elementwise_ = operation.elementwise_;
+    std::swap(elementwise_inputs_, operation.elementwise_inputs_);
+    std::swap(second_elementwise_tensor_name_,
+              operation.second_elementwise_tensor_name_);
+    std::swap(linkable_count_, operation.linkable_count_);
+    elementwise_code_ = std::move(operation.elementwise_code_);
+  }
+  return *this;
+}
+
+absl::Status GPUOperation::AddOperation(const GpuInfo& gpu_info,
+                                        GPUOperation* operation) {
+  const auto prev_type = definition_.dst_tensors[0].GetDataType();
+  definition_.dst_tensors[0] = operation->definition_.dst_tensors[0];
+  if (!elementwise_) {
+    TensorDescriptor* dst_tensor_desc;
+    RETURN_IF_ERROR(
+        GetTensorDescriptor(dst_tensors_names_[0], &dst_tensor_desc));
+    operation->definition_.dst_tensors[0].CopyWithoutData(dst_tensor_desc);
+  }
+  linkable_count_ += (operation->linkable_count_ + 1);
+  std::string code = operation->elementwise_code_;
+  std::string unique_postfix = absl::StrCat("_link", linkable_count_);
+  code = absl::StrReplaceAll(
+      code, {{"interm_value", "interm_value" + unique_postfix}});
+  operation->args_.RenameArgs(unique_postfix, &code);
+  operation->second_elementwise_tensor_name_ += unique_postfix;
+  if (elementwise_code_.empty()) {
+    elementwise_code_ = code;
+    elementwise_inputs_ = operation->elementwise_inputs_;
+    second_elementwise_tensor_name_ =
+        operation->second_elementwise_tensor_name_;
+  } else {
+    if (operation->elementwise_inputs_ == 2) {
+      if (elementwise_inputs_ == 2) {
+        // if we have fusion of 2 2-input elementwise ops, we will get 3-input
+        // elementwise, but currently we support only max 2-input elementwise.
+        // So we will resolve one input here.
+        RETURN_IF_ERROR(ResolveSecondElementwiseInput());
+      }
+      second_elementwise_tensor_name_ =
+          operation->second_elementwise_tensor_name_;
+      elementwise_inputs_ = 2;
+    }
+    const std::string new_value_name = "interm_value" + unique_postfix;
+    code = absl::StrReplaceAll(code, {{"in_value", new_value_name}});
+    elementwise_code_ =
+        absl::StrReplaceAll(elementwise_code_, {{"out_value", new_value_name}});
+    const std::string out_var_declaration =
+        "\n" + GetTypeDeclaration(gpu_info, prev_type, 4) + " " +
+        new_value_name + ";\n";
+    elementwise_code_ =
+        absl::Substitute(elementwise_code_, out_var_declaration);
+    elementwise_code_ = elementwise_code_ + "\n" + code;
+  }
+  RETURN_IF_ERROR(args_.Merge(std::move(operation->args_), unique_postfix));
+  for (int i = 0; i < operation->src_tensors_names_.size(); ++i) {
+    definition_.src_tensors.push_back(
+        operation->definition_.src_tensors[i + 1]);
+    src_tensors_names_.push_back(operation->src_tensors_names_[i] +
+                                 unique_postfix);
+  }
+  for (int i = 0; i < operation->dst_tensors_names_.size(); ++i) {
+    dst_tensors_names_.push_back(operation->dst_tensors_names_[i] +
+                                 unique_postfix);
+  }
+  return absl::OkStatus();
+}
+
+absl::Status GPUOperation::ResolveSecondElementwiseInput() {
+  if (elementwise_inputs_ != 2) {
+    return absl::FailedPreconditionError(
+        "Can not apply ResolveSecondElementwiseInput for non 2 input "
+        "elementwise");
+  }
+  TensorDescriptor* tensor_desc;
+  RETURN_IF_ERROR(
+      GetTensorDescriptor(second_elementwise_tensor_name_, &tensor_desc));
+  std::string coords = "X_COORD, Y_COORD, S_COORD";
+  if (tensor_desc->HasAxis(Axis::BATCH)) {
+    coords += ", B_COORD";
+  }
+  const std::string read_code = "args." + second_elementwise_tensor_name_ +
+                                "::type second_value = args." +
+                                second_elementwise_tensor_name_ + ".Read(" +
+                                coords + ");\n";
+  elementwise_code_ = absl::StrReplaceAll(
+      elementwise_code_,
+      {{"in2_value", "second_value"}, {"READ_SECOND_VALUE", read_code}});
+  elementwise_inputs_ = 1;
+  return absl::OkStatus();
+}
+
+absl::Status GPUOperation::GetTensorDescriptor(const std::string& tensor_name,
+                                               TensorDescriptor** resutl) {
+  GPUObjectDescriptor* desc_ptr;
+  RETURN_IF_ERROR(args_.GetDescriptor(tensor_name, &desc_ptr));
+  *resutl = static_cast<TensorDescriptor*>(desc_ptr);
+  return absl::OkStatus();
+}
+
+void GPUOperation::AddSrcTensor(const std::string& tensor_name,
+                                const TensorDescriptor& desc) {
+  src_tensors_names_.push_back(tensor_name);
+  auto desc_new = std::make_unique<TensorDescriptor>(desc);
+  args_.AddObjectRef(tensor_name, AccessType::READ, std::move(desc_new));
+}
+
+void GPUOperation::AddSrcBuffer(const std::string& buffer_name,
+                                const BufferDescriptor& desc) {
+  src_tensors_names_.push_back(buffer_name);
+  auto desc_new = std::make_unique<BufferDescriptor>(desc);
+  args_.AddObjectRef(buffer_name, AccessType::READ, std::move(desc_new));
+}
+
+void GPUOperation::AddDstTensor(const std::string& tensor_name,
+                                const TensorDescriptor& desc) {
+  dst_tensors_names_.push_back(tensor_name);
+  auto desc_new = std::make_unique<TensorDescriptor>(desc);
+  args_.AddObjectRef(tensor_name, AccessType::WRITE, std::move(desc_new));
+}
+
+absl::Status GPUOperation::AssembleCode(const GpuInfo& gpu_info) {
+  if (elementwise_inputs_ == 2) {
+    RETURN_IF_ERROR(ResolveSecondElementwiseInput());
+  }
+  if (elementwise_) {
+    src_tensors_names_.insert(src_tensors_names_.begin(), "src_tensor");
+    args_.AddObjectRef(
+        "src_tensor", AccessType::READ,
+        std::make_unique<TensorDescriptor>(definition_.src_tensors[0]));
+
+    dst_tensors_names_.insert(dst_tensors_names_.begin(), "dst_tensor");
+    args_.AddObjectRef(
+        "dst_tensor", AccessType::WRITE,
+        std::make_unique<TensorDescriptor>(definition_.dst_tensors[0]));
+
+    code_ = GetElementWiseCode(definition_);
+    elementwise_ = false;
+  }
+  std::map<std::string, LinkableContext> linkables;
+  if (!elementwise_code_.empty()) {
+    TensorDescriptor* dst_tensor_desc;
+    RETURN_IF_ERROR(
+        GetTensorDescriptor(dst_tensors_names_[0], &dst_tensor_desc));
+    linkables[dst_tensors_names_[0]] = {elementwise_code_, dst_tensor_desc};
+  }
+  RETURN_IF_ERROR(ResolveConstExprPass(gpu_info, args_, &code_));
+  RETURN_IF_ERROR(ResolveSelectorsPass(gpu_info, linkables, args_, &code_));
+  code_ = GetStringWithoutComments(code_);
+  RETURN_IF_ERROR(args_.Compile(gpu_info, &code_));
+  CalculateConstArgsSize();
+  return absl::OkStatus();
+}
+
+void GPUOperation::RecalculateWorkGroupsCount() {
+  work_groups_count_ = GetWorkGroupsCountInternal(
+      grid_dimension_, grid_size_, work_group_size_, work_group_launch_order_);
+}
+
+void GPUOperation::CalculateConstArgsSize() {
+  const_args_size_ = 0;
+  for (const auto& obj : args_.GetObjects()) {
+    const_args_size_ += obj.second->GetSizeInBytes();
+  }
+}
+
+void GPUOperation::GetPossibleDispatches(
+    TuningType tuning_type, const GpuInfo& gpu_info,
+    const KernelInfo& kernel_info,
+    std::vector<DispatchInfo>* dispatches) const {
+  std::vector<int3> work_group_sizes;
+  GetPossibleKernelWorkGroups(tuning_type, gpu_info, kernel_info,
+                              &work_group_sizes);
+  dispatches->resize(work_group_sizes.size());
+  for (int i = 0; i < work_group_sizes.size(); ++i) {
+    auto& dispatch_info = (*dispatches)[i];
+    dispatch_info.work_group_size = work_group_sizes[i];
+    dispatch_info.work_groups_count = GetWorkGroupsCountInternal(
+        grid_dimension_, grid_size_, work_group_sizes[i],
+        work_group_launch_order_);
+  }
+}
+
+void GPUOperation::GetPossibleKernelWorkGroups(
+    TuningType tuning_type, const GpuInfo& gpu_info,
+    const KernelInfo& kernel_info, std::vector<int3>* work_groups) const {
+  GetPossibleWorkGroups(tuning_type, gpu_info, kernel_info, grid_size_,
+                        work_groups);
+}
+
+int3 GPUOperation::GetGridSize() const {
+  if (tensor_to_grid_ == TensorToGrid::kWBToX_HDToY_SToZ) {
+    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);
+  }
+  if (tensor_to_grid_ == TensorToGrid::kWBToX_HDToY_ZIs1) {
+    const int grid_x = dst_[0]->Width() * dst_[0]->Batch();
+    const int grid_y = dst_[0]->Height() * dst_[0]->Depth();
+    const int grid_z = 1;
+    return int3(grid_x, grid_y, grid_z);
+  }
+  if (tensor_to_grid_ == TensorToGrid::kWBToX_HToY_DToZ) {
+    const int grid_x = dst_[0]->Width() * dst_[0]->Batch();
+    const int grid_y = dst_[0]->Height();
+    const int grid_z = dst_[0]->Depth();
+    return int3(grid_x, grid_y, grid_z);
+  }
+  if (tensor_to_grid_ == TensorToGrid::kBToX_YIs1_ZIs1) {
+    const int grid_x = dst_[0]->Batch();
+    const int grid_y = 1;
+    const int grid_z = 1;
+    return int3(grid_x, grid_y, grid_z);
+  }
+  return grid_size_;
+}
+
+GPUOperation CreateGpuOperation(const OperationDef& definition,
+                                ElementwiseDescriptor&& descriptor) {
+  const BHWC second_shape(2, 2, 2, 2);  // dummy non-broadcasted shape
+  return CreateGpuOperation(definition, std::move(descriptor), second_shape);
+}
+
+GPUOperation CreateGpuOperation(const OperationDef& definition,
+                                ElementwiseDescriptor&& descriptor,
+                                const BHWC& second_shape) {
+  GPUOperation op(definition);
+  op.elementwise_code_ = std::move(descriptor.code);
+  op.elementwise_ = true;
+  if (definition.src_tensors.size() > 1 &&
+      absl::StrContains(op.elementwise_code_, "in2_value")) {
+    const auto second_tensor_def = definition.src_tensors[1];
+    if (NeedsBroadcast(second_tensor_def, second_shape)) {
+      const std::string x_coord = second_shape.w == 1 ? "0" : "X_COORD";
+      const std::string y_coord = second_shape.h == 1 ? "0" : "Y_COORD";
+      const std::string s_coord = second_shape.c == 1 ? "0" : "S_COORD";
+      std::string coords = absl::StrCat(x_coord, ", ", y_coord, ", ", s_coord);
+      if (second_tensor_def.HasAxis(Axis::BATCH)) {
+        const std::string b_coord = second_shape.b == 1 ? "0" : "B_COORD";
+        coords += ", " + b_coord;
+      }
+      std::string read_value_code = absl::StrCat(
+          "args.src_tensor_1::type in2_value = args.src_tensor_1.Read(", coords,
+          ");\n");
+      if (second_shape.c == 1) {
+        read_value_code += "  in2_value.y = in2_value.x;\n";
+        read_value_code += "  in2_value.z = in2_value.x;\n";
+        read_value_code += "  in2_value.w = in2_value.x;\n";
+      }
+      op.elementwise_code_ =
+          "$0{" + read_value_code + op.elementwise_code_ + "}";
+      op.elementwise_code_ = absl::StrReplaceAll(
+          op.elementwise_code_, {{"in2_value", "second_value"}});
+      op.elementwise_inputs_ = 1;
+    } else {
+      op.elementwise_code_ =
+          "$0{READ_SECOND_VALUE" + op.elementwise_code_ + "}";
+      op.elementwise_inputs_ = 2;
+      op.second_elementwise_tensor_name_ = "src_tensor_1";
+    }
+  } else {
+    op.elementwise_code_ = "$0{" + op.elementwise_code_ + "}";
+    op.elementwise_inputs_ = 1;
+  }
+  op.args_ = std::move(descriptor.args);
+  for (int i = 1; i < definition.src_tensors.size(); ++i) {
+    const std::string tensor_name = "src_tensor_" + std::to_string(i);
+    op.AddSrcTensor(tensor_name, definition.src_tensors[i]);
+  }
+  op.tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+  return op;
+}
+
+absl::Status FuseElemWithElemInternal(
+    const GpuInfo& gpu_info, GPUOperation&& elem0, GPUOperation&& elem1,
+    const std::vector<std::pair<std::string, std::string>>& replacements,
+    GPUOperation* result) {
+  const int linkable_count =
+      std::max(elem0.linkable_count_, elem1.linkable_count_) + 1;
+
+  const std::string unique_postfix = absl::StrCat("_link", linkable_count);
+  elem1.args_.RenameArgs(unique_postfix, &elem1.elementwise_code_);
+
+  const auto link_value_type = elem0.definition_.dst_tensors[0].GetDataType();
+  const std::string link_value_name = "interm_value" + unique_postfix;
+  const std::string value_declaration =
+      "\n" + GetTypeDeclaration(gpu_info, link_value_type, 4) + " " +
+      link_value_name + ";\n";
+  elem0.elementwise_code_ = absl::StrReplaceAll(
+      elem0.elementwise_code_, {{"out_value", link_value_name}});
+  elem0.elementwise_code_ =
+      absl::Substitute(elem0.elementwise_code_, value_declaration);
+
+  std::vector<std::pair<const absl::string_view, std::string>> replacements_new;
+  for (int i = 0; i < replacements.size(); ++i) {
+    if (replacements[i].second == "LINK_VALUE") {
+      replacements_new.push_back(
+          {absl::string_view(replacements[i].first), link_value_name});
+    } else {
+      replacements_new.push_back(
+          {absl::string_view(replacements[i].first), replacements[i].second});
+    }
+  }
+  elem1.elementwise_code_ =
+      absl::StrReplaceAll(elem1.elementwise_code_,
+                          {{"interm_value", "interm_value" + unique_postfix}});
+  elem1.elementwise_code_ =
+      absl::StrReplaceAll(elem1.elementwise_code_, replacements_new);
+
+  OperationDef new_definition = elem0.definition_;
+  new_definition.dst_tensors[0] = elem1.definition_.dst_tensors[0];
+
+  *result = GPUOperation(new_definition);
+  result->elementwise_ = true;
+  result->elementwise_inputs_ = 1;
+  result->tensor_to_grid_ = TensorToGrid::kWBToX_HDToY_SToZ;
+  result->elementwise_code_ =
+      elem0.elementwise_code_ + "\n" + elem1.elementwise_code_;
+  result->linkable_count_ = linkable_count;
+  result->args_ = std::move(elem0.args_);
+  RETURN_IF_ERROR(result->args_.Merge(std::move(elem1.args_), unique_postfix,
+                                      {elem1.second_elementwise_tensor_name_}));
+  return absl::OkStatus();
+}
+
+absl::Status FuseSimpleElemWithSimpleElem(const GpuInfo& gpu_info,
+                                          GPUOperation&& elem0,
+                                          GPUOperation&& elem1,
+                                          GPUOperation* result) {
+  return FuseElemWithElemInternal(gpu_info, std::move(elem0), std::move(elem1),
+                                  {{"in_value", "LINK_VALUE"}}, result);
+}
+
+absl::Status Fuse2InputElemWithSimpleElemAsFirstInput(const GpuInfo& gpu_info,
+                                                      GPUOperation&& elem0,
+                                                      GPUOperation&& elem1,
+                                                      GPUOperation* result) {
+  return FuseElemWithElemInternal(gpu_info, std::move(elem0), std::move(elem1),
+                                  {{"in_value", "LINK_VALUE"},
+                                   {"READ_SECOND_VALUE", ""},
+                                   {"in2_value", "in_value"}},
+                                  result);
+}
+
+absl::Status Fuse2InputElemWithSimpleElemAsSecondInput(const GpuInfo& gpu_info,
+                                                       GPUOperation&& elem0,
+                                                       GPUOperation&& elem1,
+                                                       GPUOperation* result) {
+  return FuseElemWithElemInternal(
+      gpu_info, std::move(elem0), std::move(elem1),
+      {{"READ_SECOND_VALUE", ""}, {"in2_value", "LINK_VALUE"}}, result);
+}
+
+//      input                input           input
+//     /    \               /    \             |
+//  elem0  elem1           |    elem1          |
+//     \    /      -->      \    /      -->  elem
+//   elem_root              elem2              |
+//       |                    |                |
+//     output               output           output
+absl::Status Fuse2InputElemWith2SimpleElem(const GpuInfo& gpu_info,
+                                           GPUOperation&& elem0,
+                                           GPUOperation&& elem1,
+                                           GPUOperation&& elem_root,
+                                           GPUOperation* result) {
+  elem0.linkable_count_ =
+      std::max(elem0.linkable_count_, elem1.linkable_count_);
+  elem0.linkable_count_ =
+      std::max(elem0.linkable_count_, elem_root.linkable_count_);
+  GPUOperation elem2;
+  RETURN_IF_ERROR(
+      FuseElemWithElemInternal(gpu_info, std::move(elem0), std::move(elem_root),
+                               {{"in_value", "LINK_VALUE"}}, &elem2));
+  return FuseElemWithElemInternal(
+      gpu_info, std::move(elem1), std::move(elem2),
+      {{"READ_SECOND_VALUE", ""}, {"in2_value", "LINK_VALUE"}}, result);
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/task/gpu_operation.h

@@ -0,0 +1,278 @@
+/* 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_TASK_GPU_OPERATION_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_GPU_OPERATION_H_
+
+#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/kernel_info.h"
+#include "tensorflow/lite/delegates/gpu/common/precision.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/arguments.h"
+#include "tensorflow/lite/delegates/gpu/common/task/buffer_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/compiler_options.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_tensor.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tensor_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tflite_serialization_base_generated.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tuning_type.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+// kCustom: default value
+//   GPUOperation::GetGridSize must be overloaded
+// kWBToX_HDToY_SToZ:
+//   grid_x = dst_[0]->Width() * dst_[0]->Batch();
+//   grid_y = dst_[0]->Height() * dst_[0]->Depth();
+//   grid_z = dst_[0]->Slices();
+// kWBToX_HDToY_ZIs1:
+//   grid_x = dst_[0]->Width() * dst_[0]->Batch();
+//   grid_y = dst_[0]->Height() * dst_[0]->Depth();
+//   grid_z = 1;
+// kWBToX_HToY_DToZ:
+//   grid_x = dst_[0]->Width() * dst_[0]->Batch();
+//   grid_y = dst_[0]->Height();
+//   grid_z = dst_[0]->Depth();
+// kBToX_YIs1_ZIs1:
+//   grid_x = dst_[0]->Batch();
+//   grid_y = 1;
+//   grid_z = 1;
+enum class TensorToGrid {
+  kCustom,
+  kWBToX_HDToY_SToZ,
+  kWBToX_HDToY_ZIs1,
+  kWBToX_HToY_DToZ,
+  kBToX_YIs1_ZIs1
+};
+
+struct OperationDef {
+  CalculationsPrecision precision;
+  std::vector<TensorDescriptor> src_tensors;
+  std::vector<TensorDescriptor> dst_tensors;
+
+  // returns FLOAT32 for F32 precision and FLOAT16 for F16 precision
+  DataType GetDataType() const;
+  // Primary means the first src tensor, because first tensor usually defines
+  // the structure of kernel, all other resources(biases) types and etc.
+  DataType GetPrimaryDataType() const;
+  TensorStorageType GetPrimaryStorageType() const;
+  bool IsBatchSupported() const;
+};
+
+struct ElementwiseDescriptor {
+  Arguments args;
+  std::string code;
+};
+
+class GPUOperation {
+ public:
+  GPUOperation() = default;
+  explicit GPUOperation(const OperationDef& definition);
+  virtual ~GPUOperation() = default;
+  // Move only
+  GPUOperation(GPUOperation&& operation);
+  GPUOperation& operator=(GPUOperation&& operation);
+  GPUOperation(const GPUOperation&) = delete;
+  GPUOperation& operator=(const GPUOperation&) = delete;
+
+  absl::Status AddOperation(const GpuInfo& gpu_info, GPUOperation* operation);
+
+  int GetElementwiseInputsCount() const { return elementwise_inputs_; }
+
+  void SetSrc(GpuSpatialTensor* ptr, int index = 0);
+  void SetDst(GpuSpatialTensor* ptr, int index = 0);
+
+  struct DispatchInfo {
+    int3 work_group_size;
+    int3 work_groups_count;
+  };
+  void GetPossibleDispatches(TuningType tuning_type, const GpuInfo& gpu_info,
+                             const KernelInfo& kernel_info,
+                             std::vector<DispatchInfo>* dispatches) const;
+
+  const std::vector<std::string>& GetSrcTensorsNames() const {
+    return src_tensors_names_;
+  }
+  const std::vector<std::string>& GetDstTensorsNames() const {
+    return dst_tensors_names_;
+  }
+  const std::vector<GpuSpatialTensor*>& GetSrcTensors() const { return src_; }
+  const std::vector<GpuSpatialTensor*>& GetDstTensors() const { return dst_; }
+  const int3& GetWorkGroupsCount() const { return work_groups_count_; }
+
+  absl::Status AssembleCode(const GpuInfo& gpu_info);
+
+  const OperationDef& GetDefinition() const { return definition_; }
+  CalculationsPrecision GetPrecision() const { return definition_.precision; }
+
+  void AddSrcTensor(const std::string& tensor_name,
+                    const TensorDescriptor& desc);
+  void AddSrcBuffer(const std::string& buffer_name,
+                    const BufferDescriptor& desc);
+  void AddDstTensor(const std::string& tensor_name,
+                    const TensorDescriptor& desc);
+
+  bool IsLinkable() const { return elementwise_; }
+
+  virtual absl::Status BindArguments(ArgumentsBinder* args) {
+    return absl::OkStatus();
+  }
+  void RecalculateGridSize() { grid_size_ = GetGridSize(); }
+  void RecalculateWorkGroupsCount();
+
+  Arguments args_;
+  std::string code_;
+  int3 work_group_size_ = int3(8, 4, 1);
+  std::vector<CompilerOptions> compiler_options_;
+  // not applicable to elementwise
+  TensorToGrid tensor_to_grid_ = TensorToGrid::kCustom;
+
+  // for profiling
+  uint64_t flops_ = 0;
+  // size in bytes of constant gpu_objects inside args_
+  uint64_t const_args_size_ = 0;
+
+  // Must be called before const generic objects in args_ released.
+  void CalculateConstArgsSize();
+
+ protected:
+  friend flatbuffers::Offset<tflite::gpu::data::GPUOperation> Encode(
+      const GPUOperation& op, flatbuffers::FlatBufferBuilder* builder);
+  friend absl::Status Decode(const tflite::gpu::data::GPUOperation* fb_op,
+                             GPUOperation* op);
+  friend GPUOperation CreateGpuOperation(const OperationDef& definition,
+                                         ElementwiseDescriptor&& descriptor);
+  friend GPUOperation CreateGpuOperation(const OperationDef& definition,
+                                         ElementwiseDescriptor&& descriptor,
+                                         const BHWC& second_shape);
+
+  friend absl::Status FuseElemWithElemInternal(
+      const GpuInfo& gpu_info, GPUOperation&& elem0, GPUOperation&& elem1,
+      const std::vector<std::pair<std::string, std::string>>& replacements,
+      GPUOperation* result);
+  friend absl::Status FuseSimpleElemWithSimpleElem(const GpuInfo& gpu_info,
+                                                   GPUOperation&& elem0,
+                                                   GPUOperation&& elem1,
+                                                   GPUOperation* result);
+  friend absl::Status Fuse2InputElemWithSimpleElemAsFirstInput(
+      const GpuInfo& gpu_info, GPUOperation&& elem0, GPUOperation&& elem1,
+      GPUOperation* result);
+  friend absl::Status Fuse2InputElemWithSimpleElemAsSecondInput(
+      const GpuInfo& gpu_info, GPUOperation&& elem0, GPUOperation&& elem1,
+      GPUOperation* result);
+  friend absl::Status Fuse2InputElemWith2SimpleElem(const GpuInfo& gpu_info,
+                                                    GPUOperation&& elem0,
+                                                    GPUOperation&& elem1,
+                                                    GPUOperation&& elem_root,
+                                                    GPUOperation* result);
+
+  virtual int3 GetGridSize() const;
+  virtual void GetPossibleKernelWorkGroups(
+      TuningType tuning_type, const GpuInfo& gpu_info,
+      const KernelInfo& kernel_info, std::vector<int3>* work_groups) const;
+
+  // Defines operation calculation precision and format of src/dst tensors.
+  OperationDef definition_;
+  std::vector<GpuSpatialTensor*> src_;
+  std::vector<GpuSpatialTensor*> dst_;
+  int grid_dimension_ = 3;  // can be 1, 2 or 3
+  int3 work_group_launch_order_ = int3(0, 1, 2);
+  int3 grid_size_ = int3(0, 0, 0);
+  std::vector<std::string> src_tensors_names_;
+  std::vector<std::string> dst_tensors_names_;
+
+ private:
+  absl::Status GetTensorDescriptor(const std::string& tensor_name,
+                                   TensorDescriptor** resutl);
+  absl::Status ResolveSecondElementwiseInput();
+  int3 work_groups_count_ = int3(0, 0, 0);
+  bool elementwise_ = false;      // temporary, used during op construction
+  int elementwise_inputs_ = 0;    // can be {0, 1, 2}
+  std::string
+      second_elementwise_tensor_name_;  // used with elementwise_inputs_ = 2
+  int linkable_count_ = 0;        // temporary, used during op construction
+  std::string elementwise_code_;  // temporary, used during op construction
+};
+
+GPUOperation CreateGpuOperation(const OperationDef& definition,
+                                ElementwiseDescriptor&& descriptor);
+
+// For creating elementwise operations with 2 runtime inputs
+GPUOperation CreateGpuOperation(const OperationDef& definition,
+                                ElementwiseDescriptor&& descriptor,
+                                const BHWC& second_shape);
+
+absl::Status FuseElemWithElemInternal(
+    const GpuInfo& gpu_info, GPUOperation&& elem0, GPUOperation&& elem1,
+    const std::vector<std::pair<std::string, std::string>>& replacements,
+    GPUOperation* result);
+
+//    input       input
+//      |           |
+//    elem0         |
+//      |    -->  elem
+//    elem1         |
+//      |           |
+//    output      output
+absl::Status FuseSimpleElemWithSimpleElem(const GpuInfo& gpu_info,
+                                          GPUOperation&& elem0,
+                                          GPUOperation&& elem1,
+                                          GPUOperation* result);
+
+//      input           input
+//     /    \             |
+//  elem0    |            |
+//     \    /      -->  elem
+//     elem1              |
+//       |                |
+//     output           output
+absl::Status Fuse2InputElemWithSimpleElemAsFirstInput(const GpuInfo& gpu_info,
+                                                      GPUOperation&& elem0,
+                                                      GPUOperation&& elem1,
+                                                      GPUOperation* result);
+
+//      input           input
+//     /    \             |
+//    |    elem0          |
+//     \    /      -->  elem
+//     elem1              |
+//       |                |
+//     output           output
+absl::Status Fuse2InputElemWithSimpleElemAsSecondInput(const GpuInfo& gpu_info,
+                                                       GPUOperation&& elem0,
+                                                       GPUOperation&& elem1,
+                                                       GPUOperation* result);
+
+//      input           input
+//     /    \             |
+//  elem0  elem1          |
+//     \    /      -->  elem
+//   elem_root            |
+//       |                |
+//     output           output
+absl::Status Fuse2InputElemWith2SimpleElem(const GpuInfo& gpu_info,
+                                           GPUOperation&& elem0,
+                                           GPUOperation&& elem1,
+                                           GPUOperation&& elem_root,
+                                           GPUOperation* result);
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_GPU_OPERATION_H_

tensorflow/lite/delegates/gpu/common/task/gpu_tensor.h

@@ -0,0 +1,46 @@
+/* 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_TASK_GPU_TENSOR_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_GPU_TENSOR_H_
+
+#include "tensorflow/lite/delegates/gpu/common/task/tensor_desc.h"
+
+namespace tflite {
+namespace gpu {
+
+// Interface for GpuSpatialTensor.
+// Spatial means that it has Width/Height/Depth dimensions(or their combination)
+// and Channels dimension
+// Batch dimension optional
+class GpuSpatialTensor {
+ public:
+  GpuSpatialTensor() = default;
+  virtual ~GpuSpatialTensor() = default;
+
+  virtual int Width() const = 0;
+  virtual int Height() const = 0;
+  virtual int Depth() const = 0;
+  virtual int Channels() const = 0;
+  virtual int Slices() const = 0;
+  virtual int Batch() const = 0;
+
+  virtual TensorDescriptor GetDescriptor() const = 0;
+};
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_GPU_TENSOR_H_

tensorflow/lite/delegates/gpu/common/task/profiling_info.cc

@@ -0,0 +1,89 @@
+/* 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/task/profiling_info.h"
+
+#include <cstdint>
+#include <map>
+#include <string>
+
+#include "absl/time/time.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Duration ProfilingInfo::GetTotalTime() const {
+  absl::Duration total_time;
+  for (const auto& dispatch : dispatches) {
+    total_time += dispatch.duration;
+  }
+  return total_time;
+}
+
+std::string ProfilingInfo::GetDetailedReport() const {
+  std::string result;
+  struct OpStatistic {
+    int count;
+    double total_time;
+  };
+  std::map<std::string, OpStatistic> statistics;
+  result +=
+      "Per kernel timing(" + std::to_string(dispatches.size()) + " kernels):\n";
+  for (const auto& dispatch : dispatches) {
+    result += "  " + dispatch.label + " - " +
+              std::to_string(absl::ToDoubleMilliseconds(dispatch.duration)) +
+              " ms";
+    const double times_per_sec =
+        1000.0 / absl::ToDoubleMilliseconds(dispatch.duration);
+    if (dispatch.read_mem_size && dispatch.write_mem_size) {
+      const uint64_t total_size =
+          dispatch.read_mem_size + dispatch.write_mem_size;
+      const double giga_bytes = total_size / 1024.0 / 1024.0 / 1024.0;
+      const double giga_bytes_per_sec = times_per_sec * giga_bytes;
+      result += ", " + std::to_string(giga_bytes_per_sec) + " Gbyte/s";
+    }
+    if (dispatch.flops) {
+      const double giga_flops = dispatch.flops / 1000.0 / 1000.0 / 1000.0;
+      const double giga_flops_per_sec = times_per_sec * giga_flops;
+      result += ", " + std::to_string(giga_flops_per_sec) + " Gflops";
+    }
+    result += "\n";
+    auto name = dispatch.label.substr(0, dispatch.label.find(' '));
+    if (statistics.find(name) != statistics.end()) {
+      statistics[name].count++;
+      statistics[name].total_time +=
+          absl::ToDoubleMilliseconds(dispatch.duration);
+    } else {
+      statistics[name].count = 1;
+      statistics[name].total_time =
+          absl::ToDoubleMilliseconds(dispatch.duration);
+    }
+  }
+  result += "--------------------\n";
+  result += "Accumulated time per operation type:\n";
+  for (auto& t : statistics) {
+    auto stat = t.second;
+    result += "  " + t.first + "(x" + std::to_string(stat.count) + ") - " +
+              std::to_string(stat.total_time) + " ms\n";
+  }
+  result += "--------------------\n";
+  result += "Ideal total time: " +
+            std::to_string(absl::ToDoubleMilliseconds(GetTotalTime())) + "\n";
+  result += "--------------------\n";
+  return result;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/task/profiling_info.h

@@ -0,0 +1,58 @@
+/* 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_TASK_PROFILING_INFO_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_PROFILING_INFO_H_
+
+#include <string>
+#include <vector>
+
+#include "absl/time/time.h"
+
+namespace tflite {
+namespace gpu {
+
+struct ProfilingInfo {
+  struct DispatchInfo {
+    std::string label;
+    absl::Duration duration;
+    uint64_t read_mem_size = 0;
+    uint64_t write_mem_size = 0;
+    uint64_t flops = 0;
+  };
+
+  std::vector<DispatchInfo> dispatches;
+
+  absl::Duration GetTotalTime() const;
+
+  // Returns report (string of lines delimited by \n)
+  // This method uses GPU counters and measure GPU time only.
+  // Report has next structure:
+  // Per kernel timing(K kernels):
+  //   conv2d 3.2ms
+  //   ...
+  // --------------------
+  // Accumulated time per operation type:
+  //   conv2d - 14.5ms
+  //   ....
+  // --------------------
+  // Ideal total time: 23.4ms // Total time for all kernels
+  std::string GetDetailedReport() const;
+};
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_PROFILING_INFO_H_

tensorflow/lite/delegates/gpu/common/task/qcom_thin_filter_desc.cc

@@ -0,0 +1,51 @@
+/* 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/task/qcom_thin_filter_desc.h"
+
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+
+namespace tflite {
+namespace gpu {
+
+GPUResources QcomThinFilterDescriptor::GetGPUResources(
+    const GpuInfo& gpu_info) const {
+  GPUResources resources;
+  GPUCustomMemoryDescriptor desc;
+  desc.type_name = "__read_only qcom_weight_image_t";
+  resources.custom_memories.push_back({"filter", desc});
+  return resources;
+}
+
+absl::Status QcomThinFilterDescriptor::PerformSelector(
+    const GpuInfo& gpu_info, absl::string_view selector,
+    const std::vector<std::string>& args,
+    const std::vector<std::string>& template_args, std::string* result) const {
+  if (selector == "GetHandle" && args.empty()) {
+    *result = "filter";
+    return absl::OkStatus();
+  } else {
+    return absl::NotFoundError(absl::StrCat(
+        "QcomThinFilterDescriptor don't have selector with name - ", selector));
+  }
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/task/qcom_thin_filter_desc.h

@@ -0,0 +1,57 @@
+/* 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_TASK_QCOM_THIN_FILTER_DESC_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_QCOM_THIN_FILTER_DESC_H_
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/strings/string_view.h"
+#include "tensorflow/lite/delegates/gpu/common/status.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_object_desc.h"
+
+namespace tflite {
+namespace gpu {
+
+struct QcomThinFilterDescriptor : public GPUObjectDescriptor {
+  int kernel_size_x;
+  int kernel_size_y;
+  std::vector<uint8_t> data;
+
+  QcomThinFilterDescriptor() = default;
+  QcomThinFilterDescriptor(const QcomThinFilterDescriptor&) = default;
+  QcomThinFilterDescriptor& operator=(const QcomThinFilterDescriptor&) =
+      default;
+  QcomThinFilterDescriptor(QcomThinFilterDescriptor&& desc) = default;
+  QcomThinFilterDescriptor& operator=(QcomThinFilterDescriptor&& desc) =
+      default;
+
+  absl::Status PerformSelector(const GpuInfo& gpu_info,
+                               absl::string_view selector,
+                               const std::vector<std::string>& args,
+                               const std::vector<std::string>& template_args,
+                               std::string* result) const override;
+
+  GPUResources GetGPUResources(const GpuInfo& gpu_info) const override;
+
+  void Release() override { data.clear(); }
+};
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_QCOM_THIN_FILTER_DESC_H_

tensorflow/lite/delegates/gpu/common/task/serialization_base.cc

@@ -0,0 +1,725 @@
+/* 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/task/serialization_base.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
+#include "tensorflow/lite/delegates/gpu/common/precision.h"
+#include "tensorflow/lite/delegates/gpu/common/task/arguments.h"
+#include "tensorflow/lite/delegates/gpu/common/task/buffer_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_object_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/tflite_serialization_base_generated.h"
+
+namespace tflite {
+namespace gpu {
+
+namespace {
+data::AccessType ToFB(AccessType type) {
+  switch (type) {
+    case AccessType::READ:
+      return data::AccessType::READ;
+    case AccessType::WRITE:
+      return data::AccessType::WRITE;
+    case AccessType::READ_WRITE:
+      return data::AccessType::READ_WRITE;
+    default:
+      return data::AccessType::READ_WRITE;
+  }
+}
+
+data::DataType ToFB(DataType type) {
+  switch (type) {
+    case DataType::BOOL:
+      return data::DataType::BOOL;
+    case DataType::FLOAT16:
+      return data::DataType::FLOAT16;
+    case DataType::FLOAT32:
+      return data::DataType::FLOAT32;
+    case DataType::FLOAT64:
+      return data::DataType::FLOAT64;
+    case DataType::UINT8:
+      return data::DataType::UINT8;
+    case DataType::INT8:
+      return data::DataType::INT8;
+    case DataType::UINT16:
+      return data::DataType::UINT16;
+    case DataType::INT16:
+      return data::DataType::INT16;
+    case DataType::UINT32:
+      return data::DataType::UINT32;
+    case DataType::INT32:
+      return data::DataType::INT32;
+    case DataType::UINT64:
+      return data::DataType::UINT64;
+    case DataType::INT64:
+      return data::DataType::INT64;
+    case DataType::UNKNOWN:
+      return data::DataType::UNKNOWN;
+  }
+}
+
+data::MemoryType ToFB(MemoryType type) {
+  switch (type) {
+    case MemoryType::CONSTANT:
+      return data::MemoryType::CONSTANT;
+    case MemoryType::GLOBAL:
+      return data::MemoryType::GLOBAL;
+    case MemoryType::LOCAL:
+      return data::MemoryType::LOCAL;
+  }
+}
+
+data::TensorStorageType ToFB(TensorStorageType type) {
+  switch (type) {
+    case TensorStorageType::BUFFER:
+      return data::TensorStorageType::BUFFER;
+    case TensorStorageType::IMAGE_BUFFER:
+      return data::TensorStorageType::IMAGE_BUFFER;
+    case TensorStorageType::TEXTURE_2D:
+      return data::TensorStorageType::TEXTURE_2D;
+    case TensorStorageType::TEXTURE_ARRAY:
+      return data::TensorStorageType::TEXTURE_ARRAY;
+    case TensorStorageType::TEXTURE_3D:
+      return data::TensorStorageType::TEXTURE_3D;
+    case TensorStorageType::SINGLE_TEXTURE_2D:
+      return data::TensorStorageType::SINGLE_TEXTURE_2D;
+    case TensorStorageType::UNKNOWN:
+      return data::TensorStorageType::UNKNOWN;
+  }
+}
+
+data::Layout ToFB(Layout type) {
+  switch (type) {
+    case Layout::HWC:
+      return data::Layout::HWC;
+    case Layout::BHWC:
+      return data::Layout::BHWC;
+    case Layout::HWDC:
+      return data::Layout::HWDC;
+    case Layout::BHWDC:
+      return data::Layout::BHWDC;
+    case Layout::LINEAR:
+      return data::Layout::LINEAR;
+    case Layout::HW:
+      return data::Layout::HW;
+    default:
+      return data::Layout::UNKNOWN;
+  }
+}
+
+DataType ToEnum(data::DataType type) {
+  switch (type) {
+    case data::DataType::BOOL:
+      return DataType::BOOL;
+    case data::DataType::FLOAT16:
+      return DataType::FLOAT16;
+    case data::DataType::FLOAT32:
+      return DataType::FLOAT32;
+    case data::DataType::FLOAT64:
+      return DataType::FLOAT64;
+    case data::DataType::UINT8:
+      return DataType::UINT8;
+    case data::DataType::INT8:
+      return DataType::INT8;
+    case data::DataType::UINT16:
+      return DataType::UINT16;
+    case data::DataType::INT16:
+      return DataType::INT16;
+    case data::DataType::UINT32:
+      return DataType::UINT32;
+    case data::DataType::INT32:
+      return DataType::INT32;
+    case data::DataType::UINT64:
+      return DataType::UINT64;
+    case data::DataType::INT64:
+      return DataType::INT64;
+    case data::DataType::UNKNOWN:
+      return DataType::UNKNOWN;
+  }
+}
+
+AccessType ToEnum(data::AccessType type) {
+  switch (type) {
+    case data::AccessType::READ:
+      return AccessType::READ;
+    case data::AccessType::WRITE:
+      return AccessType::WRITE;
+    case data::AccessType::READ_WRITE:
+      return AccessType::READ_WRITE;
+  }
+}
+
+MemoryType ToEnum(data::MemoryType type) {
+  switch (type) {
+    case data::MemoryType::CONSTANT:
+      return MemoryType::CONSTANT;
+    case data::MemoryType::GLOBAL:
+      return MemoryType::GLOBAL;
+    case data::MemoryType::LOCAL:
+      return MemoryType::LOCAL;
+  }
+}
+
+TensorStorageType ToEnum(data::TensorStorageType type) {
+  switch (type) {
+    case data::TensorStorageType::BUFFER:
+      return TensorStorageType::BUFFER;
+    case data::TensorStorageType::IMAGE_BUFFER:
+      return TensorStorageType::IMAGE_BUFFER;
+    case data::TensorStorageType::TEXTURE_2D:
+      return TensorStorageType::TEXTURE_2D;
+    case data::TensorStorageType::TEXTURE_ARRAY:
+      return TensorStorageType::TEXTURE_ARRAY;
+    case data::TensorStorageType::TEXTURE_3D:
+      return TensorStorageType::TEXTURE_3D;
+    case data::TensorStorageType::SINGLE_TEXTURE_2D:
+      return TensorStorageType::SINGLE_TEXTURE_2D;
+    case data::TensorStorageType::UNKNOWN:
+      return TensorStorageType::UNKNOWN;
+  }
+}
+
+Layout ToEnum(data::Layout type) {
+  switch (type) {
+    case data::Layout::HWC:
+      return Layout::HWC;
+    case data::Layout::BHWC:
+      return Layout::BHWC;
+    case data::Layout::HWDC:
+      return Layout::HWDC;
+    case data::Layout::BHWDC:
+      return Layout::BHWDC;
+    case data::Layout::LINEAR:
+      return Layout::LINEAR;
+    case data::Layout::HW:
+      return Layout::HW;
+    default:
+      return Layout::UNKNOWN;
+  }
+}
+
+data::CalculationsPrecision ToFB(CalculationsPrecision type) {
+  switch (type) {
+    case CalculationsPrecision::F32:
+      return data::CalculationsPrecision::F32;
+    case CalculationsPrecision::F32_F16:
+      return data::CalculationsPrecision::F32_F16;
+    case CalculationsPrecision::F16:
+      return data::CalculationsPrecision::F16;
+  }
+}
+
+data::TensorToGrid ToFB(TensorToGrid type) {
+  switch (type) {
+    case TensorToGrid::kCustom:
+      return data::TensorToGrid::CUSTOM;
+    case TensorToGrid::kWBToX_HDToY_SToZ:
+      return data::TensorToGrid::WB_TO_X_HD_TO_Y_S_TO_Z;
+    case TensorToGrid::kWBToX_HDToY_ZIs1:
+      return data::TensorToGrid::WB_TO_X_HD_TO_Y_Z_IS_1;
+    case TensorToGrid::kWBToX_HToY_DToZ:
+      return data::TensorToGrid::WB_TO_X_H_TO_Y_D_TO_Z;
+    case TensorToGrid::kBToX_YIs1_ZIs1:
+      return data::TensorToGrid::B_TO_X_Y_IS_1_Z_IS_1;
+  }
+}
+
+data::CompilerOptions ToFB(CompilerOptions type) {
+  switch (type) {
+    case CompilerOptions::kAdrenoFullSimd:
+      return data::CompilerOptions::ADRENO_FULL_SIMD_LINE;
+    case CompilerOptions::kAdrenoMoreWaves:
+      return data::CompilerOptions::ADRENO_MORE_WAVES;
+    case CompilerOptions::kClFastRelaxedMath:
+      return data::CompilerOptions::CL_FAST_RELAXED_MATH;
+    case CompilerOptions::kClDisableOptimizations:
+      return data::CompilerOptions::CL_OPT_DISABLE;
+    case CompilerOptions::kCl20:
+      return data::CompilerOptions::CL_2_0;
+    case CompilerOptions::kCl30:
+      return data::CompilerOptions::CL_3_0;
+    case CompilerOptions::kClRegisterAllocation64:
+      return data::CompilerOptions::CL_REGISTER_ALLOCATION_64;
+  }
+}
+
+CalculationsPrecision ToEnum(data::CalculationsPrecision type) {
+  switch (type) {
+    case data::CalculationsPrecision::F32:
+      return CalculationsPrecision::F32;
+    case data::CalculationsPrecision::F32_F16:
+      return CalculationsPrecision::F32_F16;
+    case data::CalculationsPrecision::F16:
+      return CalculationsPrecision::F16;
+  }
+}
+
+TensorToGrid ToEnum(data::TensorToGrid type) {
+  switch (type) {
+    case data::TensorToGrid::CUSTOM:
+      return TensorToGrid::kCustom;
+    case data::TensorToGrid::WB_TO_X_HD_TO_Y_S_TO_Z:
+      return TensorToGrid::kWBToX_HDToY_SToZ;
+    case data::TensorToGrid::WB_TO_X_HD_TO_Y_Z_IS_1:
+      return TensorToGrid::kWBToX_HDToY_ZIs1;
+    case data::TensorToGrid::WB_TO_X_H_TO_Y_D_TO_Z:
+      return TensorToGrid::kWBToX_HToY_DToZ;
+    case data::TensorToGrid::B_TO_X_Y_IS_1_Z_IS_1:
+      return TensorToGrid::kBToX_YIs1_ZIs1;
+  }
+}
+
+CompilerOptions ToEnum(data::CompilerOptions type) {
+  switch (type) {
+    case data::CompilerOptions::ADRENO_FULL_SIMD_LINE:
+      return CompilerOptions::kAdrenoFullSimd;
+    case data::CompilerOptions::ADRENO_MORE_WAVES:
+      return CompilerOptions::kAdrenoMoreWaves;
+    case data::CompilerOptions::CL_FAST_RELAXED_MATH:
+      return CompilerOptions::kClFastRelaxedMath;
+    case data::CompilerOptions::CL_OPT_DISABLE:
+      return CompilerOptions::kClDisableOptimizations;
+    case data::CompilerOptions::CL_2_0:
+      return CompilerOptions::kCl20;
+    case data::CompilerOptions::CL_3_0:
+      return CompilerOptions::kCl30;
+    case data::CompilerOptions::CL_REGISTER_ALLOCATION_64:
+      return CompilerOptions::kClRegisterAllocation64;
+  }
+}
+
+}  // namespace
+
+flatbuffers::Offset<data::Int2> Encode(
+    const int2& v, flatbuffers::FlatBufferBuilder* builder) {
+  data::Int2Builder int2_builder(*builder);
+  int2_builder.add_x(v.x);
+  int2_builder.add_y(v.y);
+  return int2_builder.Finish();
+}
+
+flatbuffers::Offset<data::Int3> Encode(
+    const int3& v, flatbuffers::FlatBufferBuilder* builder) {
+  data::Int3Builder int3_builder(*builder);
+  int3_builder.add_x(v.x);
+  int3_builder.add_y(v.y);
+  int3_builder.add_z(v.z);
+  return int3_builder.Finish();
+}
+
+flatbuffers::Offset<data::GPUObjectDescriptor> Encode(
+    const GPUObjectDescriptor& desc, flatbuffers::FlatBufferBuilder* builder) {
+  std::vector<flatbuffers::Offset<data::StateVariable>> state_vars_fb;
+  for (auto& v0 : desc.state_vars_) {
+    auto key_fb = builder->CreateString(v0.first);
+    auto value_fb = builder->CreateString(v0.second);
+    data::StateVariableBuilder state_builder(*builder);
+    state_builder.add_key(key_fb);
+    state_builder.add_value(value_fb);
+    state_vars_fb.push_back(state_builder.Finish());
+  }
+  auto state_vars_fb_vec = builder->CreateVector(state_vars_fb);
+  data::GPUObjectDescriptorBuilder obj_builder(*builder);
+  obj_builder.add_state_vars(state_vars_fb_vec);
+  obj_builder.add_access_type(ToFB(desc.access_type_));
+  return obj_builder.Finish();
+}
+
+void Decode(const data::GPUObjectDescriptor* fb_obj, GPUObjectDescriptor* obj) {
+  obj->access_type_ = ToEnum(fb_obj->access_type());
+  for (auto state_fb : *fb_obj->state_vars()) {
+    std::string key(state_fb->key()->c_str(), state_fb->key()->size());
+    std::string value(state_fb->value()->c_str(), state_fb->value()->size());
+    obj->state_vars_[key] = value;
+  }
+}
+
+flatbuffers::Offset<data::BufferDescriptor> Encode(
+    const BufferDescriptor& desc, flatbuffers::FlatBufferBuilder* builder) {
+  auto obj_fb =
+      Encode(*static_cast<const GPUObjectDescriptor*>(&desc), builder);
+
+  std::vector<flatbuffers::Offset<flatbuffers::String>> attributes_fb;
+  attributes_fb.reserve(desc.attributes.size());
+  for (auto& attr : desc.attributes) {
+    attributes_fb.push_back(builder->CreateString(attr));
+  }
+  auto attributes_fb_vec = builder->CreateVector(attributes_fb);
+  auto data_fb = builder->CreateVector(desc.data);
+  data::BufferDescriptorBuilder buf_builder(*builder);
+  buf_builder.add_base_obj(obj_fb);
+  buf_builder.add_element_type(ToFB(desc.element_type));
+  buf_builder.add_element_size(desc.element_size);
+  buf_builder.add_memory_type(ToFB(desc.memory_type));
+  buf_builder.add_attributes(attributes_fb_vec);
+  buf_builder.add_size(desc.size);
+  buf_builder.add_data(data_fb);
+  return buf_builder.Finish();
+}
+
+void Decode(const data::BufferDescriptor* fb_desc, BufferDescriptor* desc) {
+  Decode(fb_desc->base_obj(), desc);
+  desc->element_type = ToEnum(fb_desc->element_type());
+  desc->element_size = fb_desc->element_size();
+  desc->memory_type = ToEnum(fb_desc->memory_type());
+  for (auto attr_fb : *fb_desc->attributes()) {
+    std::string attr(attr_fb->c_str(), attr_fb->size());
+    desc->attributes.push_back(attr);
+  }
+  desc->size = fb_desc->size();
+  desc->data =
+      std::vector<uint8_t>(fb_desc->data()->data(),
+                           fb_desc->data()->data() + fb_desc->data()->size());
+}
+
+flatbuffers::Offset<data::TensorDescriptor> Encode(
+    const TensorDescriptor& desc, flatbuffers::FlatBufferBuilder* builder) {
+  auto obj_fb =
+      Encode(*static_cast<const GPUObjectDescriptor*>(&desc), builder);
+
+  data::BHWDCBuilder shape_builder(*builder);
+  shape_builder.add_b(desc.GetBHWDCShape().b);
+  shape_builder.add_h(desc.GetBHWDCShape().h);
+  shape_builder.add_w(desc.GetBHWDCShape().w);
+  shape_builder.add_d(desc.GetBHWDCShape().d);
+  shape_builder.add_c(desc.GetBHWDCShape().c);
+  auto shape_fb = shape_builder.Finish();
+
+  auto data_fb = builder->CreateVector(desc.GetData());
+  data::TensorDescriptorBuilder tensor_builder(*builder);
+  tensor_builder.add_base_obj(obj_fb);
+  tensor_builder.add_data_type(ToFB(desc.data_type_));
+  tensor_builder.add_storage_type(ToFB(desc.storage_type_));
+  tensor_builder.add_layout(ToFB(desc.layout_));
+  tensor_builder.add_shape(shape_fb);
+  tensor_builder.add_data(data_fb);
+  tensor_builder.add_use_buffer_for_write_only_2d_texture(
+      desc.use_buffer_for_write_only_2d_texture_);
+  tensor_builder.add_use_buffer_for_write_only_image_buffer(
+      desc.use_buffer_for_write_only_image_buffer_);
+  return tensor_builder.Finish();
+}
+
+void Decode(const data::TensorDescriptor* fb_desc, TensorDescriptor* desc) {
+  Decode(fb_desc->base_obj(), desc);
+  desc->data_type_ = ToEnum(fb_desc->data_type());
+  desc->storage_type_ = ToEnum(fb_desc->storage_type());
+  desc->layout_ = ToEnum(fb_desc->layout());
+  desc->SetBHWDCShape(BHWDC(fb_desc->shape()->b(), fb_desc->shape()->h(),
+                            fb_desc->shape()->w(), fb_desc->shape()->d(),
+                            fb_desc->shape()->c()));
+  desc->SetData(
+      std::vector<uint8_t>(fb_desc->data()->data(),
+                           fb_desc->data()->data() + fb_desc->data()->size()));
+  desc->use_buffer_for_write_only_2d_texture_ =
+      fb_desc->use_buffer_for_write_only_2d_texture();
+  desc->use_buffer_for_write_only_image_buffer_ =
+      fb_desc->use_buffer_for_write_only_image_buffer();
+}
+
+absl::Status Decode(const data::Arguments* fb_args, Arguments* args) {
+  args->int_values_.clear();
+  for (auto int_values_fb : *fb_args->int_values()) {
+    Arguments::IntValue value;
+    value.value = int_values_fb->value();
+    value.active = int_values_fb->active();
+    std::string name(int_values_fb->name()->c_str(),
+                     int_values_fb->name()->size());
+    args->int_values_[name] = value;
+  }
+
+  args->float_values_.clear();
+  for (auto float_values_fb : *fb_args->float_values()) {
+    Arguments::FloatValue value;
+    value.value = float_values_fb->value();
+    value.active = float_values_fb->active();
+    std::string name(float_values_fb->name()->c_str(),
+                     float_values_fb->name()->size());
+    args->float_values_[name] = value;
+  }
+
+  args->half_values_.clear();
+  for (auto half_values_fb : *fb_args->half_values()) {
+    Arguments::HalfValue value;
+    value.value = half_values_fb->value();
+    value.active = half_values_fb->active();
+    std::string name(half_values_fb->name()->c_str(),
+                     half_values_fb->name()->size());
+    args->half_values_[name] = value;
+  }
+
+  for (auto buffer_pair_fb : *fb_args->buffer_objects()) {
+    std::string key(buffer_pair_fb->key()->c_str(),
+                    buffer_pair_fb->key()->size());
+    BufferDescriptor desc;
+    Decode(buffer_pair_fb->value(), &desc);
+    args->AddObject(key, std::make_unique<BufferDescriptor>(std::move(desc)));
+  }
+
+  for (auto tensor_pair_fb : *fb_args->tensor_objects()) {
+    std::string key(tensor_pair_fb->key()->c_str(),
+                    tensor_pair_fb->key()->size());
+    TensorDescriptor desc;
+    Decode(tensor_pair_fb->value(), &desc);
+    args->AddObject(key, std::make_unique<TensorDescriptor>(std::move(desc)));
+  }
+
+  for (auto buffer_pair_fb : *fb_args->buffer_refs()) {
+    std::string key(buffer_pair_fb->key()->c_str(),
+                    buffer_pair_fb->key()->size());
+    BufferDescriptor desc;
+    Decode(buffer_pair_fb->value(), &desc);
+    auto access_type = desc.GetAccess();
+    args->AddObjectRef(key, access_type,
+                       std::make_unique<BufferDescriptor>(std::move(desc)));
+  }
+
+  for (auto tensor_pair_fb : *fb_args->tensor_refs()) {
+    std::string key(tensor_pair_fb->key()->c_str(),
+                    tensor_pair_fb->key()->size());
+    TensorDescriptor desc;
+    Decode(tensor_pair_fb->value(), &desc);
+    auto access_type = desc.GetAccess();
+    args->AddObjectRef(key, access_type,
+                       std::make_unique<TensorDescriptor>(std::move(desc)));
+  }
+  return absl::OkStatus();
+}
+
+flatbuffers::Offset<data::Arguments> Encode(
+    const Arguments& args, flatbuffers::FlatBufferBuilder* builder) {
+  std::vector<flatbuffers::Offset<data::IntValue>> int_values_fb;
+  for (auto& value : args.int_values_) {
+    auto name_fb = builder->CreateString(value.first);
+    data::IntValueBuilder value_builder(*builder);
+    value_builder.add_name(name_fb);
+    value_builder.add_value(value.second.value);
+    value_builder.add_active(value.second.active);
+    int_values_fb.push_back(value_builder.Finish());
+  }
+
+  std::vector<flatbuffers::Offset<data::FloatValue>> float_values_fb;
+  for (auto& value : args.float_values_) {
+    auto name_fb = builder->CreateString(value.first);
+    data::FloatValueBuilder value_builder(*builder);
+    value_builder.add_name(name_fb);
+    value_builder.add_value(value.second.value);
+    value_builder.add_active(value.second.active);
+    float_values_fb.push_back(value_builder.Finish());
+  }
+
+  std::vector<flatbuffers::Offset<data::HalfValue>> half_values_fb;
+  for (auto& value : args.half_values_) {
+    auto name_fb = builder->CreateString(value.first);
+    data::HalfValueBuilder value_builder(*builder);
+    value_builder.add_name(name_fb);
+    value_builder.add_value(value.second.value);
+    value_builder.add_active(value.second.active);
+    half_values_fb.push_back(value_builder.Finish());
+  }
+
+  std::vector<flatbuffers::Offset<data::BufferDescriptorMapValue>>
+      buffer_objs_fb;
+  for (auto& value : args.objects_) {
+    const auto* buffer_desc =
+        dynamic_cast<const BufferDescriptor*>(value.second.get());
+    if (!buffer_desc) continue;
+    auto desc_fb = Encode(*buffer_desc, builder);
+    auto key_fb = builder->CreateString(value.first);
+    data::BufferDescriptorMapValueBuilder buf_map_builder(*builder);
+    buf_map_builder.add_key(key_fb);
+    buf_map_builder.add_value(desc_fb);
+    buffer_objs_fb.push_back(buf_map_builder.Finish());
+  }
+  std::vector<flatbuffers::Offset<data::TensorDescriptorMapValue>>
+      tensor_objs_fb;
+  for (auto& value : args.objects_) {
+    const auto* tensor_desc =
+        dynamic_cast<const TensorDescriptor*>(value.second.get());
+    if (!tensor_desc) continue;
+    auto desc_fb = Encode(*tensor_desc, builder);
+    auto key_fb = builder->CreateString(value.first);
+    data::TensorDescriptorMapValueBuilder ten_map_builder(*builder);
+    ten_map_builder.add_key(key_fb);
+    ten_map_builder.add_value(desc_fb);
+    tensor_objs_fb.push_back(ten_map_builder.Finish());
+  }
+
+  std::vector<flatbuffers::Offset<data::BufferDescriptorMapValue>>
+      buffer_refs_fb;
+  for (auto& value : args.object_refs_) {
+    const auto* buffer_desc =
+        dynamic_cast<const BufferDescriptor*>(value.second.get());
+    if (!buffer_desc) continue;
+    auto desc_fb = Encode(*buffer_desc, builder);
+    auto key_fb = builder->CreateString(value.first);
+    data::BufferDescriptorMapValueBuilder buf_map_builder(*builder);
+    buf_map_builder.add_key(key_fb);
+    buf_map_builder.add_value(desc_fb);
+    buffer_refs_fb.push_back(buf_map_builder.Finish());
+  }
+  std::vector<flatbuffers::Offset<data::TensorDescriptorMapValue>>
+      tensor_refs_fb;
+  for (auto& value : args.object_refs_) {
+    const auto* tensor_desc =
+        dynamic_cast<const TensorDescriptor*>(value.second.get());
+    if (!tensor_desc) continue;
+    auto desc_fb = Encode(*tensor_desc, builder);
+    auto key_fb = builder->CreateString(value.first);
+    data::TensorDescriptorMapValueBuilder ten_map_builder(*builder);
+    ten_map_builder.add_key(key_fb);
+    ten_map_builder.add_value(desc_fb);
+    tensor_refs_fb.push_back(ten_map_builder.Finish());
+  }
+
+  auto int_values_fb_vec = builder->CreateVector(int_values_fb);
+  auto float_values_fb_vec = builder->CreateVector(float_values_fb);
+  auto half_values_fb_vec = builder->CreateVector(half_values_fb);
+  auto buffer_objs_fb_vec = builder->CreateVector(buffer_objs_fb);
+  auto tensor_objs_fb_vec = builder->CreateVector(tensor_objs_fb);
+  auto buffer_refs_fb_vec = builder->CreateVector(buffer_refs_fb);
+  auto tensor_refs_fb_vec = builder->CreateVector(tensor_refs_fb);
+  data::ArgumentsBuilder arguments_builder(*builder);
+  arguments_builder.add_int_values(int_values_fb_vec);
+  arguments_builder.add_float_values(float_values_fb_vec);
+  arguments_builder.add_half_values(half_values_fb_vec);
+  arguments_builder.add_buffer_objects(buffer_objs_fb_vec);
+  arguments_builder.add_tensor_objects(tensor_objs_fb_vec);
+  arguments_builder.add_buffer_refs(buffer_refs_fb_vec);
+  arguments_builder.add_tensor_refs(tensor_refs_fb_vec);
+  return arguments_builder.Finish();
+}
+
+flatbuffers::Offset<data::OperationDef> Encode(
+    const OperationDef& def, flatbuffers::FlatBufferBuilder* builder) {
+  std::vector<flatbuffers::Offset<tflite::gpu::data::TensorDescriptor>>
+      src_tensors_fb;
+  for (auto& desc : def.src_tensors) {
+    auto desc_fb = Encode(desc, builder);
+    src_tensors_fb.push_back(desc_fb);
+  }
+
+  std::vector<flatbuffers::Offset<tflite::gpu::data::TensorDescriptor>>
+      dst_tensors_fb;
+  for (auto& desc : def.dst_tensors) {
+    auto desc_fb = Encode(desc, builder);
+    dst_tensors_fb.push_back(desc_fb);
+  }
+
+  auto src_tensors_fb_vec = builder->CreateVector(src_tensors_fb);
+  auto dst_tensors_fb_vec = builder->CreateVector(dst_tensors_fb);
+
+  data::OperationDefBuilder def_builder(*builder);
+  def_builder.add_precision(ToFB(def.precision));
+  def_builder.add_src_tensors(src_tensors_fb_vec);
+  def_builder.add_dst_tensors(dst_tensors_fb_vec);
+  return def_builder.Finish();
+}
+
+void Decode(const data::OperationDef* fb_def, OperationDef* def) {
+  for (auto src_fb : *fb_def->src_tensors()) {
+    TensorDescriptor desc;
+    Decode(src_fb, &desc);
+    def->src_tensors.push_back(std::move(desc));
+  }
+  for (auto dst_fb : *fb_def->dst_tensors()) {
+    TensorDescriptor desc;
+    Decode(dst_fb, &desc);
+    def->dst_tensors.push_back(std::move(desc));
+  }
+  def->precision = ToEnum(fb_def->precision());
+}
+
+absl::Status Decode(const data::GPUOperation* fb_op, GPUOperation* op) {
+  RETURN_IF_ERROR(Decode(fb_op->arguments(), &op->args_));
+  op->work_group_size_.x = fb_op->work_group_size()->x();
+  op->work_group_size_.y = fb_op->work_group_size()->y();
+  op->work_group_size_.z = fb_op->work_group_size()->z();
+  op->tensor_to_grid_ = ToEnum(fb_op->tensor_to_grid());
+  op->flops_ = fb_op->flops();
+  Decode(fb_op->definition(), &op->definition_);
+  op->grid_dimension_ = fb_op->grid_dimension();
+  op->work_group_launch_order_.x = fb_op->work_group_launch_order()->x();
+  op->work_group_launch_order_.y = fb_op->work_group_launch_order()->y();
+  op->work_group_launch_order_.z = fb_op->work_group_launch_order()->z();
+  op->grid_size_.x = fb_op->grid_size()->x();
+  op->grid_size_.y = fb_op->grid_size()->y();
+  op->grid_size_.z = fb_op->grid_size()->z();
+  for (auto name_fb : *fb_op->src_tensors_names()) {
+    std::string name(name_fb->c_str(), name_fb->size());
+    op->src_tensors_names_.push_back(std::move(name));
+  }
+  for (auto name_fb : *fb_op->dst_tensors_names()) {
+    std::string name(name_fb->c_str(), name_fb->size());
+    op->dst_tensors_names_.push_back(std::move(name));
+  }
+  op->work_groups_count_.x = fb_op->work_groups_count()->x();
+  op->work_groups_count_.y = fb_op->work_groups_count()->y();
+  op->work_groups_count_.z = fb_op->work_groups_count()->z();
+  op->CalculateConstArgsSize();
+  return absl::OkStatus();
+}
+
+flatbuffers::Offset<data::GPUOperation> Encode(
+    const GPUOperation& op, flatbuffers::FlatBufferBuilder* builder) {
+  auto args_fb = Encode(op.args_, builder);
+  auto work_group_size_fb = Encode(op.work_group_size_, builder);
+
+  auto def_fb = Encode(op.definition_, builder);
+  auto work_group_launch_order_fb =
+      Encode(op.work_group_launch_order_, builder);
+  auto grid_size_fb = Encode(op.grid_size_, builder);
+  auto work_groups_count_fb = Encode(op.work_groups_count_, builder);
+
+  std::vector<flatbuffers::Offset<flatbuffers::String>> src_names_fb;
+  src_names_fb.reserve(op.src_tensors_names_.size());
+  for (auto& name : op.src_tensors_names_) {
+    src_names_fb.push_back(builder->CreateString(name));
+  }
+  auto src_names_fb_vec = builder->CreateVector(src_names_fb);
+
+  std::vector<flatbuffers::Offset<flatbuffers::String>> dst_names_fb;
+  dst_names_fb.reserve(op.dst_tensors_names_.size());
+  for (auto& name : op.dst_tensors_names_) {
+    dst_names_fb.push_back(builder->CreateString(name));
+  }
+  auto dst_names_fb_vec = builder->CreateVector(dst_names_fb);
+
+  data::GPUOperationBuilder op_builder(*builder);
+  op_builder.add_arguments(args_fb);
+  op_builder.add_work_group_size(work_group_size_fb);
+  op_builder.add_tensor_to_grid(ToFB(op.tensor_to_grid_));
+  op_builder.add_flops(op.flops_);
+  op_builder.add_definition(def_fb);
+  op_builder.add_grid_dimension(op.grid_dimension_);
+  op_builder.add_work_group_launch_order(work_group_launch_order_fb);
+  op_builder.add_grid_size(grid_size_fb);
+  op_builder.add_src_tensors_names(src_names_fb_vec);
+  op_builder.add_dst_tensors_names(dst_names_fb_vec);
+  op_builder.add_work_groups_count(work_groups_count_fb);
+  return op_builder.Finish();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/task/serialization_base.h

@@ -0,0 +1,44 @@
+/* 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_TASK_SERIALIZATION_BASE_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_SERIALIZATION_BASE_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/tflite_serialization_base_generated.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+flatbuffers::Offset<data::Int2> Encode(const int2& v,
+                                       flatbuffers::FlatBufferBuilder* builder);
+
+flatbuffers::Offset<data::Int3> Encode(const int3& v,
+                                       flatbuffers::FlatBufferBuilder* builder);
+
+flatbuffers::Offset<data::TensorDescriptor> Encode(
+    const TensorDescriptor& desc, flatbuffers::FlatBufferBuilder* builder);
+void Decode(const data::TensorDescriptor* fb_desc, TensorDescriptor* desc);
+
+flatbuffers::Offset<data::GPUOperation> Encode(
+    const GPUOperation& op, flatbuffers::FlatBufferBuilder* builder);
+absl::Status Decode(const data::GPUOperation* fb_op, GPUOperation* op);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_SERIALIZATION_BASE_H_

tensorflow/lite/delegates/gpu/common/task/tensor_desc.h

@@ -0,0 +1,472 @@
+/* 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_TASK_TENSOR_DESC_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_TENSOR_DESC_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_object_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 {
+
+enum class TensorStorageType {
+  UNKNOWN,
+  BUFFER,
+  IMAGE_BUFFER,
+  TEXTURE_2D,
+  TEXTURE_3D,
+  TEXTURE_ARRAY,
+  SINGLE_TEXTURE_2D
+};
+
+class TensorDescriptor : public GPUObjectDescriptor {
+ public:
+  TensorDescriptor() = default;
+  TensorDescriptor(DataType data_type, TensorStorageType storage_type,
+                   Layout layout)
+      : data_type_(data_type), storage_type_(storage_type), layout_(layout) {}
+
+  TensorDescriptor(const TensorDescriptor&) = default;
+  TensorDescriptor& operator=(const TensorDescriptor&) = default;
+  TensorDescriptor(TensorDescriptor&& desc);
+  TensorDescriptor& operator=(TensorDescriptor&& desc);
+
+  void CopyWithoutData(TensorDescriptor* desc) const;
+
+  bool operator==(const TensorDescriptor& d) const {
+    return data_type_ == d.data_type_ && storage_type_ == d.storage_type_ &&
+           layout_ == d.layout_;
+  }
+
+  bool operator!=(const TensorDescriptor& d) const { return !(*this == d); }
+
+  void GetGpuResources(const BHWDC& tensor_shape,
+                       GenericGPUResourcesWithValue* resources) const;
+
+  absl::Status PerformConstExpr(const GpuInfo& gpu_info,
+                                absl::string_view const_expr,
+                                std::string* result) const override;
+
+  absl::Status PerformSelector(const GpuInfo& gpu_info,
+                               absl::string_view selector,
+                               const std::vector<std::string>& args,
+                               const std::vector<std::string>& template_args,
+                               std::string* result) const override;
+
+  GPUResources GetGPUResources(const GpuInfo& gpu_info) const override;
+
+  void Release() override { data_.clear(); }
+  uint64_t GetSizeInBytes() const override { return data_.size(); };
+  size_t GetSizeInBytesForShape(const BHWDC& shape5d) const;
+
+  bool HasAxis(Axis axis) const;
+
+  absl::Status GetLinkingContextFromWriteSelector(
+      const std::vector<std::string>& args, std::string* value_name,
+      std::string* x_coord, std::string* y_coord, std::string* z_coord,
+      std::string* s_coord, std::string* b_coord) const;
+
+  template <DataType T>
+  void UploadData(const tflite::gpu::Tensor<BHWC, T>& src);
+  template <DataType T>
+  void DownloadData(tflite::gpu::Tensor<BHWC, T>* dst);
+  template <DataType T>
+  void UploadData(const tflite::gpu::Tensor<BHWDC, T>& src);
+  template <DataType T>
+  void DownloadData(tflite::gpu::Tensor<BHWDC, T>* dst);
+
+  template <DataType T>
+  void UploadData(const tflite::gpu::Tensor<HWC, T>& src);
+
+  template <DataType T>
+  void UploadData(const tflite::gpu::Tensor<Linear, T>& src);
+
+  int GetLinearIndex(const BHWDC& shape5d, int b, int x, int y, int d, int s,
+                     int sub_c) const;
+
+  bool SupportsZeroClamp(const Axis& axis, const GpuInfo& gpu_info) const;
+  bool CanReadOutOfBorder(const Axis& axis) const;
+  bool IsLinear() const;
+
+  DataType GetDataType() const { return data_type_; }
+  TensorStorageType GetStorageType() const { return storage_type_; }
+
+  // applicable only for types that: IsLinear -> true.
+  // In this case for address we have 1d component - addr (int)
+  // If for addr == -1 this linear storage type returns zero value, this
+  // function returns true, otherwise false
+  bool ReturnsZeroForNegOneRead(const GpuInfo& gpu_info) const;
+
+  absl::Status CanCreateTensorWithShape(const GpuInfo& gpu_info,
+                                        const BHWDC& shape) const;
+
+  absl::Status CanCreateTensorWithShape(const GpuInfo& gpu_info,
+                                        const BHWC& shape) const;
+
+  // Can udate storage type if in the current storage type this tensor can not
+  // be allocated with shape on specified device(gpu_info)
+  // Usual scenario is to create new tensor_desc on base of another and may be
+  // update storage type for new tensor_desc shape because it can be unsuported
+  // with old storage type
+  absl::Status UpdateToSupportedStorageType(const GpuInfo& gpu_info,
+                                            const BHWC& shape);
+
+  // shape must be initialized when using this function
+  std::vector<uint64_t> GetStorageDims() const;
+  // shape must be initialized when using this function
+  int3 GetFullTensorRegion() const;
+  // shape must be initialized when using this function
+  uint64_t GetMemorySizeInBytes() const;
+  // shape must be initialized when using this function
+  int GetElementSize() const;
+
+  void SetUseBufferForWriteOnlyTexture2d(bool value) {
+    use_buffer_for_write_only_2d_texture_ = value;
+  }
+  bool GetUseBufferForWriteOnlyTexture2d() const {
+    return use_buffer_for_write_only_2d_texture_;
+  }
+
+  void SetUseBufferForWriteOnlyImageBuffer(bool value) {
+    use_buffer_for_write_only_image_buffer_ = value;
+  }
+  bool GetUseBufferForWriteOnlyImageBuffer() const {
+    return use_buffer_for_write_only_image_buffer_;
+  }
+
+  void SetBHWCShape(const BHWC& new_shape) {
+    shape_ = BHWDC(new_shape.b, new_shape.h, new_shape.w, 1, new_shape.c);
+  }
+  void SetBHWDCShape(const BHWDC& new_shape) { shape_ = new_shape; }
+  BHWC GetBHWCShape() const {
+    return BHWC(shape_.b, shape_.h, shape_.w, shape_.c);
+  }
+  BHWDC GetBHWDCShape() const { return shape_; }
+  void SetData(std::vector<uint8_t>&& new_data) { data_ = new_data; }
+  const std::vector<uint8_t>& GetData() const { return data_; }
+
+ private:
+  friend flatbuffers::Offset<data::TensorDescriptor> Encode(
+      const TensorDescriptor& desc, flatbuffers::FlatBufferBuilder* builder);
+  friend void Decode(const data::TensorDescriptor* fb_desc,
+                     TensorDescriptor* desc);
+
+  template <DataType DataTypeT>
+  friend TensorDescriptor CreateConstantLinearTensorDescriptor(
+      DataType data_type, TensorStorageType storage_type,
+      const tflite::gpu::Tensor<Linear, DataTypeT>& src);
+
+  friend TensorDescriptor CreateConstantHWVec4TensorDescriptor(
+      DataType data_type, TensorStorageType storage_type, int width, int height,
+      const uint8_t* data);
+
+  absl::Status PerformReadSelector(
+      const GpuInfo& gpu_info, const std::vector<std::string>& args,
+      const std::vector<std::string>& template_args, std::string* result) const;
+  absl::Status PerformReadNearestSelector(const GpuInfo& gpu_info,
+                                          const std::vector<std::string>& args,
+                                          std::string* result) const;
+  absl::Status PerformReadBilinearSelector(const GpuInfo& gpu_info,
+                                           const std::vector<std::string>& args,
+                                           std::string* result) const;
+  absl::Status PerformReadPerChannelSelector(
+      const GpuInfo& gpu_info, const std::vector<std::string>& args,
+      const std::vector<std::string>& template_args, std::string* result) const;
+
+  absl::Status PerformGetAddressSelector(const std::vector<std::string>& args,
+                                         std::string* result) const;
+
+  absl::Status PerformGetHandleSelector(const std::vector<std::string>& args,
+                                        std::string* result) const;
+
+  std::string StorageTypeToAddressType() const;
+
+  absl::Status PerformWriteSelector(const GpuInfo& gpu_info,
+                                    const std::vector<std::string>& args,
+                                    std::string* result) const;
+
+  absl::Status PerformWriteLinearSelector(const GpuInfo& gpu_info,
+                                          const std::vector<std::string>& args,
+                                          std::string* result) const;
+
+  absl::Status PerformWrite2DSelector(const GpuInfo& gpu_info,
+                                      const std::vector<std::string>& args,
+                                      std::string* result) const;
+
+  std::string Read(const GpuInfo& gpu_info, DataType read_as_type,
+                   const std::vector<std::string>& coords) const;
+  std::string Write(const GpuInfo& gpu_info, absl::string_view var_name,
+                    const std::vector<std::string>& coords) const;
+
+  absl::Status MaybeGetDataTypeFromTemplateArgs(
+      const std::vector<std::string>& template_args, DataType* result) const;
+
+  std::string GetGlobalAddressNoDeclaration(absl::string_view xc,
+                                            absl::string_view yc,
+                                            absl::string_view zc,
+                                            absl::string_view sc,
+                                            absl::string_view bc) const;
+
+  std::vector<std::string> GetPhysicalCoordsWHS(absl::string_view x,
+                                                absl::string_view y,
+                                                absl::string_view s) const;
+  std::vector<std::string> GetPhysicalCoordsWHSB(absl::string_view x,
+                                                 absl::string_view y,
+                                                 absl::string_view s,
+                                                 absl::string_view b) const;
+  std::vector<std::string> GetPhysicalCoordsWHDS(absl::string_view x,
+                                                 absl::string_view y,
+                                                 absl::string_view z,
+                                                 absl::string_view s) const;
+  std::vector<std::string> GetPhysicalCoordsWHDSB(absl::string_view x,
+                                                  absl::string_view y,
+                                                  absl::string_view z,
+                                                  absl::string_view s,
+                                                  absl::string_view b) const;
+  std::vector<std::string> GetPhysicalCoords(absl::string_view xc,
+                                             absl::string_view yc,
+                                             absl::string_view zc,
+                                             absl::string_view sc,
+                                             absl::string_view bc) const;
+  std::vector<std::string> GetPhysicalCoordsLinear(absl::string_view x) const;
+  std::vector<std::string> GetPhysicalCoordsHW(absl::string_view x,
+                                               absl::string_view y) const;
+
+  bool ParseCoordsFromArgs(const std::vector<std::string>& args, int offset,
+                           std::string* xc, std::string* yc, std::string* zc,
+                           std::string* sc, std::string* bc) const;
+
+  template <typename T>
+  void UploadData(const T* src);
+  template <typename T>
+  void DownloadData(T* dst);
+
+  DataType data_type_ = DataType::UNKNOWN;
+  TensorStorageType storage_type_ = TensorStorageType::UNKNOWN;
+
+  // This field describes logical layout, actual(physical) GPU layout can be
+  // totally different.
+  Layout layout_ =
+      Layout::UNKNOWN;  // Supported layouts is HWC, BHWC, HWDC, BHWDC
+                        // HW and LINEAR (for constant objects only)
+
+  // applicable only for TEXTURE_2D.
+  // When Texture 2d created from buffer, we can use it as texture or as buffer.
+  // This option allows to use texture 2d as buffer when we use it as dst
+  // tensor(write only).
+  // Currently supported only for Metal/OpenCL.
+  // By default false.
+  bool use_buffer_for_write_only_2d_texture_ = false;
+
+  // applicable only for IMAGE_BUFFER.
+  // We can use image buffer as image or as buffer.
+  // This option allows to use image buffer as buffer when we use it as dst
+  // tensor(write only).
+  // Currently supported only for Metal/OpenCL.
+  // By default true.
+  bool use_buffer_for_write_only_image_buffer_ = true;
+
+  // optional
+  BHWDC shape_;
+  std::vector<uint8_t> data_;
+};
+
+TensorDescriptor CreateBhwcTensorDescriptor(DataType data_type,
+                                            TensorStorageType storage_type,
+                                            const BHWC& shape);
+TensorDescriptor CreateHwcTensorDescriptor(DataType data_type,
+                                           TensorStorageType storage_type,
+                                           const HWC& shape);
+
+TensorStorageType GetStorageTypeForLinearTensor(const GpuInfo& gpu_info,
+                                                DataType data_type,
+                                                const Linear& shape);
+template <DataType DataTypeT>
+TensorDescriptor CreateConstantLinearTensorDescriptor(
+    DataType data_type, TensorStorageType storage_type,
+    const tflite::gpu::Tensor<Linear, DataTypeT>& src);
+
+template <DataType DataTypeT>
+TensorDescriptor CreateConstantLinearTensorDescriptor(
+    const GpuInfo& gpu_info, DataType data_type,
+    const tflite::gpu::Tensor<Linear, DataTypeT>& src);
+
+TensorDescriptor CreateConstantHWVec4TensorDescriptor(
+    DataType data_type, TensorStorageType storage_type, int width, int height,
+    const uint8_t* data);
+
+template <DataType T>
+void TensorDescriptor::UploadData(const tflite::gpu::Tensor<Linear, T>& src) {
+  shape_ = BHWDC(src.shape.v, 1, 1, 1, 1);
+  UploadData(src.data.data());
+}
+
+template <DataType T>
+void TensorDescriptor::UploadData(const tflite::gpu::Tensor<BHWC, T>& src) {
+  shape_ = BHWDC(src.shape.b, src.shape.h, src.shape.w, 1, src.shape.c);
+  UploadData(src.data.data());
+}
+
+template <DataType T>
+void TensorDescriptor::DownloadData(tflite::gpu::Tensor<BHWC, T>* dst) {
+  dst->shape = BHWC(shape_.b, shape_.h, shape_.w, shape_.c);
+  dst->data.resize(dst->shape.DimensionsProduct());
+  DownloadData(dst->data.data());
+}
+
+template <DataType T>
+void TensorDescriptor::UploadData(const tflite::gpu::Tensor<BHWDC, T>& src) {
+  shape_ = src.shape;
+  UploadData(src.data.data());
+}
+
+template <DataType T>
+void TensorDescriptor::DownloadData(tflite::gpu::Tensor<BHWDC, T>* dst) {
+  dst->shape = shape_;
+  dst->data.resize(dst->shape.DimensionsProduct());
+  DownloadData(dst->data.data());
+}
+
+template <typename T>
+void TensorDescriptor::UploadData(const T* src) {
+  data_.resize(GetSizeInBytesForShape(shape_));
+  if (layout_ == Layout::LINEAR) {
+    if (data_type_ == DataType::FLOAT16) {
+      half* gpu_data = reinterpret_cast<half*>(data_.data());
+      DataFromLinear(src, *this, gpu_data);
+    } else {
+      T* gpu_data = reinterpret_cast<T*>(data_.data());
+      DataFromLinear(src, *this, gpu_data);
+    }
+  } else {  // HWC/BHWC/HWDC/BHWDC
+    if (data_type_ == DataType::FLOAT16) {
+      half* gpu_data = reinterpret_cast<half*>(data_.data());
+      DataFromBHWDC(src, shape_, *this, gpu_data);
+    } else {
+      T* gpu_data = reinterpret_cast<T*>(data_.data());
+      DataFromBHWDC(src, shape_, *this, gpu_data);
+    }
+  }
+}
+
+template <typename T>
+void TensorDescriptor::DownloadData(T* dst) {
+  data_.resize(GetSizeInBytesForShape(shape_));
+  if (data_type_ == DataType::FLOAT16) {
+    half* gpu_data = reinterpret_cast<half*>(data_.data());
+    DataToBHWDC(gpu_data, shape_, *this, dst);
+  } else {
+    T* gpu_data = reinterpret_cast<T*>(data_.data());
+    DataToBHWDC(gpu_data, shape_, *this, dst);
+  }
+}
+
+template <typename FromType, typename ToType>
+void DataFromLinear(const FromType* src, const TensorDescriptor& desc,
+                    ToType* dst) {
+  const int element_size = desc.GetElementSize();
+  const Linear shape = Linear(desc.GetBHWCShape().c);
+  const int slices = DivideRoundUp(shape.v, element_size);
+  for (int s = 0; s < slices; ++s) {
+    for (int c = 0; c < element_size; ++c) {
+      FromType value;
+      if (s * 4 + c < shape.v) {
+        const int cpu_index = shape.LinearIndex({s * element_size + c});
+        value = src[cpu_index];
+      } else {
+        value = 0;
+      }
+      int gpu_index = s * element_size + c;
+      dst[gpu_index] = value;
+    }
+  }
+}
+
+template <typename FromType, typename ToType>
+void DataFromBHWDC(const FromType* src, const BHWDC& shape,
+                   const TensorDescriptor& desc, ToType* dst) {
+  const int channels_alignment =
+      desc.GetStorageType() == TensorStorageType::SINGLE_TEXTURE_2D ? shape.c
+                                                                    : 4;
+  const int slices = DivideRoundUp(shape.c, 4);
+  for (int b = 0; b < shape.b; ++b) {
+    for (int s = 0; s < slices; ++s) {
+      for (int y = 0; y < shape.h; ++y) {
+        for (int x = 0; x < shape.w; ++x) {
+          for (int d = 0; d < shape.d; ++d) {
+            for (int c = 0; c < channels_alignment; ++c) {
+              FromType value;
+              if (s * 4 + c < shape.c) {
+                const int cpu_index =
+                    shape.LinearIndex({b, y, x, d, s * 4 + c});
+                value = src[cpu_index];
+              } else {
+                value = 0;
+              }
+              int gpu_index = desc.GetLinearIndex(shape, b, x, y, d, s, c);
+              dst[gpu_index] = value;
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+template <typename FromType, typename ToType>
+void DataToBHWDC(const FromType* src, const BHWDC& shape,
+                 const TensorDescriptor& desc, ToType* dst) {
+  const int channels_alignment =
+      desc.GetStorageType() == TensorStorageType::SINGLE_TEXTURE_2D ? shape.c
+                                                                    : 4;
+  const int slices = DivideRoundUp(shape.c, 4);
+  for (int b = 0; b < shape.b; ++b) {
+    for (int s = 0; s < slices; ++s) {
+      for (int y = 0; y < shape.h; ++y) {
+        for (int x = 0; x < shape.w; ++x) {
+          for (int d = 0; d < shape.d; ++d) {
+            for (int c = 0; c < channels_alignment; ++c) {
+              if (s * 4 + c >= shape.c) {
+                continue;
+              }
+              int cpu_index = shape.LinearIndex({b, y, x, d, s * 4 + c});
+              int gpu_index = desc.GetLinearIndex(shape, b, x, y, d, s, c);
+              dst[cpu_index] = src[gpu_index];
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+std::string ToString(TensorStorageType type);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_TENSOR_DESC_H_

tensorflow/lite/delegates/gpu/common/task/testing_util.cc

@@ -0,0 +1,187 @@
+/* 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/task/testing_util.h"
+
+#include <cmath>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
+#include "absl/strings/str_cat.h"
+
+namespace tflite {
+namespace gpu {
+
+absl::Status TestExecutionEnvironment::ExecuteGPUOperation(
+    const std::vector<TensorDescriptor*>& src_cpu,
+    const std::vector<TensorDescriptor*>& dst_cpu,
+    std::unique_ptr<GPUOperation>&& operation) {
+  const OperationDef& op_def = operation->GetDefinition();
+  for (int i = 0; i < src_cpu.size(); ++i) {
+    auto src_shape = src_cpu[i]->GetBHWDCShape();
+    if (src_shape.b != 1 && !op_def.IsBatchSupported()) {
+      return absl::InvalidArgumentError(
+          "Layout doesn't have Batch dimension, but shape.b != 1");
+    }
+  }
+
+  for (int i = 0; i < dst_cpu.size(); ++i) {
+    auto dst_shape = dst_cpu[i]->GetBHWDCShape();
+    if (dst_shape.b != 1 && !op_def.IsBatchSupported()) {
+      return absl::InvalidArgumentError(
+          "Layout doesn't have Batch dimension, but shape.b != 1");
+    }
+  }
+  RETURN_IF_ERROR(operation->AssembleCode(GetGpuInfo()));
+  return ExecuteGpuOperationInternal(src_cpu, dst_cpu, std::move(operation));
+}
+
+template <typename DstTensorType>
+absl::Status TestExecutionEnvironment::ExecuteGPUOperation(
+    const std::vector<TensorFloat32>& src_cpu,
+    std::unique_ptr<GPUOperation>&& operation,
+    const std::vector<BHWC>& dst_sizes,
+    const std::vector<DstTensorType*>& dst_cpu) {
+  const OperationDef& op_def = operation->GetDefinition();
+  std::vector<TensorDescriptor> src_cpu_descs(src_cpu.size());
+  std::vector<TensorDescriptor*> src_cpu_desc_ptrs(src_cpu.size());
+  for (int i = 0; i < src_cpu.size(); ++i) {
+    src_cpu_descs[i] = op_def.src_tensors[i];
+    src_cpu_descs[i].UploadData(src_cpu[i]);
+    src_cpu_desc_ptrs[i] = &src_cpu_descs[i];
+  }
+  std::vector<TensorDescriptor> dst_cpu_descs(dst_cpu.size());
+  std::vector<TensorDescriptor*> dst_cpu_desc_ptrs(dst_cpu.size());
+  for (int i = 0; i < dst_cpu.size(); ++i) {
+    dst_cpu_descs[i] = op_def.dst_tensors[i];
+    dst_cpu_descs[i].SetBHWCShape(dst_sizes[i]);
+    dst_cpu_desc_ptrs[i] = &dst_cpu_descs[i];
+  }
+
+  RETURN_IF_ERROR(ExecuteGPUOperation(src_cpu_desc_ptrs, dst_cpu_desc_ptrs,
+                                      std::move(operation)));
+
+  for (int i = 0; i < dst_cpu.size(); ++i) {
+    dst_cpu_descs[i].DownloadData(dst_cpu[i]);
+  }
+  return absl::OkStatus();
+}
+
+template absl::Status TestExecutionEnvironment::ExecuteGPUOperation(
+    const std::vector<TensorFloat32>& src_cpu,
+    std::unique_ptr<GPUOperation>&& operation,
+    const std::vector<BHWC>& dst_sizes,
+    const std::vector<TensorFloat32*>& dst_cpu);
+
+template absl::Status TestExecutionEnvironment::ExecuteGPUOperation(
+    const std::vector<TensorFloat32>& src_cpu,
+    std::unique_ptr<GPUOperation>&& operation,
+    const std::vector<BHWC>& dst_sizes,
+    const std::vector<TensorInt32*>& dst_cpu);
+
+absl::Status TestExecutionEnvironment::ExecuteGPUOperation(
+    const std::vector<Tensor5DFloat32>& src_cpu,
+    std::unique_ptr<GPUOperation>&& operation,
+    const std::vector<BHWDC>& dst_sizes,
+    const std::vector<Tensor5DFloat32*>& dst_cpu) {
+  const OperationDef& op_def = operation->GetDefinition();
+  std::vector<TensorDescriptor> src_cpu_descs(src_cpu.size());
+  std::vector<TensorDescriptor*> src_cpu_desc_ptrs(src_cpu.size());
+  for (int i = 0; i < src_cpu.size(); ++i) {
+    src_cpu_descs[i] = op_def.src_tensors[i];
+    src_cpu_descs[i].UploadData(src_cpu[i]);
+    src_cpu_desc_ptrs[i] = &src_cpu_descs[i];
+  }
+  std::vector<TensorDescriptor> dst_cpu_descs(dst_cpu.size());
+  std::vector<TensorDescriptor*> dst_cpu_desc_ptrs(dst_cpu.size());
+  for (int i = 0; i < dst_cpu.size(); ++i) {
+    dst_cpu_descs[i] = op_def.dst_tensors[i];
+    dst_cpu_descs[i].SetBHWDCShape(dst_sizes[i]);
+    dst_cpu_desc_ptrs[i] = &dst_cpu_descs[i];
+  }
+
+  RETURN_IF_ERROR(ExecuteGPUOperation(src_cpu_desc_ptrs, dst_cpu_desc_ptrs,
+                                      std::move(operation)));
+
+  for (int i = 0; i < dst_cpu.size(); ++i) {
+    dst_cpu_descs[i].DownloadData(dst_cpu[i]);
+  }
+  return absl::OkStatus();
+}
+
+absl::Status PointWiseNear(const std::vector<float>& ref,
+                           const std::vector<float>& to_compare, float eps) {
+  if (ref.size() != to_compare.size()) {
+    return absl::InternalError(absl::StrCat("ref size(", ref.size(),
+                                            ") != to_compare size(",
+                                            to_compare.size(), ")"));
+  }
+  for (int i = 0; i < ref.size(); ++i) {
+    const float abs_diff = fabs(ref[i] - to_compare[i]);
+    if (abs_diff > eps) {
+      return absl::InternalError(absl::StrCat(
+          "ref[", i, "] = ", ref[i], ", to_compare[", i, "] = ", to_compare[i],
+          ", abs diff = ", abs_diff, " > ", eps, " (eps)"));
+    }
+  }
+  return absl::OkStatus();
+}
+
+template <typename DstTensorType>
+absl::Status TestExecutionEnvironment::ExecuteGpuModel(
+    const std::vector<TensorFloat32>& src_cpu,
+    const std::vector<DstTensorType*>& dst_cpu, GpuModel* gpu_model) {
+  for (int i = 0; i < gpu_model->input_ids_and_refs.size(); ++i) {
+    gpu_model->tensors[gpu_model->input_ids_and_refs[i].first].UploadData(
+        src_cpu[i]);
+  }
+
+  for (int k = 0; k < gpu_model->nodes.size(); ++k) {
+    auto& gpu_node = gpu_model->nodes[k];
+    std::vector<TensorDescriptor*> src_descs(gpu_node.inputs.size());
+    for (int i = 0; i < gpu_node.inputs.size(); ++i) {
+      src_descs[i] = &gpu_model->tensors[gpu_node.inputs[i]];
+    }
+
+    std::vector<TensorDescriptor*> dst_descs(gpu_node.outputs.size());
+    for (int i = 0; i < gpu_node.outputs.size(); ++i) {
+      dst_descs[i] = &gpu_model->tensors[gpu_node.outputs[i]];
+    }
+
+    RETURN_IF_ERROR(ExecuteGpuOperationInternal(
+        src_descs, dst_descs, std::move(gpu_model->nodes[k].gpu_operation)));
+  }
+  for (int i = 0; i < gpu_model->output_ids_and_refs.size(); ++i) {
+    gpu_model->tensors[gpu_model->output_ids_and_refs[i].first].DownloadData(
+        dst_cpu[i]);
+  }
+  return absl::OkStatus();
+}
+
+template absl::Status TestExecutionEnvironment::ExecuteGpuModel(
+    const std::vector<TensorFloat32>& src_cpu,
+    const std::vector<Tensor<BHWC, DataType::FLOAT32>*>& dst_cpu,
+    GpuModel* gpu_model);
+
+template absl::Status TestExecutionEnvironment::ExecuteGpuModel(
+    const std::vector<TensorFloat32>& src_cpu,
+    const std::vector<Tensor<BHWC, DataType::INT32>*>& dst_cpu,
+    GpuModel* gpu_model);
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/task/testing_util.h

@@ -0,0 +1,126 @@
+/* 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_TASK_TESTING_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_TESTING_UTIL_H_
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.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/gpu_model.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"
+
+namespace tflite {
+namespace gpu {
+using TensorInt32 = Tensor<BHWC, DataType::INT32>;
+class TestExecutionEnvironment {
+ public:
+  TestExecutionEnvironment() = default;
+  virtual ~TestExecutionEnvironment() = default;
+
+  virtual std::vector<CalculationsPrecision> GetSupportedPrecisions() const = 0;
+  virtual std::vector<TensorStorageType> GetSupportedStorages(
+      DataType data_type) const = 0;
+
+  virtual const GpuInfo& GetGpuInfo() const = 0;
+
+  absl::Status ExecuteGPUOperation(
+      const std::vector<TensorDescriptor*>& src_cpu,
+      const std::vector<TensorDescriptor*>& dst_cpu,
+      std::unique_ptr<GPUOperation>&& operation);
+
+  template <typename DstTensorType>
+  absl::Status ExecuteGpuModel(const std::vector<TensorFloat32>& src_cpu,
+                               const std::vector<DstTensorType*>& dst_cpu,
+                               GpuModel* gpu_model);
+
+  template <typename DstTensorType>
+  absl::Status ExecuteGPUOperation(const std::vector<TensorFloat32>& src_cpu,
+                                   std::unique_ptr<GPUOperation>&& operation,
+                                   const std::vector<BHWC>& dst_sizes,
+                                   const std::vector<DstTensorType*>& dst_cpu);
+
+  absl::Status ExecuteGPUOperation(
+      const std::vector<TensorFloat32>& src_cpu,
+      std::unique_ptr<GPUOperation>&& operation,
+      const std::vector<BHWC>& dst_sizes,
+      const std::initializer_list<TensorFloat32*>& dst_cpu) {
+    return ExecuteGPUOperation(src_cpu, std::move(operation), dst_sizes,
+                               std::vector<TensorFloat32*>(dst_cpu));
+  }
+
+  absl::Status ExecuteGPUOperation(
+      const std::vector<Tensor5DFloat32>& src_cpu,
+      std::unique_ptr<GPUOperation>&& operation,
+      const std::vector<BHWDC>& dst_sizes,
+      const std::vector<Tensor5DFloat32*>& dst_cpu);
+
+  absl::Status ExecuteGPUOperation(const TensorFloat32& src_cpu,
+                                   std::unique_ptr<GPUOperation>&& operation,
+                                   const BHWC& dst_size,
+                                   TensorFloat32* result) {
+    return ExecuteGPUOperation(std::vector<TensorFloat32>{src_cpu},
+                               std::move(operation), dst_size, result);
+  }
+
+  absl::Status ExecuteGPUOperation(const Tensor5DFloat32& src_cpu,
+                                   std::unique_ptr<GPUOperation>&& operation,
+                                   const BHWDC& dst_size,
+                                   Tensor5DFloat32* result) {
+    return ExecuteGPUOperation(std::vector<Tensor5DFloat32>{src_cpu},
+                               std::move(operation), dst_size, result);
+  }
+
+  absl::Status ExecuteGPUOperation(const std::vector<TensorFloat32>& src_cpu,
+                                   std::unique_ptr<GPUOperation>&& operation,
+                                   const BHWC& dst_size,
+                                   TensorFloat32* result) {
+    return ExecuteGPUOperation(
+        std::vector<TensorFloat32>{src_cpu}, std::move(operation),
+        std::vector<BHWC>{dst_size}, std::vector<TensorFloat32*>{result});
+  }
+
+  absl::Status ExecuteGPUOperation(const std::vector<Tensor5DFloat32>& src_cpu,
+                                   std::unique_ptr<GPUOperation>&& operation,
+                                   const BHWDC& dst_size,
+                                   Tensor5DFloat32* result) {
+    return ExecuteGPUOperation(
+        std::vector<Tensor5DFloat32>{src_cpu}, std::move(operation),
+        std::vector<BHWDC>{dst_size}, std::vector<Tensor5DFloat32*>{result});
+  }
+
+ protected:
+  virtual absl::Status ExecuteGpuOperationInternal(
+      const std::vector<TensorDescriptor*>& src_cpu,
+      const std::vector<TensorDescriptor*>& dst_cpu,
+      std::unique_ptr<GPUOperation>&& operation) = 0;
+};
+
+absl::Status PointWiseNear(const std::vector<float>& ref,
+                           const std::vector<float>& to_compare,
+                           float eps = 0.0f);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_TESTING_UTIL_H_

tensorflow/lite/delegates/gpu/common/task/tflite_serialization_base.fbs

@@ -0,0 +1,210 @@
+// 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.
+
+namespace tflite.gpu.data;
+
+table Int4 {
+  x:int32;
+  y:int32;
+  z:int32;
+  w:int32;
+}
+
+table Int3 {
+  x:int32;
+  y:int32;
+  z:int32;
+}
+
+table Int2 {
+  x:int32;
+  y:int32;
+}
+
+enum AccessType : byte {
+  READ = 0,
+  WRITE = 1,
+  READ_WRITE = 2,
+}
+
+table StateVariable {
+  key:string;
+  value:string;
+}
+
+table GPUObjectDescriptor {
+  state_vars:[StateVariable];
+  access_type:AccessType;
+}
+
+table IntValue {
+  name:string;
+  value:int32;
+  active:bool;
+}
+
+table FloatValue {
+  name:string;
+  value:float;
+  active:bool;
+}
+
+table HalfValue {
+  name:string;
+  value:float;
+  active:bool;
+}
+
+enum DataType : byte {
+  UNKNOWN = 0,
+  FLOAT16 = 1,
+  FLOAT32 = 2,
+  FLOAT64 = 3,
+  UINT8 = 4,
+  INT8 = 5,
+  UINT16 = 6,
+  INT16 = 7,
+  UINT32 = 8,
+  INT32 = 9,
+  UINT64 = 10,
+  INT64 = 11,
+  BOOL = 12,
+}
+
+enum MemoryType : byte {
+  GLOBAL = 0,
+  CONSTANT = 1,
+  LOCAL = 2,
+}
+
+table BufferDescriptor {
+  base_obj:GPUObjectDescriptor;
+  element_type:DataType;
+  element_size:int32;
+  memory_type:MemoryType;
+  attributes:[string];
+  size:int32;
+  data:[uint8];
+}
+
+enum TensorStorageType : byte {
+  UNKNOWN = 0,
+  BUFFER = 1,
+  IMAGE_BUFFER = 2,
+  TEXTURE_2D = 3,
+  TEXTURE_3D = 4,
+  TEXTURE_ARRAY = 5,
+  SINGLE_TEXTURE_2D = 6,
+}
+
+enum Layout : byte {
+  UNKNOWN = 0,
+  HWC = 1,
+  BHWC = 2,
+  HWDC = 3,
+  BHWDC = 4,
+  LINEAR = 5,
+  HW = 6,
+}
+
+table BHWDC {
+  b:int32;
+  h:int32;
+  w:int32;
+  d:int32;
+  c:int32;
+}
+
+table TensorDescriptor {
+  base_obj:GPUObjectDescriptor;
+  data_type:DataType;
+  storage_type:TensorStorageType;
+  layout:Layout;
+  shape:BHWDC;
+  data:[uint8];
+  use_buffer_for_write_only_2d_texture:bool;
+  use_buffer_for_write_only_image_buffer:bool;
+}
+
+table BufferDescriptorMapValue {
+  key:string;
+  value:BufferDescriptor;
+}
+
+table TensorDescriptorMapValue {
+  key:string;
+  value:TensorDescriptor;
+}
+
+table Arguments {
+  int_values:[IntValue];
+  float_values:[FloatValue];
+  half_values:[HalfValue];
+
+  buffer_refs:[BufferDescriptorMapValue];
+  tensor_refs:[TensorDescriptorMapValue];
+
+  buffer_objects:[BufferDescriptorMapValue];
+  tensor_objects:[TensorDescriptorMapValue];
+}
+
+enum CalculationsPrecision : byte {
+  F32 = 0,
+  F32_F16 = 1,
+  F16 = 2,
+}
+
+enum TensorToGrid : byte {
+  CUSTOM = 0,
+  WB_TO_X_HD_TO_Y_S_TO_Z = 1,
+  WB_TO_X_HD_TO_Y_Z_IS_1 = 2,
+  WB_TO_X_H_TO_Y_D_TO_Z = 3,
+  B_TO_X_Y_IS_1_Z_IS_1 = 4,
+}
+
+enum CompilerOptions : byte {
+  ADRENO_FULL_SIMD_LINE = 0,
+  ADRENO_MORE_WAVES = 1,
+  CL_FAST_RELAXED_MATH = 2,
+  CL_OPT_DISABLE = 3,
+  CL_2_0 = 4,
+  CL_3_0 = 5,
+  CL_REGISTER_ALLOCATION_64 = 6,
+}
+
+table OperationDef {
+  precision:CalculationsPrecision;
+  src_tensors:[TensorDescriptor];
+  dst_tensors:[TensorDescriptor];
+}
+
+table CompilerOption {
+  option:CompilerOptions;
+}
+
+table GPUOperation {
+  arguments:Arguments;
+  code:string;
+  work_group_size:Int3;
+  compiler_options:[CompilerOption];
+  tensor_to_grid:TensorToGrid;
+  flops:uint64;
+  definition:OperationDef;
+  grid_dimension:int32;
+  work_group_launch_order:Int3;
+  grid_size:Int3;
+  src_tensors_names:[string];
+  dst_tensors_names:[string];
+  work_groups_count:Int3;
+}

tensorflow/lite/delegates/gpu/common/task/tuning_type.h

@@ -0,0 +1,27 @@
+/* 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_TASK_TUNING_TYPE_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_TUNING_TYPE_H_
+
+namespace tflite {
+namespace gpu {
+
+enum class TuningType { kExhaustive, kFast };
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_TUNING_TYPE_H_

tensorflow/lite/delegates/gpu/common/task/util.cc

@@ -0,0 +1,347 @@
+/* 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/task/util.h"
+
+#include <cfloat>
+#include <cstddef>
+#include <map>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/ascii.h"
+#include "absl/strings/str_split.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/util.h"
+
+namespace tflite {
+namespace gpu {
+namespace {
+std::string GetGlslConversion(const GpuInfo& gpu_info, DataType src_type,
+                              DataType dst_type, int vec_size) {
+  if (src_type == dst_type) {
+    return "";
+  }
+  bool need_explicit_conversion = true;
+  switch (dst_type) {
+    case DataType::FLOAT32:
+    case DataType::FLOAT16:
+      if (gpu_info.IsGlslSupportsExplicitFp16()) {
+        if (src_type == dst_type) {
+          need_explicit_conversion = false;
+        }
+      } else {
+        if (src_type == DataType::FLOAT32 || src_type == DataType::FLOAT16) {
+          need_explicit_conversion = false;
+        }
+      }
+      break;
+    case DataType::INT32:
+    case DataType::INT16:
+    case DataType::INT8:
+      if (src_type == DataType::INT32 || src_type == DataType::INT16 ||
+          src_type == DataType::INT8) {
+        need_explicit_conversion = false;
+      }
+      break;
+    case DataType::UINT32:
+    case DataType::UINT16:
+    case DataType::UINT8:
+      if (src_type == DataType::UINT32 || src_type == DataType::UINT16 ||
+          src_type == DataType::UINT8) {
+        need_explicit_conversion = false;
+      }
+      break;
+    case DataType::BOOL:
+      need_explicit_conversion = true;
+      break;
+    default:
+      break;
+  }
+  if (need_explicit_conversion) {
+    return ToGlslShaderDataType(
+        dst_type, vec_size,
+        /*add_precision*/ false,
+        /*explicit_fp16*/ gpu_info.IsGlslSupportsExplicitFp16());
+  } else {
+    return "";
+  }
+}
+
+bool IsWordSymbol(char symbol) {
+  return absl::ascii_isalnum(symbol) || symbol == '_';
+}
+}  // namespace
+
+std::string MemoryTypeToCLType(MemoryType type) {
+  switch (type) {
+    case MemoryType::GLOBAL:
+      return "__global";
+    case MemoryType::CONSTANT:
+      return "__constant";
+    case MemoryType::LOCAL:
+      return "__local";
+  }
+  return "";
+}
+
+std::string MemoryTypeToMetalType(MemoryType type) {
+  switch (type) {
+    case MemoryType::GLOBAL:
+      return "device";
+    case MemoryType::CONSTANT:
+      return "constant";
+      break;
+    case MemoryType::LOCAL:
+      return "threadgroup";
+  }
+  return "";
+}
+
+float4 GetMaskForLastPlane(int channels) {
+  float4 mask = float4(0.0f);
+  const int reminder = channels % 4 == 0 ? 4 : channels % 4;
+  for (int i = 0; i < reminder; ++i) {
+    mask[i] = 1.0f;
+  }
+  return mask;
+}
+
+int GetRecommendedBlockSizeForConv(const GpuInfo& gpu_info,
+                                   CalculationsPrecision precision,
+                                   int task_size) {
+  const float task_size_per_cu =
+      task_size / static_cast<float>(gpu_info.GetComputeUnitsCount());
+  int block_size = 1;
+  float threshold_1 = FLT_MAX;
+  float threshold_2 = FLT_MAX;
+  float threshold_4 = FLT_MAX;
+  if (!gpu_info.IsMali()) {
+    return 1;
+  }
+  MaliInfo mali_info = gpu_info.mali_info;
+  switch (precision) {
+    case CalculationsPrecision::F16:
+      if (mali_info.IsBifrostGen1()) {
+        threshold_1 = 256.0f;
+        threshold_2 = 256.0f * 4.0f;
+        threshold_4 = 256.0f * 8.0f;
+      } else if (mali_info.IsBifrostGen2()) {
+        threshold_1 = 256.0f * 2.0f;
+        threshold_2 = 256.0f * 8.0f;
+        threshold_4 = 256.0f * 16.0f;
+      } else if (mali_info.IsBifrostGen3() || mali_info.IsValhall()) {
+        threshold_1 = 256.0f;
+        threshold_2 = 256.0f * 6.0f;
+        threshold_4 = 256.0f * 16.0f;
+      } else if (mali_info.IsMidgard()) {
+        threshold_1 = 256.0f * 4.0f;
+        threshold_2 = 256.0f * 16.0f;
+      }
+      break;
+    case CalculationsPrecision::F32_F16:
+      if (mali_info.IsBifrostGen1()) {
+        threshold_1 = 256.0f;
+        threshold_2 = 256.0f * 3.0f;
+        threshold_4 = 256.0f * 32.0f;
+      } else if (mali_info.IsBifrostGen2()) {
+        threshold_1 = 256.0f * 2.0f;
+        threshold_2 = 256.0f * 8.0f;
+      } else if (mali_info.IsBifrostGen3() || mali_info.IsValhall()) {
+        threshold_1 = 256.0f;
+        threshold_2 = 256.0f * 8.0f;
+      } else if (mali_info.IsMidgard()) {
+        threshold_1 = 256.0f * 4.0f;
+      }
+      break;
+    case CalculationsPrecision::F32:
+      if (mali_info.IsBifrostGen1()) {
+        threshold_1 = 256.0f;
+        threshold_2 = 256.0f * 4.0f;
+      } else if (mali_info.IsBifrostGen2()) {
+        threshold_1 = 128.0f;
+        threshold_2 = 256.0f * 4.0f;
+      } else if (mali_info.IsBifrostGen3() || mali_info.IsValhall()) {
+        threshold_1 = 256.0f;
+        threshold_2 = 256.0f * 12.0f;
+      } else if (mali_info.IsMidgard()) {
+        threshold_1 = 256.0f * 16.0f;
+      }
+      break;
+  }
+  if (task_size_per_cu <= threshold_1) {
+    block_size = 1;
+  } else if (task_size_per_cu <= threshold_2) {
+    block_size = 2;
+  } else if (task_size_per_cu <= threshold_4) {
+    block_size = 4;
+  } else {
+    block_size = 8;
+  }
+  return block_size;
+}
+
+int3 GetWorkGroupsCount(const int3& grid_size, const int3& work_group_size) {
+  int3 work_groups_count;
+  work_groups_count.x = DivideRoundUp(grid_size.x, work_group_size.x);
+  work_groups_count.y = DivideRoundUp(grid_size.y, work_group_size.y);
+  work_groups_count.z = DivideRoundUp(grid_size.z, work_group_size.z);
+  return work_groups_count;
+}
+
+std::string GetTypeDeclaration(const GpuInfo& gpu_info, DataType data_type,
+                               int vec_size) {
+  if (gpu_info.IsApiOpenCl()) {
+    return ToCLDataType(data_type, vec_size);
+  } else if (gpu_info.IsApiMetal()) {
+    return ToMetalDataType(data_type, vec_size);
+  } else if (gpu_info.IsGlsl()) {
+    return ToGlslShaderDataType(data_type, vec_size, true,
+                                gpu_info.IsGlslSupportsExplicitFp16());
+  } else {
+    return "";
+  }
+}
+
+std::string GetZeroValue(const GpuInfo& gpu_info, DataType data_type,
+                         int vec_size) {
+  if (gpu_info.IsApiOpenCl()) {
+    return "(" + ToCLDataType(data_type, vec_size) + ")(0)";
+  } else if (gpu_info.IsApiMetal()) {
+    return ToMetalDataType(data_type, vec_size) + "(0)";
+  } else if (gpu_info.IsGlsl()) {
+    return ToGlslShaderDataType(data_type, vec_size, false,
+                                gpu_info.IsGlslSupportsExplicitFp16()) +
+           "(0)";
+  } else {
+    return "";
+  }
+}
+
+std::string GetOneValue(const GpuInfo& gpu_info, DataType data_type,
+                        int vec_size) {
+  if (gpu_info.IsApiOpenCl()) {
+    return "(" + ToCLDataType(data_type, vec_size) + ")(1)";
+  } else if (gpu_info.IsApiMetal()) {
+    return ToMetalDataType(data_type, vec_size) + "(1)";
+  } else if (gpu_info.IsGlsl()) {
+    return ToGlslShaderDataType(data_type, vec_size, false,
+                                gpu_info.IsGlslSupportsExplicitFp16()) +
+           "(1)";
+  } else {
+    return "";
+  }
+}
+
+std::string GetTypeConversion(const GpuInfo& gpu_info, DataType src_type,
+                              DataType dst_type, int vec_size) {
+  if (src_type != dst_type) {
+    if (gpu_info.IsApiOpenCl()) {
+      if (dst_type == DataType::BOOL && vec_size != 1) {
+        // In OpenCL for bool4 we are using uchar4
+        // From OpenCL specification for "Relational and Equality Operators":
+        //   "These functions shall return a 0 if the specified relation is
+        //   false and a -1 (i.e. all bits set) if the specified relation is
+        //   true for vector argument types."
+        // (convert_uchar4((value) != 0) & (uchar4)(1))
+        return "(convert_" + ToCLDataType(DataType::UINT8, vec_size) +
+               "(($0) != " + GetZeroValue(gpu_info, src_type, vec_size) +
+               ") & " + GetOneValue(gpu_info, DataType::UINT8, vec_size) + ")";
+      } else {
+        return "convert_" + ToCLDataType(dst_type, vec_size) + "($0)";
+      }
+    } else if (gpu_info.IsApiMetal()) {
+      return ToMetalDataType(dst_type, vec_size) + "($0)";
+    } else if (gpu_info.IsGlsl()) {
+      const std::string conversion =
+          GetGlslConversion(gpu_info, src_type, dst_type, vec_size);
+      if (!conversion.empty()) {
+        return conversion + "($0)";
+      } else {
+        return "$0";
+      }
+    }
+  }
+  return "$0";
+}
+
+std::string GetNextWord(const std::string& code, size_t first_position) {
+  size_t pos = first_position;
+  char t = code[pos];
+  while (IsWordSymbol(t)) {
+    pos++;
+    t = code[pos];
+  }
+  return code.substr(first_position, pos - first_position);
+}
+
+size_t FindEnclosingBracket(const std::string& text, size_t first_pos,
+                            char bracket) {
+  const std::map<char, char> brackets = {
+      {'(', ')'},
+      {'{', '}'},
+      {'[', ']'},
+      {'<', '>'},
+  };
+  char b_open = bracket;
+  auto it = brackets.find(b_open);
+  if (it == brackets.end()) {
+    return -1;
+  }
+  char b_close = it->second;
+  size_t pos = first_pos;
+  int opened = 1;
+  int closed = 0;
+  while (opened != closed && pos < text.size()) {
+    if (text[pos] == b_open) {
+      opened++;
+    } else if (text[pos] == b_close) {
+      closed++;
+    }
+    pos++;
+  }
+  if (opened == closed) {
+    return pos;
+  } else {
+    return -1;
+  }
+}
+
+absl::Status ParseArgsInsideBrackets(const std::string& text,
+                                     size_t open_bracket_pos,
+                                     size_t* close_bracket_pos,
+                                     std::vector<std::string>* args) {
+  *close_bracket_pos =
+      FindEnclosingBracket(text, open_bracket_pos + 1, text[open_bracket_pos]);
+  if (*close_bracket_pos == -1) {
+    return absl::NotFoundError("Not found enclosing bracket");
+  }
+  std::string str_args = text.substr(open_bracket_pos + 1,
+                                     *close_bracket_pos - open_bracket_pos - 2);
+  std::vector<absl::string_view> words = absl::StrSplit(str_args, ',');
+  args->reserve(words.size());
+  for (const auto& word : words) {
+    absl::string_view arg = absl::StripAsciiWhitespace(word);
+    if (!arg.empty()) {
+      args->push_back(std::string(arg));
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/task/util.h

@@ -0,0 +1,79 @@
+/* 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_TASK_UTIL_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_UTIL_H_
+
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/delegates/gpu/common/gpu_info.h"
+#include "tensorflow/lite/delegates/gpu/common/precision.h"
+#include "tensorflow/lite/delegates/gpu/common/task/gpu_object_desc.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+
+namespace tflite {
+namespace gpu {
+
+std::string MemoryTypeToCLType(MemoryType type);
+
+std::string MemoryTypeToMetalType(MemoryType type);
+
+// Returns float4 mask for last plane(batch of 4 channels)
+// assumes that plane size is 4;
+// for example we have 7 channels, in our data structures we align it to 8
+// but 8s-channel will be empty, then last plane (batch of 4 channels) will
+// have this mask (1, 1, 1, 0).
+float4 GetMaskForLastPlane(int channels);
+
+// task_size as amount of FLT4 processed elements.
+int GetRecommendedBlockSizeForConv(const GpuInfo& gpu_info,
+                                   CalculationsPrecision precision,
+                                   int task_size);
+
+int3 GetWorkGroupsCount(const int3& grid_size, const int3& work_group_size);
+
+std::string GetTypeDeclaration(const GpuInfo& gpu_info, DataType data_type,
+                               int vec_size);
+
+std::string GetZeroValue(const GpuInfo& gpu_info, DataType data_type,
+                         int vec_size);
+
+std::string GetOneValue(const GpuInfo& gpu_info, DataType data_type,
+                        int vec_size);
+
+// Returns expression that can be substituted for converted value
+// Intended to be used with absl::Substitute
+// Example usage:
+//   auto conversion_function = GetTypeConversion(gpu_info, UINT8, FLOAT32, 4);
+//   auto code = absl::Substitute(conversion_function, "value_name");
+std::string GetTypeConversion(const GpuInfo& gpu_info, DataType src_type,
+                              DataType dst_type, int vec_size);
+
+std::string GetNextWord(const std::string& code, size_t first_position);
+
+size_t FindEnclosingBracket(const std::string& text, size_t first_pos,
+                            char bracket);
+
+absl::Status ParseArgsInsideBrackets(const std::string& text,
+                                     size_t open_bracket_pos,
+                                     size_t* close_bracket_pos,
+                                     std::vector<std::string>* args);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_UTIL_H_

tensorflow/lite/delegates/gpu/common/task/weights_conversion.cc

@@ -0,0 +1,266 @@
+/* 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/task/weights_conversion.h"
+
+#include <cassert>
+#include <cstdint>
+#include <limits>
+
+#include "absl/types/span.h"
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.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/util.h"
+
+namespace tflite {
+namespace gpu {
+uint GetTotalElementsCountForLayout(const WeightsDescription& weight_desc,
+                                    const OHWDI& shape) {
+  if (weight_desc.layout == WeightsLayout::kOSpatialIOGroupI4O4 ||
+      weight_desc.layout == WeightsLayout::kOSpatialIOGroupO4I4 ||
+      weight_desc.layout == WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4 ||
+      weight_desc.layout == WeightsLayout::k2DX4O4YIsSpatialIAndXIsOOGroupI4) {
+    uint i_aligned = static_cast<uint>(AlignByN(shape.i, 4));
+    uint o_aligned =
+        static_cast<uint>(AlignByN(shape.o, 4 * weight_desc.output_group_size));
+    uint64_t total = static_cast<uint64_t>(i_aligned) *
+                     static_cast<uint64_t>(o_aligned) *
+                     static_cast<uint64_t>(static_cast<uint>(shape.h)) *
+                     static_cast<uint64_t>(static_cast<uint>(shape.w)) *
+                     static_cast<uint64_t>(static_cast<uint>(shape.d));
+    if (total > std::numeric_limits<uint>::max()) {
+      return 0;
+    }
+    return static_cast<uint>(total);
+  } else if (weight_desc.layout == WeightsLayout::kOICustomSpatialI4O4 ||
+             weight_desc.layout == WeightsLayout::kOICustomSpatialO4I4) {
+    uint i_aligned = static_cast<uint>(AlignByN(shape.i, 4));
+    uint o_aligned = static_cast<uint>(AlignByN(shape.o, 4));
+    uint64_t total = static_cast<uint64_t>(i_aligned) *
+                     static_cast<uint64_t>(o_aligned) *
+                     weight_desc.spatial_remap.size();
+    if (total > std::numeric_limits<uint>::max()) {
+      return 0;
+    }
+    return static_cast<uint>(total);
+  } else {
+    return static_cast<uint>(-1);
+  }
+}
+
+uint GetTotalElementsCountForLayout(const WeightsDescription& weight_desc,
+                                    const OHWI& shape) {
+  const OHWDI ohwdi_shape = OHWDI(shape.o, shape.h, shape.w, 1, shape.i);
+  return GetTotalElementsCountForLayout(weight_desc, ohwdi_shape);
+}
+
+uint2 Get2dResourceSize(const WeightsDescription& weight_desc,
+                        const OHWI& shape) {
+  const OHWDI ohwdi_shape = OHWDI(shape.o, shape.h, shape.w, 1, shape.i);
+  return Get2dResourceSize(weight_desc, ohwdi_shape);
+}
+
+uint2 Get2dResourceSize(const WeightsDescription& weight_desc,
+                        const OHWDI& shape) {
+  const int dst_depth =
+      AlignByN(DivideRoundUp(shape.o, 4), weight_desc.output_group_size);
+  const int src_depth = DivideRoundUp(shape.i, 4);
+
+  return uint2(dst_depth, src_depth * shape.h * shape.w * shape.d);
+}
+
+void RearrangeWeights(
+    const tflite::gpu::Tensor<OHWI, DataType::FLOAT32>& weights,
+    const WeightsDescription& dst_weight_desc, absl::Span<uint8_t> dst) {
+  const uint flt_count =
+      GetTotalElementsCountForLayout(dst_weight_desc, weights.shape);
+  if (flt_count == 0) return;
+  if (dst_weight_desc.layout == WeightsLayout::kOSpatialIOGroupI4O4) {
+    if (dst_weight_desc.type == DataType::FLOAT32) {
+      float4* f32_ptr = reinterpret_cast<float4*>(dst.data());
+      RearrangeWeightsToOHWIOGroupI4O4(weights,
+                                       dst_weight_desc.output_group_size,
+                                       absl::MakeSpan(f32_ptr, flt_count / 4));
+    } else if (dst_weight_desc.type == DataType::FLOAT16) {
+      half4* f16_ptr = reinterpret_cast<half4*>(dst.data());
+      RearrangeWeightsToOHWIOGroupI4O4(weights,
+                                       dst_weight_desc.output_group_size,
+                                       absl::MakeSpan(f16_ptr, flt_count / 4));
+    }
+    return;
+  } else if (dst_weight_desc.layout == WeightsLayout::kOSpatialIOGroupO4I4) {
+    if (dst_weight_desc.type == DataType::FLOAT32) {
+      float4* f32_ptr = reinterpret_cast<float4*>(dst.data());
+      RearrangeWeightsToOHWIOGroupO4I4(weights,
+                                       dst_weight_desc.output_group_size,
+                                       absl::MakeSpan(f32_ptr, flt_count / 4));
+    } else if (dst_weight_desc.type == DataType::FLOAT16) {
+      half4* f16_ptr = reinterpret_cast<half4*>(dst.data());
+      RearrangeWeightsToOHWIOGroupO4I4(weights,
+                                       dst_weight_desc.output_group_size,
+                                       absl::MakeSpan(f16_ptr, flt_count / 4));
+    }
+    return;
+  } else if (dst_weight_desc.layout == WeightsLayout::kOICustomSpatialI4O4) {
+    if (dst_weight_desc.type == DataType::FLOAT32) {
+      float4* f32_ptr = reinterpret_cast<float4*>(dst.data());
+      RearrangeWeightsToOICustomSpatialI4O4(
+          weights, dst_weight_desc.spatial_remap,
+          absl::MakeSpan(f32_ptr, flt_count / 4));
+    } else if (dst_weight_desc.type == DataType::FLOAT16) {
+      half4* f16_ptr = reinterpret_cast<half4*>(dst.data());
+      RearrangeWeightsToOICustomSpatialI4O4(
+          weights, dst_weight_desc.spatial_remap,
+          absl::MakeSpan(f16_ptr, flt_count / 4));
+    }
+    return;
+  } else if (dst_weight_desc.layout == WeightsLayout::kOICustomSpatialO4I4) {
+    if (dst_weight_desc.type == DataType::FLOAT32) {
+      float4* f32_ptr = reinterpret_cast<float4*>(dst.data());
+      RearrangeWeightsToOICustomSpatialO4I4(
+          weights, dst_weight_desc.spatial_remap,
+          absl::MakeSpan(f32_ptr, flt_count / 4));
+    } else if (dst_weight_desc.type == DataType::FLOAT16) {
+      half4* f16_ptr = reinterpret_cast<half4*>(dst.data());
+      RearrangeWeightsToOICustomSpatialO4I4(
+          weights, dst_weight_desc.spatial_remap,
+          absl::MakeSpan(f16_ptr, flt_count / 4));
+    }
+    return;
+  } else if (dst_weight_desc.layout ==
+             WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4) {
+    if (dst_weight_desc.type == DataType::FLOAT32) {
+      float4* f32_ptr = reinterpret_cast<float4*>(dst.data());
+      RearrangeWeightsToI4HWIOOGroupO4(weights,
+                                       dst_weight_desc.output_group_size,
+                                       absl::MakeSpan(f32_ptr, flt_count / 4));
+    } else if (dst_weight_desc.type == DataType::FLOAT16) {
+      half4* f16_ptr = reinterpret_cast<half4*>(dst.data());
+      RearrangeWeightsToI4HWIOOGroupO4(weights,
+                                       dst_weight_desc.output_group_size,
+                                       absl::MakeSpan(f16_ptr, flt_count / 4));
+    }
+    return;
+  } else if (dst_weight_desc.layout ==
+             WeightsLayout::k2DX4O4YIsSpatialIAndXIsOOGroupI4) {
+    if (dst_weight_desc.type == DataType::FLOAT32) {
+      float4* f32_ptr = reinterpret_cast<float4*>(dst.data());
+      RearrangeWeightsToO4HWIOOGroupI4(weights,
+                                       dst_weight_desc.output_group_size,
+                                       absl::MakeSpan(f32_ptr, flt_count / 4));
+    } else if (dst_weight_desc.type == DataType::FLOAT16) {
+      half4* f16_ptr = reinterpret_cast<half4*>(dst.data());
+      RearrangeWeightsToO4HWIOOGroupI4(weights,
+                                       dst_weight_desc.output_group_size,
+                                       absl::MakeSpan(f16_ptr, flt_count / 4));
+    }
+    return;
+  }
+}
+
+void RearrangeWeights(
+    const tflite::gpu::Tensor<OHWDI, DataType::FLOAT32>& weights,
+    const WeightsDescription& dst_weight_desc, absl::Span<uint8_t> dst) {
+  const uint flt_count =
+      GetTotalElementsCountForLayout(dst_weight_desc, weights.shape);
+  if (flt_count == 0) return;
+  if (dst_weight_desc.layout == WeightsLayout::kOSpatialIOGroupI4O4) {
+    if (dst_weight_desc.type == DataType::FLOAT32) {
+      float4* f32_ptr = reinterpret_cast<float4*>(dst.data());
+      RearrangeWeightsToODHWIOGroupI4O4(weights,
+                                        dst_weight_desc.output_group_size,
+                                        absl::MakeSpan(f32_ptr, flt_count / 4));
+    } else if (dst_weight_desc.type == DataType::FLOAT16) {
+      half4* f16_ptr = reinterpret_cast<half4*>(dst.data());
+      RearrangeWeightsToODHWIOGroupI4O4(weights,
+                                        dst_weight_desc.output_group_size,
+                                        absl::MakeSpan(f16_ptr, flt_count / 4));
+    }
+    return;
+  } else if (dst_weight_desc.layout == WeightsLayout::kOSpatialIOGroupO4I4) {
+    if (dst_weight_desc.type == DataType::FLOAT32) {
+      float4* f32_ptr = reinterpret_cast<float4*>(dst.data());
+      RearrangeWeightsToODHWIOGroupO4I4(weights,
+                                        dst_weight_desc.output_group_size,
+                                        absl::MakeSpan(f32_ptr, flt_count / 4));
+    } else if (dst_weight_desc.type == DataType::FLOAT16) {
+      half4* f16_ptr = reinterpret_cast<half4*>(dst.data());
+      RearrangeWeightsToODHWIOGroupO4I4(weights,
+                                        dst_weight_desc.output_group_size,
+                                        absl::MakeSpan(f16_ptr, flt_count / 4));
+    }
+    return;
+  } else if (dst_weight_desc.layout == WeightsLayout::kOICustomSpatialI4O4) {
+    if (dst_weight_desc.type == DataType::FLOAT32) {
+      float4* f32_ptr = reinterpret_cast<float4*>(dst.data());
+      RearrangeWeightsToOICustomSpatialI4O4(
+          weights, dst_weight_desc.spatial_remap,
+          absl::MakeSpan(f32_ptr, flt_count / 4));
+    } else if (dst_weight_desc.type == DataType::FLOAT16) {
+      half4* f16_ptr = reinterpret_cast<half4*>(dst.data());
+      RearrangeWeightsToOICustomSpatialI4O4(
+          weights, dst_weight_desc.spatial_remap,
+          absl::MakeSpan(f16_ptr, flt_count / 4));
+    }
+    return;
+  } else if (dst_weight_desc.layout == WeightsLayout::kOICustomSpatialO4I4) {
+    if (dst_weight_desc.type == DataType::FLOAT32) {
+      float4* f32_ptr = reinterpret_cast<float4*>(dst.data());
+      RearrangeWeightsToOICustomSpatialO4I4(
+          weights, dst_weight_desc.spatial_remap,
+          absl::MakeSpan(f32_ptr, flt_count / 4));
+    } else if (dst_weight_desc.type == DataType::FLOAT16) {
+      half4* f16_ptr = reinterpret_cast<half4*>(dst.data());
+      RearrangeWeightsToOICustomSpatialO4I4(
+          weights, dst_weight_desc.spatial_remap,
+          absl::MakeSpan(f16_ptr, flt_count / 4));
+    }
+    return;
+  } else if (dst_weight_desc.layout ==
+             WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4) {
+    if (dst_weight_desc.type == DataType::FLOAT32) {
+      float4* f32_ptr = reinterpret_cast<float4*>(dst.data());
+      RearrangeWeightsToI4DHWIOOGroupO4(weights,
+                                        dst_weight_desc.output_group_size,
+                                        absl::MakeSpan(f32_ptr, flt_count / 4));
+    } else if (dst_weight_desc.type == DataType::FLOAT16) {
+      half4* f16_ptr = reinterpret_cast<half4*>(dst.data());
+      RearrangeWeightsToI4DHWIOOGroupO4(weights,
+                                        dst_weight_desc.output_group_size,
+                                        absl::MakeSpan(f16_ptr, flt_count / 4));
+    }
+    return;
+  } else if (dst_weight_desc.layout ==
+             WeightsLayout::k2DX4O4YIsSpatialIAndXIsOOGroupI4) {
+    if (dst_weight_desc.type == DataType::FLOAT32) {
+      float4* f32_ptr = reinterpret_cast<float4*>(dst.data());
+      RearrangeWeightsToO4DHWIOOGroupI4(weights,
+                                        dst_weight_desc.output_group_size,
+                                        absl::MakeSpan(f32_ptr, flt_count / 4));
+    } else if (dst_weight_desc.type == DataType::FLOAT16) {
+      half4* f16_ptr = reinterpret_cast<half4*>(dst.data());
+      RearrangeWeightsToO4DHWIOOGroupI4(weights,
+                                        dst_weight_desc.output_group_size,
+                                        absl::MakeSpan(f16_ptr, flt_count / 4));
+    }
+    return;
+  }
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/task/weights_conversion.h

@@ -0,0 +1,516 @@
+/* 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_TASK_WEIGHTS_CONVERSION_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_WEIGHTS_CONVERSION_H_
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/types/span.h"
+#include "tensorflow/lite/delegates/gpu/common/data_type.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/tensor.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+#include "tensorflow/lite/delegates/gpu/common/util.h"
+
+namespace tflite {
+namespace gpu {
+
+using uint = unsigned int;
+
+template <DataType S, typename T>
+void RearrangeWeightsToOHWIOGroupI4O4(
+    const tflite::gpu::Tensor<OHWI, S>& weights, int out_group_size,
+    absl::Span<T> dst) {
+  const int dst_slices = DivideRoundUp(weights.shape.o, 4);
+  const int src_slices = DivideRoundUp(weights.shape.i, 4);
+  const int dst_groups = DivideRoundUp(dst_slices, out_group_size);
+
+  int counter = 0;
+  for (int d = 0; d < dst_groups; ++d) {
+    for (int y = 0; y < weights.shape.h; ++y) {
+      for (int x = 0; x < weights.shape.w; ++x) {
+        for (int s = 0; s < src_slices; ++s) {
+          for (int d_group = 0; d_group < out_group_size; ++d_group) {
+            for (int j = 0; j < 4; ++j) {
+              T filter;
+              for (int i = 0; i < 4; ++i) {
+                const int s_ch = s * 4 + j;
+                const int d_ch = (d * out_group_size + d_group) * 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});
+                  filter[i] = weights.data[f_index];
+                } else {
+                  filter[i] = 0.0f;
+                }
+              }
+              dst[counter++] = filter;
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+template <DataType S, typename T>
+void RearrangeWeightsToODHWIOGroupI4O4(
+    const tflite::gpu::Tensor<OHWDI, S>& weights, int out_group_size,
+    absl::Span<T> dst) {
+  const int dst_slices = DivideRoundUp(weights.shape.o, 4);
+  const int src_slices = DivideRoundUp(weights.shape.i, 4);
+  const int dst_groups = DivideRoundUp(dst_slices, out_group_size);
+
+  int counter = 0;
+  for (int d = 0; d < dst_groups; ++d) {
+    for (int z = 0; z < weights.shape.d; ++z) {
+      for (int y = 0; y < weights.shape.h; ++y) {
+        for (int x = 0; x < weights.shape.w; ++x) {
+          for (int s = 0; s < src_slices; ++s) {
+            for (int d_group = 0; d_group < out_group_size; ++d_group) {
+              for (int j = 0; j < 4; ++j) {
+                T filter;
+                for (int i = 0; i < 4; ++i) {
+                  const int s_ch = s * 4 + j;
+                  const int d_ch = (d * out_group_size + d_group) * 4 + i;
+                  if (s_ch < weights.shape.i && d_ch < weights.shape.o) {
+                    const int f_index =
+                        weights.shape.LinearIndex({d_ch, y, x, z, s_ch});
+                    filter[i] = weights.data[f_index];
+                  } else {
+                    filter[i] = 0.0f;
+                  }
+                }
+                dst[counter++] = filter;
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+template <DataType S, typename T>
+void RearrangeWeightsToOHWIOGroupO4I4(
+    const tflite::gpu::Tensor<OHWI, S>& weights, int out_group_size,
+    absl::Span<T> dst) {
+  const int dst_slices = DivideRoundUp(weights.shape.o, 4);
+  const int src_slices = DivideRoundUp(weights.shape.i, 4);
+  const int dst_groups = DivideRoundUp(dst_slices, out_group_size);
+
+  int counter = 0;
+  for (int d = 0; d < dst_groups; ++d) {
+    for (int y = 0; y < weights.shape.h; ++y) {
+      for (int x = 0; x < weights.shape.w; ++x) {
+        for (int s = 0; s < src_slices; ++s) {
+          for (int d_group = 0; d_group < out_group_size; ++d_group) {
+            for (int j = 0; j < 4; ++j) {
+              T filter;
+              for (int i = 0; i < 4; ++i) {
+                const int s_ch = s * 4 + i;
+                const int d_ch = (d * out_group_size + d_group) * 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});
+                  filter[i] = weights.data[f_index];
+                } else {
+                  filter[i] = 0.0f;
+                }
+              }
+              dst[counter++] = filter;
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+template <DataType S, typename T>
+void RearrangeWeightsToODHWIOGroupO4I4(
+    const tflite::gpu::Tensor<OHWDI, S>& weights, int out_group_size,
+    absl::Span<T> dst) {
+  const int dst_slices = DivideRoundUp(weights.shape.o, 4);
+  const int src_slices = DivideRoundUp(weights.shape.i, 4);
+  const int dst_groups = DivideRoundUp(dst_slices, out_group_size);
+
+  int counter = 0;
+  for (int d = 0; d < dst_groups; ++d) {
+    for (int z = 0; z < weights.shape.d; ++z) {
+      for (int y = 0; y < weights.shape.h; ++y) {
+        for (int x = 0; x < weights.shape.w; ++x) {
+          for (int s = 0; s < src_slices; ++s) {
+            for (int d_group = 0; d_group < out_group_size; ++d_group) {
+              for (int j = 0; j < 4; ++j) {
+                T filter;
+                for (int i = 0; i < 4; ++i) {
+                  const int s_ch = s * 4 + i;
+                  const int d_ch = (d * out_group_size + d_group) * 4 + j;
+                  if (s_ch < weights.shape.i && d_ch < weights.shape.o) {
+                    const int f_index =
+                        weights.shape.LinearIndex({d_ch, y, x, z, s_ch});
+                    filter[i] = weights.data[f_index];
+                  } else {
+                    filter[i] = 0.0f;
+                  }
+                }
+                dst[counter++] = filter;
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+template <DataType S, typename T>
+void RearrangeWeightsToI4HWIOOGroupO4(
+    const tflite::gpu::Tensor<OHWI, S>& weights, int out_group_size,
+    absl::Span<T> dst) {
+  const int dst_slices = DivideRoundUp(weights.shape.o, 4);
+  const int src_slices = DivideRoundUp(weights.shape.i, 4);
+  const int dst_groups = DivideRoundUp(dst_slices, out_group_size);
+
+  int counter = 0;
+  for (int j = 0; j < 4; ++j) {
+    for (int y = 0; y < weights.shape.h; ++y) {
+      for (int x = 0; x < weights.shape.w; ++x) {
+        for (int s = 0; s < src_slices; ++s) {
+          for (int d = 0; d < dst_groups; ++d) {
+            for (int d_group = 0; d_group < out_group_size; ++d_group) {
+              T filter;
+              for (int i = 0; i < 4; ++i) {
+                const int s_ch = s * 4 + j;
+                const int d_ch = (d * out_group_size + d_group) * 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});
+                  filter[i] = weights.data[f_index];
+                } else {
+                  filter[i] = 0.0f;
+                }
+              }
+              dst[counter++] = filter;
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+template <DataType S, typename T>
+void RearrangeWeightsToI4DHWIOOGroupO4(
+    const tflite::gpu::Tensor<OHWDI, S>& weights, int out_group_size,
+    absl::Span<T> dst) {
+  const int dst_slices = DivideRoundUp(weights.shape.o, 4);
+  const int src_slices = DivideRoundUp(weights.shape.i, 4);
+  const int dst_groups = DivideRoundUp(dst_slices, out_group_size);
+
+  int counter = 0;
+  for (int j = 0; j < 4; ++j) {
+    for (int z = 0; z < weights.shape.d; ++z) {
+      for (int y = 0; y < weights.shape.h; ++y) {
+        for (int x = 0; x < weights.shape.w; ++x) {
+          for (int s = 0; s < src_slices; ++s) {
+            for (int d = 0; d < dst_groups; ++d) {
+              for (int d_group = 0; d_group < out_group_size; ++d_group) {
+                T filter;
+                for (int i = 0; i < 4; ++i) {
+                  const int s_ch = s * 4 + j;
+                  const int d_ch = (d * out_group_size + d_group) * 4 + i;
+                  if (s_ch < weights.shape.i && d_ch < weights.shape.o) {
+                    const int f_index =
+                        weights.shape.LinearIndex({d_ch, y, x, z, s_ch});
+                    filter[i] = weights.data[f_index];
+                  } else {
+                    filter[i] = 0.0f;
+                  }
+                }
+                dst[counter++] = filter;
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+template <DataType S, typename T>
+void RearrangeWeightsToO4HWIOOGroupI4(
+    const tflite::gpu::Tensor<OHWI, S>& weights, int out_group_size,
+    absl::Span<T> dst) {
+  const int dst_slices = DivideRoundUp(weights.shape.o, 4);
+  const int src_slices = DivideRoundUp(weights.shape.i, 4);
+  const int dst_groups = DivideRoundUp(dst_slices, out_group_size);
+
+  int counter = 0;
+  for (int j = 0; j < 4; ++j) {
+    for (int y = 0; y < weights.shape.h; ++y) {
+      for (int x = 0; x < weights.shape.w; ++x) {
+        for (int s = 0; s < src_slices; ++s) {
+          for (int d = 0; d < dst_groups; ++d) {
+            for (int d_group = 0; d_group < out_group_size; ++d_group) {
+              T filter;
+              for (int i = 0; i < 4; ++i) {
+                const int s_ch = s * 4 + i;
+                const int d_ch = (d * out_group_size + d_group) * 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});
+                  filter[i] = weights.data[f_index];
+                } else {
+                  filter[i] = 0.0f;
+                }
+              }
+              dst[counter++] = filter;
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+template <DataType S, typename T>
+void RearrangeWeightsToO4DHWIOOGroupI4(
+    const tflite::gpu::Tensor<OHWDI, S>& weights, int out_group_size,
+    absl::Span<T> dst) {
+  const int dst_slices = DivideRoundUp(weights.shape.o, 4);
+  const int src_slices = DivideRoundUp(weights.shape.i, 4);
+  const int dst_groups = DivideRoundUp(dst_slices, out_group_size);
+
+  int counter = 0;
+  for (int j = 0; j < 4; ++j) {
+    for (int z = 0; z < weights.shape.d; ++z) {
+      for (int y = 0; y < weights.shape.h; ++y) {
+        for (int x = 0; x < weights.shape.w; ++x) {
+          for (int s = 0; s < src_slices; ++s) {
+            for (int d = 0; d < dst_groups; ++d) {
+              for (int d_group = 0; d_group < out_group_size; ++d_group) {
+                T filter;
+                for (int i = 0; i < 4; ++i) {
+                  const int s_ch = s * 4 + i;
+                  const int d_ch = (d * out_group_size + d_group) * 4 + j;
+                  if (s_ch < weights.shape.i && d_ch < weights.shape.o) {
+                    const int f_index =
+                        weights.shape.LinearIndex({d_ch, y, x, z, s_ch});
+                    filter[i] = weights.data[f_index];
+                  } else {
+                    filter[i] = 0.0f;
+                  }
+                }
+                dst[counter++] = filter;
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+template <DataType S, typename T>
+void RearrangeWeightsToOICustomSpatialI4O4(
+    const tflite::gpu::Tensor<OHWI, S>& weights,
+    const std::vector<int>& spatial_remap, absl::Span<T> dst) {
+  const int dst_slices = DivideRoundUp(weights.shape.o, 4);
+  const int src_slices = DivideRoundUp(weights.shape.i, 4);
+
+  int counter = 0;
+  for (int d = 0; d < dst_slices; ++d) {
+    for (int s = 0; s < src_slices; ++s) {
+      for (int y = 0; y < weights.shape.h; ++y) {
+        for (int x = 0; x < weights.shape.w; ++x) {
+          const int kernel_index = spatial_remap[y * weights.shape.w + x];
+          const int kernel_index_x = kernel_index % weights.shape.w;
+          const int kernel_index_y = kernel_index / weights.shape.w;
+          for (int i = 0; i < 4; ++i) {
+            T filter;
+            for (int j = 0; j < 4; ++j) {
+              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, kernel_index_y, kernel_index_x, s_ch});
+                filter[j] = weights.data[f_index];
+              } else {
+                filter[j] = 0.0f;
+              }
+            }
+            dst[counter++] = filter;
+          }
+        }
+      }
+    }
+  }
+}
+
+template <DataType S, typename T>
+void RearrangeWeightsToOICustomSpatialI4O4(
+    const tflite::gpu::Tensor<OHWDI, S>& weights,
+    const std::vector<int>& spatial_remap, absl::Span<T> dst) {
+  const int dst_slices = DivideRoundUp(weights.shape.o, 4);
+  const int src_slices = DivideRoundUp(weights.shape.i, 4);
+
+  int counter = 0;
+  for (int d = 0; d < dst_slices; ++d) {
+    for (int s = 0; s < src_slices; ++s) {
+      for (int z = 0; z < weights.shape.d; ++z) {
+        for (int y = 0; y < weights.shape.h; ++y) {
+          for (int x = 0; x < weights.shape.w; ++x) {
+            int kernel_index =
+                spatial_remap[(z * weights.shape.h + y) * weights.shape.w + x];
+            const int kernel_index_x = kernel_index % weights.shape.w;
+            kernel_index /= weights.shape.w;
+            const int kernel_index_y = kernel_index % weights.shape.h;
+            const int kernel_index_z = kernel_index / weights.shape.h;
+            for (int i = 0; i < 4; ++i) {
+              T filter;
+              for (int j = 0; j < 4; ++j) {
+                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, kernel_index_y, kernel_index_x, kernel_index_z,
+                       s_ch});
+                  filter[j] = weights.data[f_index];
+                } else {
+                  filter[j] = 0.0f;
+                }
+              }
+              dst[counter++] = filter;
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+template <DataType S, typename T>
+void RearrangeWeightsToOICustomSpatialO4I4(
+    const tflite::gpu::Tensor<OHWI, S>& weights,
+    const std::vector<int>& spatial_remap, absl::Span<T> dst) {
+  const int dst_slices = DivideRoundUp(weights.shape.o, 4);
+  const int src_slices = DivideRoundUp(weights.shape.i, 4);
+
+  int counter = 0;
+  for (int d = 0; d < dst_slices; ++d) {
+    for (int s = 0; s < src_slices; ++s) {
+      for (int y = 0; y < weights.shape.h; ++y) {
+        for (int x = 0; x < weights.shape.w; ++x) {
+          const int kernel_index = spatial_remap[y * weights.shape.w + x];
+          const int kernel_index_x = kernel_index % weights.shape.w;
+          const int kernel_index_y = kernel_index / weights.shape.w;
+          for (int i = 0; i < 4; ++i) {
+            T filter;
+            for (int j = 0; j < 4; ++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, kernel_index_y, kernel_index_x, s_ch});
+                filter[j] = weights.data[f_index];
+              } else {
+                filter[j] = 0.0f;
+              }
+            }
+            dst[counter++] = filter;
+          }
+        }
+      }
+    }
+  }
+}
+
+template <DataType S, typename T>
+void RearrangeWeightsToOICustomSpatialO4I4(
+    const tflite::gpu::Tensor<OHWDI, S>& weights,
+    const std::vector<int>& spatial_remap, absl::Span<T> dst) {
+  const int dst_slices = DivideRoundUp(weights.shape.o, 4);
+  const int src_slices = DivideRoundUp(weights.shape.i, 4);
+
+  int counter = 0;
+  for (int d = 0; d < dst_slices; ++d) {
+    for (int s = 0; s < src_slices; ++s) {
+      for (int z = 0; z < weights.shape.d; ++z) {
+        for (int y = 0; y < weights.shape.h; ++y) {
+          for (int x = 0; x < weights.shape.w; ++x) {
+            int kernel_index =
+                spatial_remap[(z * weights.shape.h + y) * weights.shape.w + x];
+            const int kernel_index_x = kernel_index % weights.shape.w;
+            kernel_index /= weights.shape.w;
+            const int kernel_index_y = kernel_index % weights.shape.h;
+            const int kernel_index_z = kernel_index / weights.shape.h;
+            for (int i = 0; i < 4; ++i) {
+              T filter;
+              for (int j = 0; j < 4; ++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, kernel_index_y, kernel_index_x, kernel_index_z,
+                       s_ch});
+                  filter[j] = weights.data[f_index];
+                } else {
+                  filter[j] = 0.0f;
+                }
+              }
+              dst[counter++] = filter;
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+uint GetTotalElementsCountForLayout(const WeightsDescription& weight_desc,
+                                    const OHWI& shape);
+uint GetTotalElementsCountForLayout(const WeightsDescription& weight_desc,
+                                    const OHWDI& shape);
+
+// Applicable to:
+//   k2DX4I4YIsSpatialIAndXIsOOGroupO4
+//   k2DX4O4YIsSpatialIAndXIsOOGroupI4
+uint2 Get2dResourceSize(const WeightsDescription& weight_desc,
+                        const OHWI& shape);
+// Applicable to:
+//   k2DX4I4YIsSpatialIAndXIsOOGroupO4
+//   k2DX4O4YIsSpatialIAndXIsOOGroupI4
+uint2 Get2dResourceSize(const WeightsDescription& weight_desc,
+                        const OHWDI& shape);
+
+void RearrangeWeights(
+    const tflite::gpu::Tensor<OHWI, DataType::FLOAT32>& weights,
+    const WeightsDescription& dst_weight_desc, absl::Span<uint8_t> dst);
+
+void RearrangeWeights(
+    const tflite::gpu::Tensor<OHWDI, DataType::FLOAT32>& weights,
+    const WeightsDescription& dst_weight_desc, absl::Span<uint8_t> dst);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_WEIGHTS_CONVERSION_H_

tensorflow/lite/delegates/gpu/common/task/weights_conversion_test.cc

@@ -0,0 +1,188 @@
+/* Copyright 2026 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/task/weights_conversion.h"
+
+#include <cstdint>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "absl/types/span.h"
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+#include "tensorflow/lite/delegates/gpu/common/shape.h"
+#include "tensorflow/lite/delegates/gpu/common/task/weights_layout.h"
+#include "tensorflow/lite/delegates/gpu/common/tensor.h"
+
+namespace tflite {
+namespace gpu {
+namespace {
+
+TEST(WeightsConversionTest, GetTotalElementsCountForLayoutOverflow_k2DX4I4) {
+  WeightsDescription weight_desc;
+  weight_desc.type = DataType::FLOAT32;
+  weight_desc.layout = WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4;
+  weight_desc.output_group_size = 1;
+
+  OHWI shape(1, 1, 1, 2000000000);
+  EXPECT_EQ(GetTotalElementsCountForLayout(weight_desc, shape), 0);
+}
+
+TEST(WeightsConversionTest, GetTotalElementsCountForLayoutOverflow_kOICustom) {
+  WeightsDescription weight_desc;
+  weight_desc.type = DataType::FLOAT32;
+  weight_desc.layout = WeightsLayout::kOICustomSpatialI4O4;
+  weight_desc.spatial_remap.resize(1);
+
+  OHWI shape(1, 1, 1, 2000000000);
+  EXPECT_EQ(GetTotalElementsCountForLayout(weight_desc, shape), 0);
+}
+
+TEST(WeightsConversionTest,
+     GetTotalElementsCountForLayoutOverflow_OHWDI_k2DX4I4) {
+  WeightsDescription weight_desc;
+  weight_desc.type = DataType::FLOAT32;
+  weight_desc.layout = WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4;
+  weight_desc.output_group_size = 1;
+
+  OHWDI shape(1, 1, 1, 1, 2000000000);
+  EXPECT_EQ(GetTotalElementsCountForLayout(weight_desc, shape), 0);
+}
+
+TEST(WeightsConversionTest,
+     GetTotalElementsCountForLayoutOverflow_OHWDI_kOICustom) {
+  WeightsDescription weight_desc;
+  weight_desc.type = DataType::FLOAT32;
+  weight_desc.layout = WeightsLayout::kOICustomSpatialI4O4;
+  weight_desc.spatial_remap.resize(1);
+
+  OHWDI shape(1, 1, 1, 1, 2000000000);
+  EXPECT_EQ(GetTotalElementsCountForLayout(weight_desc, shape), 0);
+}
+
+TEST(WeightsConversionTest, RearrangeWeightsEarlyReturnOnOverflow_k2DX4I4) {
+  WeightsDescription weight_desc;
+  weight_desc.type = DataType::FLOAT32;
+  weight_desc.layout = WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4;
+  weight_desc.output_group_size = 1;
+
+  Tensor<OHWI, DataType::FLOAT32> weights;
+  weights.shape = OHWI(1, 1, 1, 2000000000);
+
+  std::vector<uint8_t> dummy_output(4, 0xAB);
+  RearrangeWeights(weights, weight_desc, absl::MakeSpan(dummy_output));
+  EXPECT_EQ(dummy_output[0], 0xAB);
+  EXPECT_EQ(dummy_output[1], 0xAB);
+  EXPECT_EQ(dummy_output[2], 0xAB);
+  EXPECT_EQ(dummy_output[3], 0xAB);
+}
+
+TEST(WeightsConversionTest, RearrangeWeightsEarlyReturnOnOverflow_OHWDI) {
+  WeightsDescription weight_desc;
+  weight_desc.type = DataType::FLOAT32;
+  weight_desc.layout = WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4;
+  weight_desc.output_group_size = 1;
+
+  Tensor<OHWDI, DataType::FLOAT32> weights;
+  weights.shape =
+      OHWDI(1, 1, 1, 1, 2000000000);  // O, H, W, D, I? Or O, D, H, W, I?
+
+  std::vector<uint8_t> dummy_output(4, 0xAB);
+  RearrangeWeights(weights, weight_desc, absl::MakeSpan(dummy_output));
+  EXPECT_EQ(dummy_output[0], 0xAB);
+  EXPECT_EQ(dummy_output[1], 0xAB);
+  EXPECT_EQ(dummy_output[2], 0xAB);
+  EXPECT_EQ(dummy_output[3], 0xAB);
+}
+
+TEST(WeightsConversionTest, GetTotalElementsCountForLayoutNormal_k2DX4I4) {
+  WeightsDescription weight_desc;
+  weight_desc.type = DataType::FLOAT32;
+  weight_desc.layout = WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4;
+  weight_desc.output_group_size = 1;
+
+  OHWI shape(2, 2, 2, 8);
+  EXPECT_EQ(GetTotalElementsCountForLayout(weight_desc, shape), 128);
+}
+
+TEST(WeightsConversionTest, GetTotalElementsCountForLayoutNormal_kOICustom) {
+  WeightsDescription weight_desc;
+  weight_desc.type = DataType::FLOAT32;
+  weight_desc.layout = WeightsLayout::kOICustomSpatialI4O4;
+  weight_desc.spatial_remap.resize(3);
+
+  OHWI shape(2, 2, 2, 8);
+  EXPECT_EQ(GetTotalElementsCountForLayout(weight_desc, shape), 96);
+}
+
+TEST(WeightsConversionTest, GetTotalElementsCountForLayoutUnknown) {
+  WeightsDescription weight_desc;
+  weight_desc.type = DataType::FLOAT32;
+  weight_desc.layout = WeightsLayout::kUnknown;
+
+  OHWI shape(1, 1, 1, 1);
+  EXPECT_EQ(GetTotalElementsCountForLayout(weight_desc, shape),
+            static_cast<uint>(-1));
+}
+
+TEST(WeightsConversionTest, RearrangeWeightsNormal_k2DX4I4) {
+  WeightsDescription weight_desc;
+  weight_desc.type = DataType::FLOAT32;
+  weight_desc.layout = WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4;
+  weight_desc.output_group_size = 1;
+
+  Tensor<OHWI, DataType::FLOAT32> weights;
+  weights.shape = OHWI(1, 1, 1, 1);
+  weights.data.resize(1);
+
+  std::vector<uint8_t> dummy_output(64);
+  RearrangeWeights(weights, weight_desc, absl::MakeSpan(dummy_output));
+}
+
+TEST(WeightsConversionTest, GetTotalElementsCountForLayout_Valid_k2DX4I4) {
+  WeightsDescription weight_desc;
+  weight_desc.type = DataType::FLOAT32;
+  weight_desc.layout = WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4;
+  weight_desc.output_group_size = 1;
+
+  OHWI shape(1, 1, 1, 1);
+  // i_aligned = 4, o_aligned = 4, h=1, w=1, d=1
+  // total = 4 * 4 * 1 * 1 * 1 = 16
+  EXPECT_EQ(GetTotalElementsCountForLayout(weight_desc, shape), 16);
+}
+
+TEST(WeightsConversionTest, GetTotalElementsCountForLayout_Valid_kOICustom) {
+  WeightsDescription weight_desc;
+  weight_desc.type = DataType::FLOAT32;
+  weight_desc.layout = WeightsLayout::kOICustomSpatialI4O4;
+  weight_desc.spatial_remap.resize(2);
+
+  OHWI shape(1, 1, 1, 1);
+  // i_aligned = 4, o_aligned = 4, remap_size = 2
+  // total = 4 * 4 * 2 = 32
+  EXPECT_EQ(GetTotalElementsCountForLayout(weight_desc, shape), 32);
+}
+
+TEST(WeightsConversionTest, GetTotalElementsCountForLayout_UnknownLayout) {
+  WeightsDescription weight_desc;
+  weight_desc.type = DataType::FLOAT32;
+  weight_desc.layout = WeightsLayout::kUnknown;
+  OHWI shape(1, 1, 1, 1);
+  EXPECT_EQ(GetTotalElementsCountForLayout(weight_desc, shape),
+            static_cast<uint>(-1));
+}
+
+}  // namespace
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/task/weights_layout.cc

@@ -0,0 +1,57 @@
+/* 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/task/weights_layout.h"
+
+namespace tflite {
+namespace gpu {
+
+int WeightsDescription::GetOutputGroupSize() const {
+  if (layout == WeightsLayout::kOSpatialIOGroupI4O4 ||
+      layout == WeightsLayout::kOSpatialIOGroupO4I4 ||
+      layout == WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4 ||
+      layout == WeightsLayout::k2DX4O4YIsSpatialIAndXIsOOGroupI4) {
+    return output_group_size;
+  } else {
+    return 1;
+  }
+}
+
+bool WeightsDescription::IsI4O4() const {
+  return layout == WeightsLayout::kOSpatialIOGroupI4O4 ||
+         layout == WeightsLayout::kOICustomSpatialI4O4 ||
+         layout == WeightsLayout::k2DX4I4YIsSpatialIAndXIsOOGroupO4;
+}
+
+bool WeightsDescription::IsO4I4() const {
+  return layout == WeightsLayout::kOSpatialIOGroupO4I4 ||
+         layout == WeightsLayout::kOICustomSpatialO4I4 ||
+         layout == WeightsLayout::k2DX4O4YIsSpatialIAndXIsOOGroupI4;
+}
+
+bool WeightsDescription::IsCustomSpatial() const {
+  return layout == WeightsLayout::kOICustomSpatialI4O4 ||
+         layout == WeightsLayout::kOICustomSpatialO4I4;
+}
+
+bool WeightsDescription::operator==(const WeightsDescription& t) const {
+  const bool equal_spatial_remap =
+      IsCustomSpatial() ? spatial_remap == t.spatial_remap : true;
+  return type == t.type && layout == t.layout &&
+         GetOutputGroupSize() == t.GetOutputGroupSize() && equal_spatial_remap;
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/task/weights_layout.h

@@ -0,0 +1,57 @@
+/* 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_TASK_WEIGHTS_LAYOUT_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_WEIGHTS_LAYOUT_H_
+
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/data_type.h"
+
+namespace tflite {
+namespace gpu {
+
+enum class WeightsLayout {
+  kUnknown,
+  // Spatial is DHW/HW depending on amount of spatial dimensions (Depth, Height,
+  // Width).
+  kOSpatialIOGroupI4O4,
+  kOSpatialIOGroupO4I4,
+  kOICustomSpatialI4O4,
+  kOICustomSpatialO4I4,
+  k2DX4I4YIsSpatialIAndXIsOOGroupO4,
+  k2DX4O4YIsSpatialIAndXIsOOGroupI4,
+};
+
+struct WeightsDescription {
+  DataType type;
+  WeightsLayout layout;
+  // applicable with layouts that have OGroup.
+  int output_group_size;  // OGroup size
+  // applicable with layouts that have CustomSpatial
+  std::vector<int> spatial_remap;
+
+  int GetOutputGroupSize() const;
+  bool IsI4O4() const;
+  bool IsO4I4() const;
+  bool IsCustomSpatial() const;
+
+  bool operator==(const WeightsDescription& t) const;
+};
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_WEIGHTS_LAYOUT_H_

tensorflow/lite/delegates/gpu/common/task/work_group_picking.cc

@@ -0,0 +1,331 @@
+/* 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/task/work_group_picking.h"
+
+#include <algorithm>
+#include <limits>
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/util.h"
+
+namespace tflite {
+namespace gpu {
+
+namespace {
+
+std::vector<int2> Get2DWorkgroupsEqualTo128() {
+  return {{128, 1}, {64, 2}, {32, 4}, {16, 8},
+          {8, 16},  {4, 32}, {2, 64}, {1, 128}};
+}
+
+std::vector<int3> GenerateWorkGroupSizesXYMultipleOf(
+    int multiplier, int3 grid, const KernelInfo& kernel_info,
+    const GpuInfo& gpu_info, WorkGroupSizeAlignment z_alignment) {
+  std::vector<int3> work_groups;
+  work_groups.reserve(32);
+
+  std::vector<int> possible_z_sizes = GetPossibleSizes(grid.z, z_alignment);
+
+  for (int x = 1; x <= kernel_info.max_work_group_size; x *= 2) {
+    for (int y = 1; y <= kernel_info.max_work_group_size; y *= 2) {
+      int work_group_size_xy = x * y;
+      if (work_group_size_xy % multiplier != 0 ||
+          work_group_size_xy > kernel_info.max_work_group_size) {
+        continue;
+      }
+      for (auto z : possible_z_sizes) {
+        if (work_group_size_xy * z > kernel_info.max_work_group_size) {
+          continue;
+        }
+        if (x <= gpu_info.GetMaxWorkGroupSizeForX() &&
+            y <= gpu_info.GetMaxWorkGroupSizeForY() &&
+            z <= gpu_info.GetMaxWorkGroupSizeForZ()) {
+          work_groups.push_back({x, y, z});
+        }
+      }
+    }
+  }
+  return work_groups;
+}
+
+std::vector<int3> GenerateWorkGroupSizesXMultipleOf(
+    int multiplier, int3 grid, const KernelInfo& kernel_info,
+    const GpuInfo& gpu_info, WorkGroupSizeAlignment z_alignment) {
+  std::vector<int3> work_groups;
+  work_groups.reserve(32);
+
+  std::vector<int> possible_z_sizes = GetPossibleSizes(grid.z, z_alignment);
+  std::vector<int> possible_y_sizes =
+      GetPossibleSizes(grid.y, WorkGroupSizeAlignment::PRECISE);
+
+  for (int x = multiplier;
+       x <= kernel_info.max_work_group_size && x < grid.x + multiplier;
+       x += multiplier) {
+    for (auto y : possible_y_sizes) {
+      for (auto z : possible_z_sizes) {
+        if (x <= gpu_info.GetMaxWorkGroupSizeForX() &&
+            y <= gpu_info.GetMaxWorkGroupSizeForY() &&
+            z <= gpu_info.GetMaxWorkGroupSizeForZ() &&
+            x * y * z <= kernel_info.max_work_group_size) {
+          work_groups.push_back({x, y, z});
+        }
+      }
+    }
+  }
+  return work_groups;
+}
+
+void GetWorkGroupsAlignedToGrid(const GpuInfo& gpu_info,
+                                const KernelInfo& kernel_info, const int3& grid,
+                                std::vector<int3>* work_groups) {
+  int3 max_wg_size;
+  max_wg_size.x = gpu_info.GetMaxWorkGroupSizeForX();
+  max_wg_size.y = gpu_info.GetMaxWorkGroupSizeForY();
+  max_wg_size.z = gpu_info.GetMaxWorkGroupSizeForZ();
+  GenerateWorkGroupSizesAlignedToGrid(
+      grid, max_wg_size, kernel_info.max_work_group_size, work_groups);
+}
+
+int GetPenalty(int grid_size, int group_size) {
+  const int reminder = grid_size % group_size;
+  return reminder == 0 ? 0 : group_size - reminder;
+}
+
+int GetPenalty(int2 grid_size, int2 group_size) {
+  const int p_x = GetPenalty(grid_size.x, group_size.x);
+  const int p_y = GetPenalty(grid_size.y, group_size.y);
+  return p_x * grid_size.y + p_y * grid_size.x + p_x * p_y;
+}
+
+int GetMaxSizeWithMinPenalty(int size, int max_size) {
+  int best_size = 128;
+  int min_penalty = GetPenalty(size, best_size);
+  for (int i = 2; i * 128 <= max_size; ++i) {
+    if (GetPenalty(size, i * 128) == min_penalty) {
+      best_size = i * 128;
+    }
+  }
+  return best_size;
+}
+
+int2 GetMaxSizeWithMinPenalty(int2 size, int max_size) {
+  std::vector<int2> base_groups = Get2DWorkgroupsEqualTo128();
+  int min_penalty = std::numeric_limits<int>::max();
+  for (const auto& group : base_groups) {
+    min_penalty = std::min(GetPenalty(size, group), min_penalty);
+  }
+  for (const auto& group : base_groups) {
+    for (int y = 1; y * group.y <= max_size; ++y) {
+      int new_group_y = y * group.y;
+      for (int x = 1; x * group.x <= max_size; ++x) {
+        int new_group_x = x * group.x;
+        if (new_group_x * new_group_y > max_size) {
+          break;
+        }
+        if (GetPenalty(size, int2(new_group_x, new_group_y)) == min_penalty) {
+          return int2(new_group_x, new_group_y);
+        }
+      }
+    }
+  }
+  return int2(0, 0);
+}
+
+int GetBiggestDividerWithPriority(int number, int max_divider) {
+  if (number % 8 == 0 && 8 <= max_divider) {
+    return 8;
+  }
+  if (number % 4 == 0 && 4 <= max_divider) {
+    return 4;
+  }
+  if (number % 2 == 0 && 2 <= max_divider) {
+    return 2;
+  }
+  for (int i = max_divider; i != 0; i--) {
+    if (number % i == 0) {
+      return i;
+    }
+  }
+  return 1;
+}
+
+int GetBiggestDivider(int number, int max_divider) {
+  for (int i = max_divider; i != 0; i--) {
+    if (number % i == 0) {
+      return i;
+    }
+  }
+  return 1;
+}
+
+int GetOptimalSizeForApple(int grid_size) {
+  if (grid_size % 8 == 0 || grid_size % 8 >= 4 || grid_size >= 16) {
+    return 8;
+  }
+  if (grid_size % 4 == 0 || grid_size % 4 >= 2 || grid_size >= 8) {
+    return 4;
+  }
+  if (grid_size % 2 == 0 || grid_size >= 4) {
+    return 2;
+  }
+  return 1;
+}
+
+int3 GetWorkGroupSizeForApple(const int3& grid_size) {
+  int x_size = GetOptimalSizeForApple(grid_size.x);
+  int y_size = GetOptimalSizeForApple(grid_size.y);
+  int z_size = std::max(1, 32 / (x_size * y_size));
+  z_size = std::min(z_size, static_cast<int>(grid_size.z));
+  return {x_size, y_size, z_size};
+}
+
+}  // namespace
+
+int3 GetWorkGroupXY128ConvLinear(const int3& grid) {
+  int grid_z = GetBiggestDividerWithPriority(grid.z, 4);
+  if (grid.x <= 128) {
+    return int3(128, 1, grid_z);
+  }
+  int grid_x = GetMaxSizeWithMinPenalty(grid.x, 512 / grid_z);
+  return {grid_x, 1, grid_z};
+}
+
+int3 GetWorkGroupXY128Conv(const int3& grid) {
+  int grid_z = GetBiggestDividerWithPriority(grid.z, 4);
+  if (grid.x <= 16 && grid.y <= 8) {
+    return int3(16, 8, grid_z);
+  }
+  int2 grid_xy = GetMaxSizeWithMinPenalty(int2(grid.x, grid.y), 512 / grid_z);
+  return int3(grid_xy.x, grid_xy.y, grid_z);
+}
+
+int3 GetWorkGroupXY128Simple(const int3& grid) { return int3(16, 8, 1); }
+
+int3 GetWorkGroup(const int3& grid, int max_size) {
+  int wg_z = GetBiggestDividerWithPriority(grid.z, 8);
+  int wg_xy_size = max_size / wg_z;
+  int wg_x = std::min(DivideRoundUp(grid.x, 2), wg_xy_size);
+  int wg_y = std::min(wg_xy_size / wg_x, grid.y);
+  return int3(wg_x, wg_y, wg_z);
+}
+
+int3 GetWorkGroupConv(const int3& grid, int max_size, int max_z_size) {
+  int wg_z = GetBiggestDivider(grid.z, max_z_size);
+  int wg_xy_size = std::min(256, max_size) / wg_z;
+  int wg_x = std::min(grid.x, wg_xy_size);
+  int wg_y = std::min(wg_xy_size / wg_x, grid.y);
+  if (wg_y == grid.y && grid.y % 2 == 0) {
+    wg_y = grid.y / 2;
+  }
+  return int3(wg_x, wg_y, wg_z);
+}
+
+void GetPossibleWorkGroupsXYMultipleOf(int multiplier, const GpuInfo& gpu_info,
+                                       const KernelInfo& kernel_info,
+                                       const int3& grid,
+                                       WorkGroupSizeAlignment z_alignment,
+                                       std::vector<int3>* work_groups) {
+  *work_groups = GenerateWorkGroupSizesXYMultipleOf(
+      multiplier, grid, kernel_info, gpu_info, z_alignment);
+}
+
+void GetPossibleWorkGroupsXMultipleOf(int multiplier, const GpuInfo& gpu_info,
+                                      const KernelInfo& kernel_info,
+                                      const int3& grid,
+                                      WorkGroupSizeAlignment z_alignment,
+                                      std::vector<int3>* work_groups) {
+  *work_groups = GenerateWorkGroupSizesXMultipleOf(
+      multiplier, grid, kernel_info, gpu_info, z_alignment);
+}
+
+bool XY128RequiresMoreWorkGroupsThenXY128Linear(int width, int height) {
+  int planar_work_groups = DivideRoundUp(width * height, 128);
+  auto base_work_groups = Get2DWorkgroupsEqualTo128();
+  bool have_equal_work_groups = false;
+  for (auto& work_group : base_work_groups) {
+    int x_groups = DivideRoundUp(width, work_group.x);
+    int y_groups = DivideRoundUp(height, work_group.y);
+    int xy_groups = x_groups * y_groups;
+    if (xy_groups == planar_work_groups) {
+      have_equal_work_groups = true;
+      break;
+    }
+  }
+  return !have_equal_work_groups;
+}
+
+void GetPossibleWorkGroups(TuningType tuning_type, const GpuInfo& gpu_info,
+                           const KernelInfo& kernel_info, const int3& grid,
+                           std::vector<int3>* work_groups) {
+  if (gpu_info.IsApple()) {
+    work_groups->push_back(GetWorkGroupSizeForApple(grid));
+    return;
+  }
+  switch (tuning_type) {
+    case TuningType::kFast:
+      work_groups->push_back(
+          GetWorkGroup(grid, kernel_info.max_work_group_size));
+      return;
+    case TuningType::kExhaustive: {
+      GetWorkGroupsAlignedToGrid(gpu_info, kernel_info, grid, work_groups);
+      return;
+    }
+    default:
+      work_groups->push_back({8, 4, 1});
+      return;
+  }
+}
+
+void GetPossibleWorkGroupsConv(TuningType tuning_type, const GpuInfo& gpu_info,
+                               const KernelInfo& kernel_info, const int3& grid,
+                               std::vector<int3>* work_groups) {
+  if (gpu_info.IsApple()) {
+    work_groups->push_back(GetWorkGroupSizeForApple(grid));
+    return;
+  }
+  switch (tuning_type) {
+    case TuningType::kFast: {
+      int max_z_size = 16;
+      if (gpu_info.IsAdreno()) {
+        max_z_size = gpu_info.adreno_info.IsAdreno3xx() ? 16 : 64;
+      }
+      max_z_size = std::min(max_z_size, gpu_info.GetMaxWorkGroupSizeForZ());
+      work_groups->push_back(
+          GetWorkGroupConv(grid, kernel_info.max_work_group_size, max_z_size));
+      return;
+    }
+    case TuningType::kExhaustive: {
+      GetWorkGroupsAlignedToGrid(gpu_info, kernel_info, grid, work_groups);
+      return;
+    }
+    default:
+      work_groups->push_back({8, 4, 1});
+      return;
+  }
+}
+
+int3 GetFirstSuitableWorkGroup(const std::vector<int3>& wgs, int max_wg_size) {
+  for (const auto& wg : wgs) {
+    const int wg_size = wg.x * wg.y * wg.z;
+    if (wg_size <= max_wg_size) {
+      return wg;
+    }
+  }
+  return {1, 1, 1};
+}
+
+}  // namespace gpu
+}  // namespace tflite

tensorflow/lite/delegates/gpu/common/task/work_group_picking.h

@@ -0,0 +1,65 @@
+/* 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_TASK_WORK_GROUP_PICKING_H_
+#define TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_WORK_GROUP_PICKING_H_
+
+#include <vector>
+
+#include "tensorflow/lite/delegates/gpu/common/gpu_info.h"
+#include "tensorflow/lite/delegates/gpu/common/kernel_info.h"
+#include "tensorflow/lite/delegates/gpu/common/task/tuning_type.h"
+#include "tensorflow/lite/delegates/gpu/common/types.h"
+#include "tensorflow/lite/delegates/gpu/common/workgroup_selection.h"
+
+namespace tflite {
+namespace gpu {
+
+// multiplier can be power of two only
+void GetPossibleWorkGroupsXYMultipleOf(int multiplier, const GpuInfo& gpu_info,
+                                       const KernelInfo& kernel_info,
+                                       const int3& grid,
+                                       WorkGroupSizeAlignment z_alignment,
+                                       std::vector<int3>* work_groups);
+
+void GetPossibleWorkGroupsXMultipleOf(int multiplier, const GpuInfo& gpu_info,
+                                      const KernelInfo& kernel_info,
+                                      const int3& grid,
+                                      WorkGroupSizeAlignment z_alignment,
+                                      std::vector<int3>* work_groups);
+
+int3 GetWorkGroupXY128ConvLinear(const int3& grid);
+
+int3 GetWorkGroupXY128Simple(const int3& grid);
+int3 GetWorkGroupXY128Conv(const int3& grid);
+
+bool XY128RequiresMoreWorkGroupsThenXY128Linear(int width, int height);
+
+void GetPossibleWorkGroups(TuningType tuning_type, const GpuInfo& gpu_info,
+                           const KernelInfo& kernel_info, const int3& grid,
+                           std::vector<int3>* work_groups);
+
+void GetPossibleWorkGroupsConv(TuningType tuning_type, const GpuInfo& gpu_info,
+                               const KernelInfo& kernel_info, const int3& grid,
+                               std::vector<int3>* work_groups);
+
+// returns first work group from wgs that has size not bigger than max_wg_size
+// if no suitable groups among wgs, returns {1, 1, 1}
+int3 GetFirstSuitableWorkGroup(const std::vector<int3>& wgs, int max_wg_size);
+
+}  // namespace gpu
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_DELEGATES_GPU_COMMON_TASK_WORK_GROUP_PICKING_H_