Home
last modified time | relevance | path

Searched refs:HloSortInstruction (Results 1 – 16 of 16) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dtopk_rewriter.h28 explicit TopkRewriter(std::function<bool(const HloSortInstruction*, int64)> in TopkRewriter() argument
39 std::function<bool(const HloSortInstruction*, int64)>
Dtopk_rewriter_test.cc80 TopkRewriter rewriter([](const HloSortInstruction*, int64) { return true; }); in TEST_F() argument
110 TopkRewriter rewriter([](const HloSortInstruction*, int64) { return true; }); in TEST_F() argument
140 TopkRewriter rewriter([](const HloSortInstruction*, int64) { return true; }); in TEST_F() argument
Dstable_sort_expander.cc39 auto* sort = Cast<HloSortInstruction>(instruction); in ExpandInstruction()
125 sort = Cast<HloSortInstruction>(new_sort); in ExpandInstruction()
202 Cast<HloSortInstruction>(instruction)->is_stable(); in InstructionMatchesPattern()
Dtopk_rewriter.cc70 HloSortInstruction* sort = DynCast<HloSortInstruction>(inst); in Run()
Dhlo_instructions.cc945 HloSortInstruction::HloSortInstruction( in HloSortInstruction() function in xla::HloSortInstruction
958 HloInstructionProto HloSortInstruction::ToProto() const { in ToProto()
967 std::vector<string> HloSortInstruction::ExtraAttributesToStringImpl( in ExtraAttributesToStringImpl()
977 bool HloSortInstruction::IdenticalSlowPath( in IdenticalSlowPath()
981 const auto& casted_other = static_cast<const HloSortInstruction&>(other); in IdenticalSlowPath()
991 std::unique_ptr<HloInstruction> HloSortInstruction::CloneWithNewOperandsImpl( in CloneWithNewOperandsImpl()
994 return absl::make_unique<HloSortInstruction>( in CloneWithNewOperandsImpl()
Ddynamic_dimension_inference.cc347 HloSortInstruction* sort = Cast<HloSortInstruction>(hlo); in HandleSort()
Ddynamic_padder.cc1196 HloSortInstruction* sort = Cast<HloSortInstruction>(hlo); in RewriteDynamicSort()
Dhlo_instructions.h627 class HloSortInstruction : public HloInstruction {
629 explicit HloSortInstruction(const Shape& shape, int64 dimension,
Dhlo_evaluator.cc2240 if (Cast<HloSortInstruction>(sort)->is_stable()) { in HandleSort()
Dhlo_instruction.cc1480 return absl::make_unique<HloSortInstruction>(shape, dimension, operands, in CreateSort()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dcpu_compiler.cc374 pipeline.AddPass<TopkRewriter>([](const HloSortInstruction* sort, int64) { in RunHloPassesThroughLayoutAssn()
Dir_emitter.cc546 const HloSortInstruction* sort = Cast<HloSortInstruction>(hlo); in HandleSort()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_util.cc1156 HloSortInstruction* sort = DynCast<HloSortInstruction>(hlo); in GetKValueInTopKWhenPartitionSortDim()
Dspmd_partitioner.cc1657 HloSortInstruction* sort = DynCast<HloSortInstruction>(hlo); in HandleSort()
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dhlo_function_importer.cc487 auto sort_instruction = Cast<HloSortInstruction>(instruction); in ImportInstructionImpl()
/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
Dmhlo_to_lhlo_with_xla.cc419 auto* sort_instr = xla::Cast<xla::HloSortInstruction>(instr); in EmitSortOp()