| 1047 | } |
| 1048 | |
| 1049 | Status IrEmitterUnnested::HandleSort(HloInstruction* sort) { |
| 1050 | std::vector<std::unique_ptr<Thunk>> thunks; |
| 1051 | Shape keys_shape = sort->operand(0)->shape(); |
| 1052 | int64 dimension_to_sort = sort->dimensions(0); |
| 1053 | for (int64 i = 0; i < sort->operand_count(); ++i) { |
| 1054 | ShapeIndex shape_index = |
| 1055 | sort->operand_count() > 1 ? ShapeIndex({i}) : ShapeIndex({}); |
| 1056 | // We assume that the layout of all involved operands and outputs is the |
| 1057 | // same. |
| 1058 | TF_RET_CHECK(LayoutUtil::LayoutsInShapesEqual(keys_shape, |
| 1059 | sort->operand(i)->shape())); |
| 1060 | TF_RET_CHECK(LayoutUtil::LayoutsInShapesEqual( |
| 1061 | keys_shape, ShapeUtil::GetSubshape(sort->shape(), shape_index))); |
| 1062 | |
| 1063 | // If possible, we share buffers. If that is not possible, we need to copy |
| 1064 | // the values, because the emitter does the sorting in-place. |
| 1065 | auto destination_buffer = GetAllocationSlice(*sort, shape_index); |
| 1066 | auto source_address = GetAllocationSlice(*sort->operand(i)); |
| 1067 | if (destination_buffer != source_address) { |
| 1068 | // TODO(b/26783907): Figure out why we never seem to share buffers for |
| 1069 | // key/value sort. |
| 1070 | thunks.push_back(absl::make_unique<DeviceToDeviceCopyThunk>( |
| 1071 | /*source_address=*/source_address, |
| 1072 | /*destination_buffer=*/destination_buffer, |
| 1073 | /*mem_size=*/ShapeUtil::ByteSizeOf(sort->operand(i)->shape()), |
| 1074 | nullptr)); |
| 1075 | } |
| 1076 | } |
| 1077 | |
| 1078 | uint64 dimension_to_sort_bound = keys_shape.dimensions(dimension_to_sort); |
| 1079 | int64 num_stages = tensorflow::Log2Ceiling(dimension_to_sort_bound); |
| 1080 | CHECK_GE(1ULL << num_stages, dimension_to_sort_bound); |
| 1081 | CHECK_LT(1ULL << (num_stages - 1), dimension_to_sort_bound); |
| 1082 | |
| 1083 | // Naive C++ code for the outer loops: |
| 1084 | // |
| 1085 | // for (int64 stage = 0; stage < Log2Ceiling(dimension_to_sort_bound); |
| 1086 | // ++stage) { |
| 1087 | // int64 first_xor_mask = (1LL << (stage + 1)) - 1; |
| 1088 | // SortInPlace(first_xor_mask); |
| 1089 | // for (int64 mask = stage - 1; mask >= 0; --mask) { |
| 1090 | // int64 later_xor_mask = 1LL << mask; |
| 1091 | // SortInPlace(later_xor_mask); |
| 1092 | // } |
| 1093 | // } |
| 1094 | // |
| 1095 | // This follows the alternative representation of the algorithm described on |
| 1096 | // Wikipedia: https://en.wikipedia.org/wiki/Bitonic_sorter |
| 1097 | // |
| 1098 | // Each mask specifies how to derive from one position in the array the |
| 1099 | // position with which it should be compared (we calculate the xor of the |
| 1100 | // position with the mask). |
| 1101 | // As an optimization, we can move the 'mask' loop to inside the |
| 1102 | // sorting/comparison loop if the comparisons happen within a small block of |
| 1103 | // the array. To make this work, we collect all consecutive masks that are |
| 1104 | // smaller than our chosen power of 2 tile size, and pass them to SortInPlace. |
| 1105 | // Each thread then processes one tile of data. |
| 1106 |
nothing calls this directly
no test coverage detected