| 292 | } // namespace |
| 293 | |
| 294 | Status EmitSortInPlace( |
| 295 | int64 dimension_to_sort, const std::vector<IrArray>& values_arrays, |
| 296 | absl::string_view name, absl::Span<const int64> xor_masks, |
| 297 | llvm::IRBuilder<>* b, const gpu::LaunchDimensions& launch_dimensions, |
| 298 | int64 num_iterations_in_sort_dim, const int64 tile_size, |
| 299 | const EmitCallToNestedComputationCallback& emit_compare_callback) { |
| 300 | // Iterate through the keys shape in physical order, but skip the dimension to |
| 301 | // sort and make it the innermost loop which is the loop where the comparisons |
| 302 | // happen. In the dimension to sort, if we use tiling, we iterate through it |
| 303 | // in tiles of 64 elements each, so we use another loop that happens within |
| 304 | // one thread to process this tile worth of data (thereby combining several |
| 305 | // comparison stages of the bitonic sort algorithm because they all happen |
| 306 | // within those 64 elements and are therefore independent of the other |
| 307 | // comparisons). |
| 308 | |
| 309 | const Shape& keys_shape = values_arrays[0].GetShape(); |
| 310 | int64 rank = keys_shape.rank(); |
| 311 | int64 dimension_to_sort_bound = keys_shape.dimensions(dimension_to_sort); |
| 312 | std::vector<int64> dimensions_in_iteration_order(rank); |
| 313 | std::vector<int64> iteration_order_to_logical_order(rank); |
| 314 | int64 dim = 0; |
| 315 | for (int64 dimension : LayoutUtil::MinorToMajor(keys_shape)) { |
| 316 | if (dimension != dimension_to_sort) { |
| 317 | dimensions_in_iteration_order[dim] = keys_shape.dimensions(dimension); |
| 318 | iteration_order_to_logical_order[dim++] = dimension; |
| 319 | } |
| 320 | } |
| 321 | dimensions_in_iteration_order[dim] = num_iterations_in_sort_dim; |
| 322 | iteration_order_to_logical_order[dim] = dimension_to_sort; |
| 323 | |
| 324 | Shape iteration_shape = ShapeUtil::MakeShape(keys_shape.element_type(), |
| 325 | dimensions_in_iteration_order); |
| 326 | |
| 327 | // Allocate shared memory for the tiled compare loop. |
| 328 | std::vector<llvm::Value*> param_shmem_buffers(values_arrays.size(), nullptr); |
| 329 | if (xor_masks.size() > 1) { |
| 330 | llvm::Module* module = b->GetInsertBlock()->getParent()->getParent(); |
| 331 | for (int64 i = 0; i < values_arrays.size(); ++i) { |
| 332 | llvm::Type* tile_type = llvm::ArrayType::get( |
| 333 | llvm_ir::PrimitiveTypeToIrType( |
| 334 | values_arrays[i].GetShape().element_type(), module), |
| 335 | tile_size); |
| 336 | param_shmem_buffers[i] = llvm_ir::AllocateSharedMemoryTile( |
| 337 | module, tile_type, absl::StrCat(name, "_tile_param_", i)); |
| 338 | } |
| 339 | } |
| 340 | |
| 341 | auto compare_loop_body_emitter = |
| 342 | [&](const IrArray::Index& tiles_index) -> Status { |
| 343 | // Naive C++ code for the inner compare loop: |
| 344 | // |
| 345 | // for (int64 i = 0; i < dimension_to_sort_bound; ++i) { |
| 346 | // int64 j = i ^ xor_mask; |
| 347 | // /* emitted in EmitCompareLoopBody() */ |
| 348 | // if (i < j && j < dimension_to_sort_bound) { |
| 349 | // int64 min_key = std::min(keys[i], keys[j]); |
| 350 | // keys[j] = std::max(keys[i], keys[j]); |
| 351 | // keys[i] = min_key; |
no test coverage detected