Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dir_emitter_unnested.h445 mlir::Operation* unnested_hlo, mlir::Operation* first_reduce);
587 mlir::Operation* unnested_hlo, const Shape& reduction_operand_shape,
601 mlir::Operation* unnested_hlo, absl::Span<const int> instr_index_group,
610 llvm::Type* index_ty, mlir::Operation* unnested_hlo,
618 void EmitIRForReduction(mlir::Operation* unnested_hlo,
Dir_emitter_unnested.cc383 const HloInstruction* unnested_hlo = hlo; in GetIndexTypeForKernel() local
386 unnested_hlo = computation->FusionInstruction(); in GetIndexTypeForKernel()
408 if (!shape_in_range(unnested_hlo->shape())) { in GetIndexTypeForKernel()
417 if (!absl::c_all_of(unnested_hlo->operands(), hlo_shape_in_range)) { in GetIndexTypeForKernel()
422 if (unnested_hlo->opcode() == HloOpcode::kFusion) { in GetIndexTypeForKernel()
424 unnested_hlo->fused_instructions_computation()->instructions(), in GetIndexTypeForKernel()
4200 static mlir::Operation* GetReduceFromUnnestedMlir(mlir::Operation* unnested_hlo, in GetReduceFromUnnestedMlir() argument
4202 if (mlir::isa<mlir::lmhlo::ReduceOp>(unnested_hlo)) { in GetReduceFromUnnestedMlir()
4204 return unnested_hlo; in GetReduceFromUnnestedMlir()
4206 if (auto fusion = mlir::dyn_cast<mlir::lmhlo::FusionOp>(unnested_hlo)) { in GetReduceFromUnnestedMlir()
[all …]
Dir_emission_utils.h171 bool IsInputFusibleSlices(mlir::Operation* unnested_hlo,
Dir_emission_utils.cc372 bool IsInputFusibleSlices(mlir::Operation* unnested_hlo, in IsInputFusibleSlices() argument
374 auto fusion = mlir::dyn_cast<mlir::lmhlo::FusionOp>(unnested_hlo); in IsInputFusibleSlices()