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 "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
24 #include "tensorflow/core/platform/dynamic_annotations.h"
25 #include "tensorflow/core/platform/macros.h"
26 #include "tensorflow/core/platform/types.h"
27
28 namespace {
29 using tensorflow::int32;
30 using tensorflow::int64;
31 } // namespace
32
__xla_cpu_runtime_KeyValueSort(int64 a,int64 b,int64 c,char ** values,int32 values_count,int32 * values_primitive_type_size_in_bytes,bool is_stable,char * run_options,int64 * prof_counters,void (* less_than)(char *,char *,char **,char **,tensorflow::int64 *))33 TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSort(
34 int64 a, int64 b, int64 c, char** values, int32 values_count,
35 int32* values_primitive_type_size_in_bytes, bool is_stable,
36 char* run_options, int64* prof_counters,
37 void (*less_than)(char*, char*, char**, char**, tensorflow::int64*)) {
38 // 'values' and 'values_primitive_type_size_in_bytes' are managed by the JIT
39 // code, so msan can't tell they are initialized.
40 TF_ANNOTATE_MEMORY_IS_INITIALIZED(values, values_count * sizeof(char*));
41 TF_ANNOTATE_MEMORY_IS_INITIALIZED(values_primitive_type_size_in_bytes,
42 values_count * sizeof(int32));
43
44 // High-level idea of the iteration/sorting logic:
45 // Conceptually we have a 3-dimensional shape [a, b, c]. b corresponds to the
46 // dimension to sort, c is the product of the more minor dimensions (set to 1
47 // if b is the most minor dimension), and a is the product of the more major
48 // dimensions (set to 1 if b is the most major dimension). There are a * c
49 // many rows that we need to sort. We iterate through these, calculate a
50 // 'base_offset' value which points to the first element in that row, and add
51 // i * c for accessing the 'i'-th element in that row.
52
53 int64 sort_dimension_elements = b;
54 int64 num_iteration_elements = a * c;
55 int64 sort_dimension_offset = c;
56
57 std::unique_ptr<int64[]> indices(new int64[sort_dimension_elements]);
58 std::unique_ptr<char*[]> comparison_values(new char*[2 * values_count]);
59 std::iota(indices.get(), indices.get() + sort_dimension_elements, 0);
60 std::unique_ptr<std::string[]> reordered_values(
61 new std::string[sort_dimension_elements]);
62 for (int64 index = 0; index < num_iteration_elements; ++index) {
63 // 'index' can be split into two values which index into the 'c' dimension
64 // and the 'a' dimension, respectively. 'index' % 'c' is the index into the
65 // 'c' dimension, 'index' / 'c' is the index into the 'a' dimension. When
66 // calculating the base offset, we need to multiply the index into the 'a'
67 // dimension with 'b' * 'c'.
68 // 'index' / 'c' * 'c' * 'b' = ('index' - 'index' % 'c') * 'b'.
69 int64 base_offset =
70 index % sort_dimension_offset +
71 (index - index % sort_dimension_offset) * sort_dimension_elements;
72 auto compare_function = [&](int64 a, int64 b) -> bool {
73 int64 memory_index_lhs = (base_offset + a * sort_dimension_offset) *
74 values_primitive_type_size_in_bytes[0];
75 int64 memory_index_rhs = (base_offset + b * sort_dimension_offset) *
76 values_primitive_type_size_in_bytes[0];
77 for (int32 i = 0; i < values_count; ++i) {
78 comparison_values[i * 2] = values[i] + memory_index_lhs;
79 comparison_values[i * 2 + 1] = values[i] + memory_index_rhs;
80 }
81 char result = 0; // Overwritten by less_than.
82 less_than(&result, run_options, comparison_values.get(), nullptr,
83 prof_counters);
84 return result != 0u;
85 };
86 if (is_stable) {
87 std::stable_sort(indices.get(), indices.get() + sort_dimension_elements,
88 compare_function);
89 } else {
90 std::sort(indices.get(), indices.get() + sort_dimension_elements,
91 compare_function);
92 }
93
94 // Reorder the values according to the order defined by 'indices'.
95 for (int32 idx = 0; idx < values_count; ++idx) {
96 for (int64 i = 0; i < sort_dimension_elements; ++i) {
97 int64 memory_index =
98 (base_offset + indices[i] * sort_dimension_offset) *
99 values_primitive_type_size_in_bytes[idx];
100
101 reordered_values[i] =
102 std::string(values[idx] + memory_index,
103 values_primitive_type_size_in_bytes[idx]);
104 }
105 for (int64 i = 0; i < sort_dimension_elements; ++i) {
106 int64 memory_index = (base_offset + i * sort_dimension_offset) *
107 values_primitive_type_size_in_bytes[idx];
108 memcpy(values[idx] + memory_index, reordered_values[i].c_str(),
109 values_primitive_type_size_in_bytes[idx]);
110 }
111 }
112 }
113 }
114