diff options
| author |  Adrian Kuegel <akuegel@google.com> | 2018-09-19 01:55:55 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-09-19 01:59:05 -0700 |
| commit | 0d8942fcbcc9cb3a05be8acc843d1fc4b6dfc9f1 (patch) | |
| tree | 28cccc167c27fa13e207860a4f339d9c28cd3e74 /tensorflow/compiler/xla/service/cpu/ir_emitter.h | |
| parent | 6967287715a097c8b009b52010c53247ab658232 (diff) | |
Implement sort op for CPU.
Also don't allow parallelization for the sort op in parallel_task_assignment.
PiperOrigin-RevId: 213592046
Diffstat (limited to 'tensorflow/compiler/xla/service/cpu/ir_emitter.h')
| -rw-r--r-- | tensorflow/compiler/xla/service/cpu/ir_emitter.h | 6 |
1 files changed, 6 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.h b/tensorflow/compiler/xla/service/cpu/ir_emitter.h index 3df99464ba..daafef4eb3 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emitter.h +++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.h @@ -163,6 +163,12 @@ class IrEmitter : public DfsHloVisitorWithDefault, Status Preprocess(HloInstruction* hlo) override; Status Postprocess(HloInstruction* hlo) override; + // A convenient helper for calling BufferAssignment::GetUniqueSlice. + BufferAllocation::Slice GetAllocationSlice( + const HloInstruction& hlo, const ShapeIndex& index = {}) const { + return assignment_.GetUniqueSlice(&hlo, index).ConsumeValueOrDie(); + } + private: // Private helper to initialize an IR function for the computation. void InitializeIrFunction(const string& function_name); |