experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Apply new TraceMe APIs.

Browse Source 
PiperOrigin-RevId: 317169381
Change-Id: I2259895a8dde21e25661a239b9d4f5911a454adb
This commit is contained in:
Jiho Choi 2020-06-18 13:30:08 -07:00 committed by TensorFlower Gardener
parent 2663edb669
commit 6558da5a66
3 changed files with 11 additions and 8 deletions
tensorflow
compiler/xla/pjrt
BUILDpjrt_client.cc
core/profiler/lib
connected_traceme.h

2
tensorflow/compiler/xla/pjrt/BUILD
View File

@ -141,7 +141,9 @@ cc_library(
"//tensorflow/compiler/xla/service/gpu:gpu_executable_run_options", "//tensorflow/compiler/xla/service/gpu:gpu_executable_run_options",
"//tensorflow/core:allocator", "//tensorflow/core:allocator",
"//tensorflow/core:lib", "//tensorflow/core:lib",
"//tensorflow/core/profiler/lib:connected_traceme",
"//tensorflow/core/profiler/lib:traceme", "//tensorflow/core/profiler/lib:traceme",
"//tensorflow/core/profiler/lib:traceme_encode",
"//tensorflow/stream_executor:event", "//tensorflow/stream_executor:event",
"//tensorflow/stream_executor:stream", "//tensorflow/stream_executor:stream",
"//tensorflow/stream_executor/host:host_platform_id", "//tensorflow/stream_executor/host:host_platform_id",

16
tensorflow/compiler/xla/pjrt/pjrt_client.cc
View File

@ -98,7 +98,9 @@ limitations under the License.
#include "tensorflow/core/platform/mem.h" #include "tensorflow/core/platform/mem.h"
#include "tensorflow/core/platform/status.h" #include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/types.h" #include "tensorflow/core/platform/types.h"
#include "tensorflow/core/profiler/lib/connected_traceme.h"
#include "tensorflow/core/profiler/lib/traceme.h" #include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/profiler/lib/traceme_encode.h"
#include "tensorflow/stream_executor/device_memory.h" #include "tensorflow/stream_executor/device_memory.h"
#include "tensorflow/stream_executor/device_memory_allocator.h" #include "tensorflow/stream_executor/device_memory_allocator.h"
#include "tensorflow/stream_executor/event.h" #include "tensorflow/stream_executor/event.h"
@ -1429,10 +1431,9 @@ StatusOr<ScopedShapedBuffer> PjRtExecutable::EnqueueExecution(
int executable_idx, const RunId& run_id, const ExecuteOptions& options, int executable_idx, const RunId& run_id, const ExecuteOptions& options,
Device* device, std::vector<PjRtBuffer::ScopedHold>* device_buffers) const { Device* device, std::vector<PjRtBuffer::ScopedHold>* device_buffers) const {
int device_ordinal = device->local_device_state()->device_ordinal(); int device_ordinal = device->local_device_state()->device_ordinal();
tensorflow::profiler::TraceMe traceme([&] { tensorflow::profiler::TraceMeConsumer activity(
return absl::StrCat("LocalExecutable::Execute#run_id=", run_id.ToInt(), "LocalExecutable::Execute", tensorflow::profiler::ContextType::kPjRt,
"#"); run_id.ToInt());
});
VLOG(3) << "Replica " << replica << ", partition " << partition VLOG(3) << "Replica " << replica << ", partition " << partition
<< " mapped to device ordinal for execution: " << device_ordinal; << " mapped to device ordinal for execution: " << device_ordinal;
@ -1721,10 +1722,9 @@ PjRtExecutable::ExecuteOnLocalDevices(
absl::Span<const std::vector<PjRtBuffer*>> argument_handles, absl::Span<const std::vector<PjRtBuffer*>> argument_handles,
const ExecuteOptions& options) const { const ExecuteOptions& options) const {
RunId run_id; RunId run_id;
tensorflow::profiler::TraceMe traceme([&] { tensorflow::profiler::TraceMeProducer activity(
return absl::StrCat( "LocalExecutable::ExecuteOnLocalDevices",
"LocalExecutable::ExecuteOnLocalDevices#run_id=", run_id.ToInt(), "#"); tensorflow::profiler::ContextType::kPjRt, run_id.ToInt());
});
const int num_local_devices = local_devices_.size(); const int num_local_devices = local_devices_.size();

1
tensorflow/core/profiler/lib/connected_traceme.h
View File

@ -29,6 +29,7 @@ enum class ContextType : int {
kGeneric, kGeneric,
kTfExecutor, kTfExecutor,
kSharedBatchScheduler, kSharedBatchScheduler,
kPjRt,
}; };
/* /*