1 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 #include "tensorflow/compiler/xla/service/cpu/runtime_key_value_sort.h"
16
17 #include <algorithm>
18 #include <cstring>
19 #include <memory>
20 #include <numeric>
21 #include <string>
22
23 #include "absl/base/dynamic_annotations.h"
24 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
25
__xla_cpu_runtime_KeyValueSort(int64_t a,int64_t b,int64_t c,char ** values,int32_t values_count,int32_t * values_primitive_type_size_in_bytes,bool is_stable,char * run_options,int64_t * prof_counters,void (* less_than)(char *,char *,char **,char **,int64_t *))26 ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSort(
27 int64_t a, int64_t b, int64_t c, char** values, int32_t values_count,
28 int32_t* values_primitive_type_size_in_bytes, bool is_stable,
29 char* run_options, int64_t* prof_counters,
30 void (*less_than)(char*, char*, char**, char**, int64_t*)) {
31 // 'values' and 'values_primitive_type_size_in_bytes' are managed by the JIT
32 // code, so msan can't tell they are initialized.
33 ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(values, values_count * sizeof(char*));
34 ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(values_primitive_type_size_in_bytes,
35 values_count * sizeof(int32_t));
36
37 // High-level idea of the iteration/sorting logic:
38 // Conceptually we have a 3-dimensional shape [a, b, c]. b corresponds to the
39 // dimension to sort, c is the product of the more minor dimensions (set to 1
40 // if b is the most minor dimension), and a is the product of the more major
41 // dimensions (set to 1 if b is the most major dimension). There are a * c
42 // many rows that we need to sort. We iterate through these, calculate a
43 // 'base_offset' value which points to the first element in that row, and add
44 // i * c for accessing the 'i'-th element in that row.
45
46 int64_t sort_dimension_elements = b;
47 int64_t num_iteration_elements = a * c;
48 int64_t sort_dimension_offset = c;
49
50 std::unique_ptr<int64_t[]> indices(new int64_t[sort_dimension_elements]);
51 std::unique_ptr<char*[]> comparison_values(new char*[2 * values_count]);
52 std::iota(indices.get(), indices.get() + sort_dimension_elements, 0);
53 std::unique_ptr<std::string[]> reordered_values(
54 new std::string[sort_dimension_elements]);
55 for (int64_t index = 0; index < num_iteration_elements; ++index) {
56 // If the sort should be stable, we have to reinitialize indices to iota to
57 // guarantee that we still keep the relative order in case of ties.
58 if (is_stable && index > 0) {
59 std::iota(indices.get(), indices.get() + sort_dimension_elements, 0);
60 }
61 // 'index' can be split into two values which index into the 'c' dimension
62 // and the 'a' dimension, respectively. 'index' % 'c' is the index into the
63 // 'c' dimension, 'index' / 'c' is the index into the 'a' dimension. When
64 // calculating the base offset, we need to multiply the index into the 'a'
65 // dimension with 'b' * 'c'.
66 // 'index' / 'c' * 'c' * 'b' = ('index' - 'index' % 'c') * 'b'.
67 int64_t base_offset =
68 index % sort_dimension_offset +
69 (index - index % sort_dimension_offset) * sort_dimension_elements;
70 auto compare_function = [&](int64_t a, int64_t b) -> bool {
71 for (int32_t i = 0; i < values_count; ++i) {
72 int64_t memory_index_lhs = (base_offset + a * sort_dimension_offset) *
73 values_primitive_type_size_in_bytes[i];
74 int64_t memory_index_rhs = (base_offset + b * sort_dimension_offset) *
75 values_primitive_type_size_in_bytes[i];
76 comparison_values[i * 2] = values[i] + memory_index_lhs;
77 comparison_values[i * 2 + 1] = values[i] + memory_index_rhs;
78 }
79 char result = 0; // Overwritten by less_than.
80 less_than(&result, run_options, comparison_values.get(), nullptr,
81 prof_counters);
82 return result != 0u;
83 };
84 if (is_stable) {
85 std::stable_sort(indices.get(), indices.get() + sort_dimension_elements,
86 compare_function);
87 } else {
88 std::sort(indices.get(), indices.get() + sort_dimension_elements,
89 compare_function);
90 }
91
92 // Reorder the values according to the order defined by 'indices'.
93 for (int32_t idx = 0; idx < values_count; ++idx) {
94 for (int64_t i = 0; i < sort_dimension_elements; ++i) {
95 int64_t memory_index =
96 (base_offset + indices[i] * sort_dimension_offset) *
97 values_primitive_type_size_in_bytes[idx];
98
99 reordered_values[i] =
100 std::string(values[idx] + memory_index,
101 values_primitive_type_size_in_bytes[idx]);
102 }
103 for (int64_t i = 0; i < sort_dimension_elements; ++i) {
104 int64_t memory_index = (base_offset + i * sort_dimension_offset) *
105 values_primitive_type_size_in_bytes[idx];
106 memcpy(values[idx] + memory_index, reordered_values[i].c_str(),
107 values_primitive_type_size_in_bytes[idx]);
108 }
109 }
110 }
111 }
112