| 565 | } |
| 566 | |
| 567 | Status IrEmitter::HandleSort(HloInstruction* hlo) { |
| 568 | const HloSortInstruction* sort = Cast<HloSortInstruction>(hlo); |
| 569 | TF_RETURN_IF_ERROR(EmitTargetAddressForOp(sort)); |
| 570 | Shape keys_shape = sort->keys()->shape(); |
| 571 | PrimitiveType keys_type = keys_shape.element_type(); |
| 572 | switch (keys_type) { |
| 573 | case PRED: |
| 574 | case S8: |
| 575 | case U8: |
| 576 | case S16: |
| 577 | case U16: |
| 578 | case BF16: |
| 579 | case F16: |
| 580 | case S32: |
| 581 | case U32: |
| 582 | case F32: |
| 583 | case S64: |
| 584 | case U64: |
| 585 | case F64: |
| 586 | break; |
| 587 | default: |
| 588 | return Unimplemented( |
| 589 | "Element type %s not supported in the Sort op on CPU.", |
| 590 | PrimitiveType_Name(keys_type)); |
| 591 | } |
| 592 | std::vector<llvm::Value*> destination_addresses(sort->operand_count()); |
| 593 | for (int64 i = 0; i < sort->operand_count(); ++i) { |
| 594 | ShapeIndex shape_index = |
| 595 | sort->values_count() > 0 ? ShapeIndex({i}) : ShapeIndex({}); |
| 596 | const HloInstruction* operand = sort->operand(i); |
| 597 | // We assume that the layout of all involved operands and outputs is the |
| 598 | // same. |
| 599 | TF_RET_CHECK( |
| 600 | LayoutUtil::LayoutsInShapesEqual(keys_shape, operand->shape())); |
| 601 | TF_RET_CHECK(LayoutUtil::LayoutsInShapesEqual( |
| 602 | keys_shape, ShapeUtil::GetSubshape(sort->shape(), shape_index))); |
| 603 | |
| 604 | // The sort is implemented in-place, therefore we first copy the operand |
| 605 | // buffer to the output buffer if they are not the same. |
| 606 | auto destination_buffer = GetAllocationSlice(*sort, shape_index); |
| 607 | destination_addresses[i] = |
| 608 | EmitBufferPointer(destination_buffer, operand->shape()); |
| 609 | auto source_address = GetAllocationSlice(*operand); |
| 610 | if (destination_buffer != source_address) { |
| 611 | int64 primitive_type_size = |
| 612 | ShapeUtil::ByteSizeOfPrimitiveType(operand->shape().element_type()); |
| 613 | auto source_buffer = GetEmittedValueFor(operand); |
| 614 | int64 size = ByteSizeOf(operand->shape()); |
| 615 | MemCpy(destination_addresses[i], /*DstAlign=*/primitive_type_size, |
| 616 | source_buffer, |
| 617 | /*SrcAlign=*/primitive_type_size, size); |
| 618 | } |
| 619 | } |
| 620 | |
| 621 | // Normalize the shape and the dimension to sort. |
| 622 | Shape normalized_keys_shape = |
| 623 | ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout(keys_shape); |
| 624 | int64 physical_dimension_to_sort = LayoutUtil::MakeLogicalToPhysical( |
nothing calls this directly
no test coverage detected