Home
last modified time | relevance | path

Searched refs:unnested_hlo (Results 1 – 3 of 3) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dir_emitter_unnested.cc296 const HloInstruction* unnested_hlo = hlo; in GetIndexTypeForKernel() local
299 unnested_hlo = computation->FusionInstruction(); in GetIndexTypeForKernel()
321 if (!shape_in_range(unnested_hlo->shape())) { in GetIndexTypeForKernel()
330 if (!absl::c_all_of(unnested_hlo->operands(), hlo_shape_in_range)) { in GetIndexTypeForKernel()
335 if (unnested_hlo->opcode() == HloOpcode::kFusion) { in GetIndexTypeForKernel()
337 unnested_hlo->fused_instructions_computation()->instructions(), in GetIndexTypeForKernel()
707 const HloInstruction* unnested_hlo, const IrArray::Index& index, in EmitExtraOutputsForReduce() argument
712 GetIrArray(*unnested_hlo, *unnested_hlo, extra_output_gens[i].second) in EmitExtraOutputsForReduce()
2643 HloInstruction* unnested_hlo, HloInstruction* reduce_inst, int reduce_idx, in EmitPrologueForOneReduction() argument
2675 if (unnested_hlo->opcode() == HloOpcode::kFusion) { in EmitPrologueForOneReduction()
[all …]
Dir_emitter_unnested.h211 const HloInstruction* unnested_hlo, const llvm_ir::IrArray::Index& index,
218 Status EmitReductionToVector(HloInstruction* unnested_hlo);
226 ComputeMappingSchemeAndReductionKind(const HloInstruction* unnested_hlo,
247 LaunchDimensions EmitKernel(HloInstruction* unnested_hlo,
270 void EmitTileElementForReduction(HloInstruction* unnested_hlo,
276 void EmitPrologueForReduction(HloInstruction* unnested_hlo,
278 void EmitPrologueForOneReduction(HloInstruction* unnested_hlo,
284 void EmitEpilogueForReduction(HloInstruction* unnested_hlo,
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dir_emitter.h236 HloInstruction* unnested_hlo) { in GetGeneratorForOperandIrArrays() argument
237 return [=]() { return GetIrArraysForOperandsOf(unnested_hlo); }; in GetGeneratorForOperandIrArrays()