commit 12b359caed1f19e73a82b02e813bcc02951564c3
author Raman Sarokin <sorokin@google.com> Fri Mar 04 02:37:36 2022 -0800
committer TensorFlower Gardener <gardener@tensorflow.org> Fri Mar 04 02:43:07 2022 -0800
tree b180bf94bf39e2e06863887cce3da4001c7dd7ee
parent 1cd7e608c986ba0757400372f022843603e32b88

Added support of reading short4/ushort4/char4/uchar4 from buffers in Glsl.

PiperOrigin-RevId: 432399398

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

diff --git a/tensorflow/lite/delegates/gpu/common/task/tensor_desc.cc b/tensorflow/lite/delegates/gpu/common/task/tensor_desc.cc
index 639db2a..9fc01ab 100644
--- a/tensorflow/lite/delegates/gpu/common/task/tensor_desc.cc
+++ b/tensorflow/lite/delegates/gpu/common/task/tensor_desc.cc
@@ -493,6 +493,22 @@
           return absl::StrCat("vec4(unpackHalf2x16(buffer[", coords[0],
                               "].x), unpackHalf2x16(buffer[", coords[0],
                               "].y))");
+        } else if (data_type == DataType::INT16 ||
+                   data_type == DataType::UINT16) {
+          const std::string vec_type =
+              data_type == DataType::INT16 ? "ivec4" : "uvec4";
+          return absl::Substitute(
+              "$1(buffer[$0].x & 0xffff, (buffer[$0].x >> 16) & 0xffff, "
+              "buffer[$0].y & 0xffff, (buffer[$0].y >> 16) & 0xffff)",
+              coords[0], vec_type);
+        } else if (data_type == DataType::INT8 ||
+                   data_type == DataType::UINT8) {
+          const std::string vec_type =
+              data_type == DataType::INT8 ? "ivec4" : "uvec4";
+          return absl::Substitute(
+              "$1(buffer[$0] & 0xff, (buffer[$0] >> 8) & 0xff, "
+              "(buffer[$0] >> 16) & 0xff, (buffer[$0] >> 24) & 0xff)",
+              coords[0], vec_type);
         } else {
           return absl::StrCat("buffer[", coords[0], "]");
         }