Home
last modified time | relevance | path

Searched refs:stream_assignment (Results 1 – 8 of 8) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dstream_assignment.cc67 const HloInstruction& hlo, const StreamAssignment& stream_assignment, in ComputeStreamToAssign() argument
92 if (stream_assignment.HasStreamAssigned(*operand)) { in ComputeStreamToAssign()
94 stream_assignment.StreamNumberForHlo(*operand)); in ComputeStreamToAssign()
109 int stream_num = stream_assignment.StreamNumberForHlo(*seen_gemm); in ComputeStreamToAssign()
116 for (int stream_num = 0; stream_num < stream_assignment.StreamCount(); in ComputeStreamToAssign()
122 return stream_assignment.StreamCount(); in ComputeStreamToAssign()
128 auto stream_assignment = absl::make_unique<StreamAssignment>(); in AssignStreams() local
146 : ComputeStreamToAssign(*hlo, *stream_assignment, in AssignStreams()
149 stream_assignment->AssignStreamToHlo(hlo, stream_num); in AssignStreams()
159 return stream_assignment; in AssignStreams()
Dgpu_hlo_schedule.cc39 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()
Dthunk_schedule.cc52 std::unique_ptr<StreamAssignment> stream_assignment, in ThunkSchedule() argument
55 stream_assignment_(std::move(stream_assignment)), in ThunkSchedule()
Dgpu_hlo_schedule.h44 const HloModule& module, const StreamAssignment& stream_assignment,
Dthunk_schedule.h55 std::unique_ptr<StreamAssignment> stream_assignment,
Dgpu_compiler.cc521 std::unique_ptr<StreamAssignment> stream_assignment = in RunHloPassesAndBufferAssignement() local
525 GpuHloSchedule::Build(*hlo_module, *stream_assignment, pointer_size_)); in RunHloPassesAndBufferAssignement()
564 std::unique_ptr<StreamAssignment> stream_assignment = in CompileModuleToLlvmIrImpl() local
568 GpuHloSchedule::Build(*hlo_module, *stream_assignment, pointer_size)); in CompileModuleToLlvmIrImpl()
655 std::move(stream_assignment), std::move(thunk_to_hlo)); in CompileModuleToLlvmIrImpl()
DBUILD138 name = "stream_assignment",
139 srcs = ["stream_assignment.cc"],
140 hdrs = ["stream_assignment.h"],
160 ":stream_assignment",
616 ":stream_assignment",
1264 ":stream_assignment",
1550 ":stream_assignment",
1570 ":stream_assignment",
/external/tensorflow/tensorflow/compiler/mlir/g3doc/
Dxla_gpu_codegen.md257 …om/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/service/gpu/stream_assignment.cc#L140)