| 57 | } |
| 58 | |
| 59 | XlaOp TopKWithPartitions(XlaOp input, int64 k, int64 num_partitions) { |
| 60 | XlaBuilder* const builder = input.builder(); |
| 61 | return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 62 | TF_ASSIGN_OR_RETURN(Shape input_shape, builder->GetShape(input)); |
| 63 | int last_dim = input_shape.dimensions_size() - 1; |
| 64 | // Calculate per partition size. |
| 65 | auto input_dims = input_shape.dimensions(); |
| 66 | int64 last_dim_size = input_shape.dimensions(last_dim); |
| 67 | const int64 per_partition_size = CeilOfRatio(last_dim_size, num_partitions); |
| 68 | // Do normal TopK when per partition size is smaller than or equal to k. |
| 69 | if (k >= per_partition_size) { |
| 70 | return TopK(input, k); |
| 71 | } |
| 72 | |
| 73 | Shape iota_shape = |
| 74 | ShapeUtil::MakeShape(S32, AsInt64Slice(input_shape.dimensions())); |
| 75 | XlaOp iota_s32 = Iota(builder, iota_shape, last_dim); |
| 76 | for (int64 i = 0; i < input_shape.rank(); ++i) { |
| 77 | if (input_shape.is_dynamic_dimension(i)) { |
| 78 | // Propagate dynamic dimension from inputs to iota. |
| 79 | iota_s32 = SetDimensionSize(iota_s32, GetDimensionSize(input, i), i); |
| 80 | } |
| 81 | } |
| 82 | |
| 83 | XlaOp values, indices; |
| 84 | for (int64 partition = 0; partition < num_partitions; partition++) { |
| 85 | std::vector<int64> start_indices(input_shape.dimensions_size(), 0); |
| 86 | std::vector<int64> limit_indices(input_dims.begin(), input_dims.end()); |
| 87 | std::vector<int64> strides(input_shape.dimensions_size(), 1); |
| 88 | start_indices[last_dim] = partition * per_partition_size; |
| 89 | limit_indices[last_dim] = |
| 90 | std::min((partition + 1) * per_partition_size, last_dim_size); |
| 91 | // Slice value and indices for this partition.. |
| 92 | XlaOp sliced_input = Slice(input, start_indices, limit_indices, strides); |
| 93 | XlaOp sliced_indices = |
| 94 | Slice(iota_s32, start_indices, limit_indices, strides); |
| 95 | // Concat with previous results. |
| 96 | if (partition > 0) { |
| 97 | sliced_input = ConcatInDim(builder, {values, sliced_input}, last_dim); |
| 98 | sliced_indices = |
| 99 | ConcatInDim(builder, {indices, sliced_indices}, last_dim); |
| 100 | } |
| 101 | // Sort this slice |
| 102 | XlaOp sort_result = |
| 103 | Sort({sliced_input, sliced_indices}, |
| 104 | CreateScalarGtComputation({input_shape.element_type(), S32}, |
| 105 | sliced_indices.builder()), |
| 106 | last_dim, /*is_stable=*/true); |
| 107 | // Slice topk. |
| 108 | start_indices[last_dim] = 0; |
| 109 | limit_indices[last_dim] = k; |
| 110 | values = Slice(GetTupleElement(sort_result, 0), start_indices, |
| 111 | limit_indices, strides); |
| 112 | indices = Slice(GetTupleElement(sort_result, 1), start_indices, |
| 113 | limit_indices, strides); |
| 114 | } |
| 115 | return Tuple(builder, {values, indices}); |
| 116 | }); |