Searched refs:NcclAllGatherThunk (Results 1 – 4 of 4) sorted by relevance
40 /*static*/ bool NcclAllGatherThunk::CanImplement(const HloInstruction* hlo) { in CanImplement()51 /*static*/ bool NcclAllGatherThunk::CanImplement(mlir::lmhlo::AllGatherOp op) { in CanImplement()61 NcclAllGatherThunk::NcclAllGatherThunk( in NcclAllGatherThunk() function in xla::gpu::NcclAllGatherThunk63 std::vector<NcclAllGatherThunk::Buffer> buffers) in NcclAllGatherThunk()70 Status NcclAllGatherThunk::RunNcclCollective(const ExecuteParams& params, in RunNcclCollective()
35 class NcclAllGatherThunk : public NcclCollectiveThunk {37 NcclAllGatherThunk(ThunkInfo thunk_info, mlir::lmhlo::AllGatherOp op,
159 return !NcclAllGatherThunk::CanImplement(&ag); in OptimizeHloModule()
3174 return EmitNcclThunkFromMlir<NcclAllGatherThunk, mlir::lmhlo::AllGatherOp>( in HandleAllGather()