Searched refs:stream_assignment (Results 1 – 7 of 7) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | stream_assignment.cc | 66 const HloInstruction& hlo, const StreamAssignment& stream_assignment, in ComputeStreamToAssign() argument
87 if (stream_assignment.HasStreamAssigned(*operand)) { in ComputeStreamToAssign()
89 stream_assignment.StreamNumberForHlo(*operand)); in ComputeStreamToAssign()
104 int stream_num = stream_assignment.StreamNumberForHlo(*seen_gemm); in ComputeStreamToAssign()
111 for (int stream_num = 0; stream_num < stream_assignment.StreamCount(); in ComputeStreamToAssign()
117 return stream_assignment.StreamCount(); in ComputeStreamToAssign()
123 auto stream_assignment = absl::make_unique<StreamAssignment>(); in AssignStreams() local
141 : ComputeStreamToAssign(*hlo, *stream_assignment, in AssignStreams()
144 stream_assignment->AssignStreamToHlo(hlo, stream_num); in AssignStreams()
154 return stream_assignment; in AssignStreams()
|
| D | gpu_hlo_schedule.cc | 39 const StreamAssignment& stream_assignment,
58 const HloModule* module, const StreamAssignment& stream_assignment, in GpuHloOrdering() argument
62 if (stream_assignment.StreamCount() == 1) { in GpuHloOrdering()
84 stream_assignment.StreamCount(), nullptr); in GpuHloOrdering()
88 if (stream_assignment.HasStreamAssigned(*hlo)) { in GpuHloOrdering()
100 const int stream_no = stream_assignment.StreamNumberForHlo(*hlo); in GpuHloOrdering()
193 const HloModule& module, const StreamAssignment& stream_assignment, in Build() argument
199 if (stream_assignment.StreamCount() == 1) { in Build()
215 &module, stream_assignment, schedule->thunk_launch_order_); in Build()
|
| D | thunk_schedule.cc | 52 std::unique_ptr<StreamAssignment> stream_assignment, in ThunkSchedule() argument
55 stream_assignment_(std::move(stream_assignment)) { in ThunkSchedule()
|
| D | gpu_hlo_schedule.h | 44 const HloModule& module, const StreamAssignment& stream_assignment,
|
| D | thunk_schedule.h | 50 std::unique_ptr<StreamAssignment> stream_assignment,
|
| D | BUILD | 88 name = "stream_assignment",
89 srcs = ["stream_assignment.cc"],
90 hdrs = ["stream_assignment.h"],
108 ":stream_assignment",
378 ":stream_assignment",
829 ":stream_assignment",
981 ":stream_assignment",
1000 ":stream_assignment",
|
| D | nvptx_compiler.cc | 659 std::unique_ptr<StreamAssignment> stream_assignment = AssignStreams(*module); in RunBackend() local
662 GpuHloSchedule::Build(*module, *stream_assignment, pointer_size_)); in RunBackend()
767 ir_emitter.ConsumeThunkSequence(), std::move(stream_assignment), in RunBackend()
|