Introduce additional XLA TPU Ops to open source

PiperOrigin-RevId: 326343558
Change-Id: I47da1dc0c96cdf8223ccebef012e2a5088a857a4

tensorflow/core/tpu/kernels/BUILD

@@ -38,6 +38,7 @@ tf_kernel_library(
         ":tpu_execute_op",
         ":tpu_handle_to_key_op",
         ":transfer_ops",
+        "//tensorflow/core/tpu/kernels/xla:xla_ops",
     ],
 )
 

tensorflow/core/tpu/kernels/xla/BUILD

@@ -0,0 +1,52 @@
+# XLA Ops for TPUs
+
+package(
+    licenses = ["notice"],  # Apache 2.0
+)
+
+cc_library(
+    name = "xla_ops",
+    srcs = [
+        "get_item_op.cc",
+        "host_compute_ops.cc",
+        "index_ops.cc",
+        "infeed_op.cc",
+        "inplace_ops.cc",
+        "outfeed_ops.cc",
+        "segment_reduction_ops.cc",
+        "where_op.cc",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/compiler/tf2xla:common",
+        "//tensorflow/compiler/tf2xla:sharding_util",
+        "//tensorflow/compiler/tf2xla:side_effect_util",
+        "//tensorflow/compiler/tf2xla:xla_compiler",
+        "//tensorflow/compiler/tf2xla:xla_context",
+        "//tensorflow/compiler/tf2xla:xla_helpers",
+        "//tensorflow/compiler/tf2xla:xla_op_registry",
+        "//tensorflow/compiler/tf2xla/kernels:if_op",
+        "//tensorflow/compiler/tf2xla/kernels:while_op",
+        "//tensorflow/compiler/tf2xla/kernels:xla_ops",
+        "//tensorflow/compiler/tf2xla/lib:scatter",
+        "//tensorflow/compiler/xla:shape_util",
+        "//tensorflow/compiler/xla:util",
+        "//tensorflow/compiler/xla:xla_data_proto_cc",
+        "//tensorflow/compiler/xla/client:xla_builder",
+        "//tensorflow/compiler/xla/client/lib:arithmetic",
+        "//tensorflow/compiler/xla/client/lib:comparators",
+        "//tensorflow/compiler/xla/client/lib:constants",
+        "//tensorflow/core:core_cpu_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:graph",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/tpu:tpu_api",
+        "//tensorflow/core/tpu:tpu_defs",
+        "//tensorflow/core/tpu/kernels:cross_replica_ops",
+        "//tensorflow/stream_executor/tpu:c_api_conversions",
+        "//tensorflow/stream_executor/tpu:c_api_decl",
+        "@com_google_absl//absl/strings",
+    ],
+    alwayslink = 1,
+)

tensorflow/core/tpu/kernels/xla/get_item_op.cc

@@ -0,0 +1,75 @@
+/* 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.
+==============================================================================*/
+
+#define EIGEN_USE_THREADS
+
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/compiler/xla/client/xla_builder.h"
+#include "tensorflow/core/framework/kernel_def_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor_util.h"
+
+namespace tensorflow {
+namespace {
+
+// The Xla kernel to build up the computation for get_item(data, index).
+class GetItemXlaOp : public XlaOpKernel {
+ public:
+  explicit GetItemXlaOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape& data_shape = ctx->InputShape(0);
+    const TensorShape& index_shape = ctx->InputShape(1);
+    OP_REQUIRES(
+        ctx, TensorShapeUtils::IsVectorOrHigher(data_shape),
+        errors::InvalidArgument("data must be at least 1 dimensional."));
+    OP_REQUIRES(ctx, index_shape.dims() == 1 && index_shape.dim_size(0) == 1,
+                errors::InvalidArgument("index must be a vector of size 1."));
+
+    // NOTE(pbar) Use Concat to extend the indices to match cl/142279605.
+    // This isn't the simplest way to emit the indices, but the code for
+    // dynamic slice needs to be able to see that minor dims are const zero.
+    auto const_zero = xla::ConstantR0(ctx->builder(), 0);
+    std::vector<xla::XlaOp> operands;
+    operands.push_back(xla::Reshape(ctx->Input(1), {}));
+    for (int i = 1; i < data_shape.dims(); i++) {
+      operands.push_back(const_zero);
+    }
+
+    std::vector<int64> dims = {0};
+    std::vector<int64> slice_sizes = {1};
+    std::vector<int64> out_sizes = {};
+    for (int i = 1; i < data_shape.dims(); i++) {
+      dims.push_back(i);
+      auto size = data_shape.dim_size(i);
+      slice_sizes.push_back(size);
+      out_sizes.push_back(size);
+    }
+    // NOTE: DynamicSlice here doesn't raise an error or wraps the index
+    // if its out-of-range.
+    auto slice = xla::DynamicSlice(ctx->Input(0), operands, slice_sizes);
+    // In-order collapse to remove the 1st dim.
+    auto reshape = xla::Reshape(slice, dims, out_sizes);
+    ctx->SetOutput(0, reshape);
+  }
+};
+
+REGISTER_XLA_OP(Name("GetItem"), GetItemXlaOp);
+
+}  // namespace
+}  // namespace tensorflow