From fc52e643284c56a3165eba1217f80434949f1c79 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 2 Dec 2020 16:11:57 -0800 Subject: [PATCH] Distinguish between not having a cross-program prefetch address and having a nullptr cross-program prefetch address. PiperOrigin-RevId: 345334049 Change-Id: I150066c6ae18e762231b55388159c479ef0fb84f --- tensorflow/core/tpu/tpu_ops_c_api.h | 1 + tensorflow/stream_executor/tpu/tpu_executable.cc | 2 ++ 2 files changed, 3 insertions(+) diff --git a/tensorflow/core/tpu/tpu_ops_c_api.h b/tensorflow/core/tpu/tpu_ops_c_api.h index f49438bde85..77e5ddb406c 100644 --- a/tensorflow/core/tpu/tpu_ops_c_api.h +++ b/tensorflow/core/tpu/tpu_ops_c_api.h @@ -122,6 +122,7 @@ typedef struct TpuExecutable_LoadProgramAndEnqueueToStream_Params { SE_DeviceMemoryBase* arguments; size_t arguments_len; SE_DeviceMemoryBase* result; + bool has_cross_program_prefetch_addr; SE_DeviceMemoryBase* cross_program_prefetch_addr; int32_t rng_seed; XLA_DeviceAssignment* device_assignment; diff --git a/tensorflow/stream_executor/tpu/tpu_executable.cc b/tensorflow/stream_executor/tpu/tpu_executable.cc index a251f6711c6..6408d37b990 100644 --- a/tensorflow/stream_executor/tpu/tpu_executable.cc +++ b/tensorflow/stream_executor/tpu/tpu_executable.cc @@ -84,6 +84,8 @@ Status TpuExecutable::LoadProgramAndEnqueueToStream( params.arguments = arguments_bases; params.arguments_len = arguments.size(); params.result = &result_base; + params.has_cross_program_prefetch_addr = + cross_program_prefetch_addr.has_value(); params.cross_program_prefetch_addr = cross_program_prefetch_addr.has_value() ? &prefetch_base : nullptr; params.rng_seed = rng_seed;