1 /* Copyright 2022 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
16 #include "tensorflow/compiler/xla/runtime/custom_call.h"
17
18 #include <string>
19 #include <string_view>
20
21 #include "llvm/Support/raw_ostream.h"
22 #include "tensorflow/compiler/xla/runtime/type_id.h"
23
24 namespace xla {
25 namespace runtime {
26
27 using llvm::ArrayRef;
28 using llvm::raw_ostream;
29
30 template <typename T>
31 using TensorRef = CustomCall::TensorRef<T>;
32
PrintArr(raw_ostream & os,std::string_view name,llvm::ArrayRef<int64_t> arr)33 static void PrintArr(raw_ostream& os, std::string_view name,
34 llvm::ArrayRef<int64_t> arr) {
35 os << " " << name << ": [";
36 auto i64_to_string = [](int64_t v) { return std::to_string(v); };
37 os << llvm::join(llvm::map_range(arr, i64_to_string), ", ");
38 os << "]";
39 }
40
operator <<(raw_ostream & os,const StridedMemrefView & view)41 raw_ostream& operator<<(raw_ostream& os, const StridedMemrefView& view) {
42 os << "StridedMemrefView: dtype: " << view.dtype;
43 PrintArr(os, "sizes", view.sizes);
44 PrintArr(os, "strides", view.strides);
45 return os;
46 }
47
operator <<(raw_ostream & os,const MemrefView & view)48 raw_ostream& operator<<(raw_ostream& os, const MemrefView& view) {
49 os << "MemrefView: dtype: " << view.dtype;
50 PrintArr(os, "sizes", view.sizes);
51 return os;
52 }
53
operator <<(raw_ostream & os,const FlatMemrefView & view)54 raw_ostream& operator<<(raw_ostream& os, const FlatMemrefView& view) {
55 return os << "FlatMemrefView: dtype: " << view.dtype
56 << " size_in_bytes: " << view.size_in_bytes;
57 }
58
PopulateCustomCallTypeIdNames(TypeIDNameRegistry & r)59 void PopulateCustomCallTypeIdNames(TypeIDNameRegistry& r) {
60 r.Register<Tagged<llvm::StringRef>>("__type_id_string");
61
62 r.Register<Tagged<uint8_t>>("__type_id_uint8");
63 r.Register<Tagged<uint32_t>>("__type_id_uint32");
64 r.Register<Tagged<uint64_t>>("__type_id_uint64");
65 r.Register<Tagged<bool>>("__type_id_bool");
66 r.Register<Tagged<int32_t>>("__type_id_int32");
67 r.Register<Tagged<int64_t>>("__type_id_int64");
68 r.Register<Tagged<Eigen::half>>("__type_id_half");
69 r.Register<Tagged<float>>("__type_id_float");
70 r.Register<Tagged<double>>("__type_id_double");
71
72 r.Register<Tagged<MemrefView>>("__type_id_memref_view");
73 r.Register<Tagged<StridedMemrefView>>("__type_id_strided_memref_view");
74 r.Register<Tagged<EmptyArrayRef>>("__type_id_empty_array");
75
76 r.Register<Tagged<ArrayRef<int8_t>>>("__type_id_arrayr_int8");
77 r.Register<Tagged<ArrayRef<int16_t>>>("__type_id_array_int16");
78 r.Register<Tagged<ArrayRef<int32_t>>>("__type_id_array_int32");
79 r.Register<Tagged<ArrayRef<int64_t>>>("__type_id_array_int64");
80 r.Register<Tagged<ArrayRef<float>>>("__type_id_array_float");
81 r.Register<Tagged<ArrayRef<double>>>("__type_id_array_double");
82
83 r.Register<Tagged<TensorRef<int32_t>>>("__type_id__tensor_int32_t");
84 r.Register<Tagged<TensorRef<int64_t>>>("__type_id_tensor_int64_t");
85 r.Register<Tagged<TensorRef<float>>>("__type_id_tensor_float");
86 r.Register<Tagged<TensorRef<double>>>("__type_id_tensor_double");
87 }
88
89 } // namespace runtime
90 } // namespace xla
91
92 XLA_RUNTIME_DEFINE_EXPLICIT_TYPE_ID(llvm::StringRef);
93 XLA_RUNTIME_DEFINE_EXPLICIT_TYPE_ID(xla::runtime::StridedMemrefView);
94 XLA_RUNTIME_DEFINE_EXPLICIT_TYPE_ID(xla::runtime::MemrefView);
95 XLA_RUNTIME_DEFINE_EXPLICIT_TYPE_ID(xla::runtime::FlatMemrefView);
96 XLA_RUNTIME_DEFINE_EXPLICIT_TYPE_ID(int32_t);
97 XLA_RUNTIME_DEFINE_EXPLICIT_TYPE_ID(int64_t);
98 XLA_RUNTIME_DEFINE_EXPLICIT_TYPE_ID(float);
99 XLA_RUNTIME_DEFINE_EXPLICIT_TYPE_ID(double);
100 XLA_RUNTIME_DEFINE_EXPLICIT_TYPE_ID(ArrayRef<int32_t>);
101 XLA_RUNTIME_DEFINE_EXPLICIT_TYPE_ID(ArrayRef<int64_t>);
102 XLA_RUNTIME_DEFINE_EXPLICIT_TYPE_ID(ArrayRef<float>);
103 XLA_RUNTIME_DEFINE_EXPLICIT_TYPE_ID(ArrayRef<double>);
104