[XLA:TPU] HostOffloader: Refactor HandleInputStreaming to occur befor…

…e tracing the graph for custom calls.

Currently, input streaming gets handled in arbitrary order, after we have already possibly searched the graph for the corresponding streamed custom calls. This refactors the input streaming to occur before any graph searching, which makes more sense and prepares for output streaming.

PiperOrigin-RevId: 629208199

tensorflow/core/data/service/client/data_service_client.cc

@@ -42,6 +42,7 @@ limitations under the License.
 #include "tensorflow/core/data/utils.h"
 #include "tensorflow/core/framework/allocator.h"
 #include "tensorflow/core/framework/dataset.h"
+#include "tensorflow/core/framework/device_base.h"
 #include "tensorflow/core/framework/metrics.h"
 #include "tensorflow/core/framework/model.h"
 #include "tensorflow/core/platform/env.h"
@@ -102,7 +103,10 @@ DataServiceClient::~DataServiceClient() {
           << iteration_client_id_;
 }
 
-Status DataServiceClient::Initialize(Allocator* allocator) {
+Status DataServiceClient::Initialize(
+    const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info,
+    Allocator* allocator) {
+  accelerator_device_info_ = accelerator_device_info;
   allocator_ = allocator;
   TF_RETURN_IF_ERROR(ValidateDataServiceParams(params_));
   VLOG(3) << "Connecting to " << params_.address
@@ -343,7 +347,7 @@ DataServiceClient::CreateWorkerClient(const std::string& protocol,
   TF_ASSIGN_OR_RETURN(DataTransferServerInfo transfer_server,
                       GetTransferServer(protocol, task_info));
   return CreateDataServiceWorkerClient(params_.protocol, transfer_server,
-                                       allocator_);
+                                       accelerator_device_info_, allocator_);
 }
 
 absl::StatusOr<std::unique_ptr<DataServiceWorkerClient>>
@@ -356,7 +360,7 @@ DataServiceClient::CreateAlternativeWorkerClientWithGrpcFallback(
     const DataTransferServerInfo& transfer_server, const TaskInfo& task_info) {
   absl::StatusOr<std::unique_ptr<DataServiceWorkerClient>> worker =
       CreateDataServiceWorkerClient(params_.protocol, transfer_server,
-                                    allocator_);
+                                    accelerator_device_info_, allocator_);
   if (worker.ok()) {
     LOG(INFO) << "Successfully started client for data transfer protocol '"
               << transfer_server.protocol() << "' for worker '"
@@ -383,7 +387,8 @@ DataServiceClient::CreateWorkerClient(const TaskInfo& task_info) {
     DataTransferServerInfo info;
     info.set_protocol(kLocalTransferProtocol);
     info.set_address(task_info.worker_address());
-    return CreateDataServiceWorkerClient(params_.protocol, info, allocator_);
+    return CreateDataServiceWorkerClient(params_.protocol, info,
+                                         accelerator_device_info_, allocator_);
   }
   if (!params_.data_transfer_protocol.empty()) {
     TF_ASSIGN_OR_RETURN(

tensorflow/core/data/service/client/data_service_client.h

@@ -80,7 +80,9 @@ class DataServiceClient {
   DataServiceClient& operator=(const DataServiceClient&) = delete;
 
   // Initializes the client.
-  Status Initialize(Allocator* allocator);
+  Status Initialize(
+      const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info,
+      Allocator* allocator);
 
   // Reads the next element from tf.data workers. Blocks if the next element is
   // not ready.
@@ -246,6 +248,7 @@ class DataServiceClient {
   int64_t job_id_;
   int64_t iteration_client_id_;
   std::unique_ptr<DataServiceDispatcherClient> dispatcher_;
+  const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info_;
   Allocator* allocator_;
 
   int64_t get_next_index_ TF_GUARDED_BY(mu_) = 0;

tensorflow/core/data/service/client/data_service_client_test.cc

@@ -134,7 +134,8 @@ TEST(DataServiceClientTest, NoSharding) {
   DataServiceParams params = GetDataServiceParams(
       dataset_id, test_cluster.DispatcherAddress(), ProcessingModeDef::OFF);
   DataServiceClient client(params);
-  TF_ASSERT_OK(client.Initialize(/*allocator=*/nullptr));
+  TF_ASSERT_OK(client.Initialize(/*accelerator_device_info=*/nullptr,
+                                 /*allocator=*/nullptr));
   EXPECT_THAT(GetResults<int64_t>(client),
               IsOkAndHolds(ElementsAreArray(Range(10))));
   client.Cancel();
@@ -150,7 +151,8 @@ TEST(DataServiceClientTest, DynamicSharding) {
   DataServiceParams params = GetDataServiceParams(
       dataset_id, test_cluster.DispatcherAddress(), ProcessingModeDef::DYNAMIC);
   DataServiceClient client(params);
-  TF_ASSERT_OK(client.Initialize(/*allocator=*/nullptr));
+  TF_ASSERT_OK(client.Initialize(/*accelerator_device_info=*/nullptr,
+                                 /*allocator=*/nullptr));
   EXPECT_THAT(GetResults<int64_t>(client),
               IsOkAndHolds(UnorderedElementsAreArray(Range(10))));
   client.Cancel();
@@ -167,7 +169,8 @@ TEST(DataServiceClientTest, StaticSharding) {
       GetDataServiceParams(dataset_id, test_cluster.DispatcherAddress(),
                            ProcessingModeDef::FILE_OR_DATA);
   DataServiceClient client(params);
-  TF_ASSERT_OK(client.Initialize(/*allocator=*/nullptr));
+  TF_ASSERT_OK(client.Initialize(/*accelerator_device_info=*/nullptr,
+                                 /*allocator=*/nullptr));
   EXPECT_THAT(GetResults<int64_t>(client),
               IsOkAndHolds(UnorderedElementsAreArray(Range(10))));
   client.Cancel();
@@ -183,7 +186,8 @@ TEST(DataServiceClientTest, RecordBufferEvents) {
   DataServiceParams params = GetDataServiceParams(
       dataset_id, test_cluster.DispatcherAddress(), ProcessingModeDef::OFF);
   DataServiceClient client(params);
-  TF_ASSERT_OK(client.Initialize(/*allocator=*/nullptr));
+  TF_ASSERT_OK(client.Initialize(/*accelerator_device_info=*/nullptr,
+                                 /*allocator=*/nullptr));
 
   auto mock_context = std::make_unique<TestDataServiceContext>();
   TestDataServiceContext* ctx = mock_context.get();
@@ -206,7 +210,8 @@ TEST(DataServiceClientTest, Cancel) {
   DataServiceParams params = GetDataServiceParams(
       dataset_id, test_cluster.DispatcherAddress(), ProcessingModeDef::OFF);
   DataServiceClient client(params);
-  TF_ASSERT_OK(client.Initialize(/*allocator=*/nullptr));
+  TF_ASSERT_OK(client.Initialize(/*accelerator_device_info=*/nullptr,
+                                 /*allocator=*/nullptr));
   client.Cancel();
   EXPECT_THAT(client.GetNext(GetTestDataServiceContext),
               StatusIs(error::CANCELLED));
@@ -218,7 +223,8 @@ TEST(DataServiceClientTest, ValidationError) {
   params.target_workers = TARGET_WORKERS_LOCAL;
   DataServiceClient client(params);
   EXPECT_THAT(
-      client.Initialize(/*allocator=*/nullptr),
+      client.Initialize(/*accelerator_device_info=*/nullptr,
+                        /*allocator=*/nullptr),
       StatusIs(
           error::INVALID_ARGUMENT,
           HasSubstr(

tensorflow/core/data/service/data_transfer.h

@@ -70,6 +70,7 @@ class DataTransferClient {
   struct Config {
     absl::string_view protocol;
     std::string address;
+    const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info;
     Allocator* allocator;
   };
   using ClientFactoryT =

tensorflow/core/data/service/test_cluster.h

@@ -168,7 +168,8 @@ DatasetClient<T>::DatasetClient(const TestCluster& cluster)
   for (size_t i = 0; i < cluster.NumWorkers(); ++i) {
     worker_clients_[cluster_.WorkerAddress(i)] =
         std::make_unique<DataServiceWorkerClient>(
-            cluster_.WorkerAddress(i), "grpc", "grpc", /*allocator=*/nullptr);
+            cluster_.WorkerAddress(i), "grpc", "grpc",
+            /*accelerator_device_info=*/nullptr, /*allocator=*/nullptr);
   }
 }
 

tensorflow/core/data/service/worker_client.cc

@@ -38,6 +38,7 @@ limitations under the License.
 #include "tensorflow/core/data/service/worker_impl.h"
 #include "tensorflow/core/framework/allocator.h"
 #include "tensorflow/core/framework/dataset.pb.h"
+#include "tensorflow/core/framework/device_base.h"
 #include "tensorflow/core/framework/metrics.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
@@ -56,11 +57,13 @@ namespace tensorflow {
 namespace data {
 
 StatusOr<std::unique_ptr<DataServiceWorkerClient>>
-CreateDataServiceWorkerClient(const std::string& dispatcher_protocol,
-                              const DataTransferServerInfo& info,
-                              Allocator* allocator) {
+CreateDataServiceWorkerClient(
+    const std::string& dispatcher_protocol, const DataTransferServerInfo& info,
+    const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info,
+    Allocator* allocator) {
   auto client = std::make_unique<DataServiceWorkerClient>(
-      info.address(), dispatcher_protocol, info.protocol(), allocator);
+      info.address(), dispatcher_protocol, info.protocol(),
+      accelerator_device_info, allocator);
   TF_RETURN_IF_ERROR(client->Initialize());
   TF_RETURN_WITH_CONTEXT_IF_ERROR(
       client->CheckCompatibility(info.compatibility_info()),
@@ -82,7 +85,8 @@ Status DataServiceWorkerClient::EnsureInitialized() {
     return absl::OkStatus();
   }
   TF_RETURN_IF_ERROR(DataTransferClient::Build(
-      GetDataTransferProtocol(), {protocol_, address_, allocator_}, &client_));
+      GetDataTransferProtocol(),
+      {protocol_, address_, accelerator_device_info_, allocator_}, &client_));
   return absl::OkStatus();
 }
 

tensorflow/core/data/service/worker_client.h

@@ -37,12 +37,14 @@ constexpr const char kGrpcTransferProtocol[] = "grpc";
 // Client for communicating with the tf.data service worker.
 class DataServiceWorkerClient : public DataServiceClientBase {
  public:
-  DataServiceWorkerClient(const std::string& address,
-                          const std::string& protocol,
-                          const std::string& transfer_protocol,
-                          Allocator* allocator)
+  DataServiceWorkerClient(
+      const std::string& address, const std::string& protocol,
+      const std::string& transfer_protocol,
+      const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info,
+      Allocator* allocator)
       : DataServiceClientBase(address, protocol),
         transfer_protocol_(transfer_protocol),
+        accelerator_device_info_(accelerator_device_info),
         allocator_(allocator) {}
 
   // Fetches an element from the worker.
@@ -66,6 +68,7 @@ class DataServiceWorkerClient : public DataServiceClientBase {
 
  private:
   std::string transfer_protocol_;
+  const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info_;
   Allocator* allocator_;
 
   mutex mu_;
@@ -77,9 +80,10 @@ class DataServiceWorkerClient : public DataServiceClientBase {
 // Creates and initializes a new tf.data service worker client to read
 // from the data transfer server specified in `info`.
 StatusOr<std::unique_ptr<DataServiceWorkerClient>>
-CreateDataServiceWorkerClient(const std::string& dispatcher_protocol,
-                              const DataTransferServerInfo& info,
-                              Allocator* allocator);
+CreateDataServiceWorkerClient(
+    const std::string& dispatcher_protocol, const DataTransferServerInfo& info,
+    const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info,
+    Allocator* allocator);
 
 }  // namespace data
 }  // namespace tensorflow

tensorflow/core/data/service/worker_client_test.cc

@@ -107,6 +107,7 @@ class WorkerClientTest : public ::testing::Test {
     info.set_address(GetWorkerAddress());
     info.set_protocol(data_transfer_protocol);
     return CreateDataServiceWorkerClient(kProtocol, info,
+                                         /*accelerator_device_info=*/nullptr,
                                          /*allocator=*/nullptr);
   }
 

tensorflow/core/framework/dataset.h

@@ -667,7 +667,8 @@ class IteratorContext {
  public:
   struct Params {
     explicit Params(IteratorContext* ctx)
-        : allocator_getter(ctx->allocator_getter()),
+        : accelerator_device_info(ctx->accelerator_device_info()),
+          allocator_getter(ctx->allocator_getter()),
           cancellation_manager(ctx->cancellation_manager()),
           collective_executor(ctx->collective_executor()),
           env(ctx->env()),
@@ -697,6 +698,7 @@ class IteratorContext {
       // NOTE: need reinterpret_cast because function.h forward-declares Device.
       DeviceBase* device =
           reinterpret_cast<DeviceBase*>(ctx->function_library()->device());
+      accelerator_device_info = device->tensorflow_accelerator_device_info();
       allocator_getter = [device](AllocatorAttributes attrs) {
         return device->GetAllocator(attrs);
       };
@@ -719,6 +721,9 @@ class IteratorContext {
           *ctx->runner(), std::placeholders::_1);
     }
 
+    // If non-null, information about the GPU or TPU on which the op is placed.
+    const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info = nullptr;
+
     // The Allocator to be used to allocate the output of an iterator.
     std::function<Allocator*(AllocatorAttributes)> allocator_getter = nullptr;
 
@@ -825,6 +830,10 @@ class IteratorContext {
     return params_.id_registry;
   }
 
+  const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info() {
+    return params_.accelerator_device_info;
+  }
+
   Allocator* allocator(AllocatorAttributes attrs) {
     return params_.allocator_getter(attrs);
   }

tensorflow/core/kernels/data/experimental/data_service_dataset_op.cc

@@ -340,7 +340,8 @@ class DataServiceDatasetOp::Dataset : public DatasetBase {
       TF_RETURN_IF_ERROR(RegisterCancellationCallback(
           ctx->cancellation_manager(),
           [this]() { data_service_client_.Cancel(); }, &deregister_fn_));
-      return data_service_client_.Initialize(ctx->allocator(/*attrs=*/{}));
+      return data_service_client_.Initialize(ctx->accelerator_device_info(),
+                                             ctx->allocator(/*attrs=*/{}));
     }
 
     Status GetNextInternal(IteratorContext* ctx,

tensorflow/core/profiler/protobuf/task.proto

@@ -9,7 +9,7 @@ option cc_enable_arenas = true;
 // 'Task' contains information about a task that profiler traced.
 message Task {
   // The most recent changelist number from the client that built the binary.
-  optional int32 changelist = 1;
+  optional int64 changelist = 1;
   // True if the client that built the binary was mint (no local changes).
   optional bool clean_build = 2;
   // Build time (in ns relative to the Unix epoch).

third_party/stablehlo/temporary.patch

@@ -506,6 +506,17 @@ diff --ruN a/stablehlo/stablehlo/dialect/TypeInference.cpp b/stablehlo/stablehlo
    return success();
  }
 
+@@ -3861,8 +3819,8 @@
+   if (SmallVector<int64_t> shape; operandType.hasStaticShape() &&
+                                   matchInts(outputShape, shape).succeeded()) {
+     int64_t operandCount = operandType.getNumElements();
+-    int64_t shapeCount = std::accumulate(shape.begin(), shape.end(), 1,
+-                                         std::multiplies<int64_t>());
++    int64_t shapeCount = std::accumulate(
++        shape.begin(), shape.end(), int64_t{1}, std::multiplies<int64_t>());
+     if (operandCount != shapeCount) {
+       return emitOptionalError(location,
+                                "output_shape is incompatible with input type "
 diff --ruN a/stablehlo/stablehlo/experimental/BUILD.bazel b/stablehlo/stablehlo/experimental/BUILD.bazel
 --- stablehlo/stablehlo/experimental/BUILD.bazel
 +++ stablehlo/stablehlo/experimental/BUILD.bazel

third_party/xla/third_party/stablehlo/temporary.patch

@@ -506,6 +506,17 @@ diff --ruN a/stablehlo/stablehlo/dialect/TypeInference.cpp b/stablehlo/stablehlo
    return success();
  }
 
+@@ -3861,8 +3819,8 @@
+   if (SmallVector<int64_t> shape; operandType.hasStaticShape() &&
+                                   matchInts(outputShape, shape).succeeded()) {
+     int64_t operandCount = operandType.getNumElements();
+-    int64_t shapeCount = std::accumulate(shape.begin(), shape.end(), 1,
+-                                         std::multiplies<int64_t>());
++    int64_t shapeCount = std::accumulate(
++        shape.begin(), shape.end(), int64_t{1}, std::multiplies<int64_t>());
+     if (operandCount != shapeCount) {
+       return emitOptionalError(location,
+                                "output_shape is incompatible with input type "
 diff --ruN a/stablehlo/stablehlo/experimental/BUILD.bazel b/stablehlo/stablehlo/experimental/BUILD.bazel
 --- stablehlo/stablehlo/experimental/BUILD.bazel
 +++ stablehlo/stablehlo/experimental/BUILD.bazel

third_party/xla/xla/service/host_offloader.cc

@@ -755,12 +755,12 @@ absl::StatusOr<bool> HostOffloader::Run(
   // Run HloAliasAnalysis on module.
   TF_ASSIGN_OR_RETURN(alias_analysis_, HloAliasAnalysis::Run(module));
 
+  // Handle streamed parameters first.
+  TF_RETURN_IF_ERROR(HandleInputStreaming(module->entry_computation()));
+
   // Iterate over all instructions and look for XLA host offload annotations.
   for (HloComputation* computation :
        module->MakeNonfusionComputations(execution_threads)) {
-    if (computation->IsEntryComputation()) {
-      TF_RETURN_IF_ERROR(HandleInputStreaming(computation));
-    }
     for (HloInstruction* instruction :
          computation->MakeInstructionPostOrder()) {
       if (instruction->opcode() != HloOpcode::kCustomCall) {