From 811e7b67c88b4a64ddc2eacb67f809b6a8f79cda Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 Oct 2019 12:49:28 -0700
Subject: [PATCH] Added support of batch to all operations with any storage
 types. Changed layout for BUFFER and IMAGE_BUFFER

PiperOrigin-RevId: 272498500
---
 tensorflow/lite/delegates/gpu/cl/environment.cc |  8 --------
 tensorflow/lite/delegates/gpu/cl/environment.h  |  1 -
 .../lite/delegates/gpu/cl/kernels/util.cc       | 17 ++++++++++++++---
 tensorflow/lite/delegates/gpu/cl/tensor.h       |  2 --
 4 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/cl/environment.cc b/tensorflow/lite/delegates/gpu/cl/environment.cc
index d9978a6739e..7677e06c2f9 100644
--- a/tensorflow/lite/delegates/gpu/cl/environment.cc
+++ b/tensorflow/lite/delegates/gpu/cl/environment.cc
@@ -182,14 +182,6 @@ std::vector<TensorStorageType> Environment::GetSupportedTextureStorages()
   return storage_types;
 }
 
-std::vector<TensorStorageType> Environment::GetSupportedBatchStorages() const {
-  std::vector<TensorStorageType> storage_types = {TensorStorageType::BUFFER};
-  if (device_.IsAdreno() && device_.SupportsImageBuffer()) {
-    storage_types.push_back(TensorStorageType::IMAGE_BUFFER);
-  }
-  return storage_types;
-}
-
 std::vector<TensorStorageType> Environment::GetSupportedStorages() const {
   std::vector<TensorStorageType> storage_types = {TensorStorageType::TEXTURE_2D,
                                                   TensorStorageType::BUFFER};
diff --git a/tensorflow/lite/delegates/gpu/cl/environment.h b/tensorflow/lite/delegates/gpu/cl/environment.h
index 99eb702133a..80df6a96deb 100644
--- a/tensorflow/lite/delegates/gpu/cl/environment.h
+++ b/tensorflow/lite/delegates/gpu/cl/environment.h
@@ -56,7 +56,6 @@ class Environment {
   std::vector<CalculationsPrecision> GetSupportedPrecisions() const;
   bool IsSupported(CalculationsPrecision precision) const;
   std::vector<TensorStorageType> GetSupportedTextureStorages() const;
-  std::vector<TensorStorageType> GetSupportedBatchStorages() const;
   std::vector<TensorStorageType> GetSupportedStorages() const;
 
   void SetHighPerformance() const;
diff --git a/tensorflow/lite/delegates/gpu/cl/kernels/util.cc b/tensorflow/lite/delegates/gpu/cl/kernels/util.cc
index a36da7050f0..dc7fe8f566b 100644
--- a/tensorflow/lite/delegates/gpu/cl/kernels/util.cc
+++ b/tensorflow/lite/delegates/gpu/cl/kernels/util.cc
@@ -234,9 +234,20 @@ std::string TensorCodeGenerator::GetGlobalAddressNoDeclaration(
   switch (descriptor_.storage_type) {
     case TensorStorageType::BUFFER:
     case TensorStorageType::IMAGE_BUFFER:
-      return absl::Substitute("(((($3) * $4 + $2) * $5 + ($1)) * $6 + ($0))", x,
-                              y, z, b, sizes_.depth, sizes_.height,
-                              sizes_.width);
+      return absl::Substitute("(((($3) * $4 + $2) * $5 + ($1)) * $6 + ($0))", b,
+                              x, y, z, sizes_.height, sizes_.width,
+                              sizes_.batch_size);
+    case TensorStorageType::TEXTURE_2D:
+      return absl::Substitute("(int2)(($0) * ($4) + ($1), ($2) * $5 + ($3))", x,
+                              b, y, z, sizes_.batch_size, sizes_.depth);
+    case TensorStorageType::SINGLE_TEXTURE_2D:
+      return absl::Substitute("(int2)(($0) * ($3) + ($1), ($2))", x, b, y,
+                              sizes_.batch_size);
+    case TensorStorageType::TEXTURE_ARRAY:
+      return absl::Substitute("(int4)(($0) * ($4) + ($1), ($2), ($3), 0)", x, b,
+                              y, z, sizes_.batch_size);
+    case TensorStorageType::UNKNOWN:
+      return "error";
     default:
       return "error";
   }
diff --git a/tensorflow/lite/delegates/gpu/cl/tensor.h b/tensorflow/lite/delegates/gpu/cl/tensor.h
index 050ad4fc392..136dda298c2 100644
--- a/tensorflow/lite/delegates/gpu/cl/tensor.h
+++ b/tensorflow/lite/delegates/gpu/cl/tensor.h
@@ -95,8 +95,6 @@ class Tensor {
     switch (descriptor_.storage_type) {
       case TensorStorageType::BUFFER:
       case TensorStorageType::IMAGE_BUFFER:
-        return (((b * Depth() + d) * shape_.h + y) * shape_.w + x) * 4 +
-               sub_d;  // BDHWC4
       case TensorStorageType::TEXTURE_ARRAY:
         return (((d * shape_.h + y) * shape_.w + x) * shape_.b + b) * 4 +
                sub_d;  // DHWBC4