Open source global_iter_id.cc for SparseCore.

PiperOrigin-RevId: 565097655

tensorflow/core/tpu/kernels/BUILD

 # TPU Kernel Implementations
 
-load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
-load(
-    "//tensorflow/core/platform:build_config.bzl",
-    "tf_proto_library",
-)
 load(
     "//tensorflow:tensorflow.bzl",
     "if_libtpu",
     "tf_cc_test",
     "tf_copts",
+    "tf_gen_op_wrapper_py",
 )
 load(
     "//tensorflow:tensorflow.default.bzl",
     "tf_grpc_cc_dependencies",
     "tf_kernel_library",
 )
+load(
+    "//tensorflow/core/platform:build_config.bzl",
+    "tf_proto_library",
+)
 load(
     "//tensorflow/core/platform:build_config_root.bzl",
     "if_static",
 )
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 load(
     "//tensorflow/core/tpu:build_defs.bzl",
     "if_libtpu_tf_status",
@@ -1270,6 +1271,24 @@ cc_library(
     alwayslink = 1,
 )
 
+tf_kernel_library(
+    name = "global_iter_id_op",
+    srcs = ["global_iter_id.cc"],
+    deps = [
+        "//tensorflow/core:framework",
+        "//tensorflow/core/kernels:partitioned_function_ops",
+    ],
+)
+
+tf_gen_op_wrapper_py(
+    name = "gen_global_iter_id_op",
+    out = "gen_global_iter_id_op.py",
+    op_allowlist = [
+        "GlobalIterId",
+    ],
+    deps = [":global_iter_id_op"],
+)
+
 tf_cc_test(
     name = "sharding_util_ops_test",
     srcs = ["sharding_util_ops_test.cc"],

tensorflow/core/tpu/kernels/global_iter_id.cc

+/* Copyright 2023 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/core/framework/function.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.h"
+
+namespace tensorflow {
+
+namespace {
+
+class GlobalIterId : public OpKernel {
+ public:
+  explicit GlobalIterId(OpKernelConstruction* ctx) : OpKernel(ctx) {}
+  void Compute(OpKernelContext* ctx) override {
+    ctx->set_output(0, Tensor(ctx->frame_iter().iter_id));
+  }
+};
+REGISTER_OP("GlobalIterId").Output("iter_id: int64").SetIsStateful();
+REGISTER_KERNEL_BUILDER(Name("GlobalIterId").Device(DEVICE_CPU), GlobalIterId);
+
+}  // anonymous namespace
+}  // namespace tensorflow