Searched refs:body_b (Results 1 – 2 of 2) sorted by relevance
239 SpmdBuilder body_b("fft_collective_permute_body", hlo); in GetFinalFftUsingCollectivePermute() local240 auto param = body_b.AddInstruction(HloInstruction::CreateParameter( in GetFinalFftUsingCollectivePermute()246 auto dest_transform = body_b.AddInstruction( in GetFinalFftUsingCollectivePermute()248 auto source_transform = body_b.AddInstruction( in GetFinalFftUsingCollectivePermute()251 body_b.AddInstruction(HloInstruction::CreateGetTupleElement( in GetFinalFftUsingCollectivePermute()254 body_b.AddInstruction(HloInstruction::CreateGetTupleElement( in GetFinalFftUsingCollectivePermute()256 auto i = body_b.AddInstruction( in GetFinalFftUsingCollectivePermute()264 auto constant_factor = body_b.AddInstruction(HloInstruction::CreateConstant( in GetFinalFftUsingCollectivePermute()267 constant_factor = body_b.AddInstruction(HloInstruction::CreateBinary( in GetFinalFftUsingCollectivePermute()270 constant_factor = body_b.AddInstruction(HloInstruction::CreateBinary( in GetFinalFftUsingCollectivePermute()[all …]
659 SpmdBuilder body_b("windowed_dot_general_body", original_hlo); in PartitionBaseCase() local667 lhs.state().collective_ops_creator.create_partition_id(&body_b); in PartitionBaseCase()669 body_b.AddInstruction(HloInstruction::CreateConstant( in PartitionBaseCase()672 body_b.AddInstruction(HloInstruction::CreateBinary( in PartitionBaseCase()675 body_b.AddInstruction(HloInstruction::CreateBinary( in PartitionBaseCase()679 body_b.AddInstruction(HloInstruction::CreateBinary( in PartitionBaseCase()681 body_b.AddInstruction(HloInstruction::CreateConstant( in PartitionBaseCase()684 body_b.AddInstruction(HloInstruction::CreateBinary( in PartitionBaseCase()686 body_b.AddInstruction(HloInstruction::CreateConstant( in PartitionBaseCase()690 body_b.AddInstruction(HloInstruction::CreateBinary( in PartitionBaseCase()[all …]