1 /* Copyright 2017 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 #define EIGEN_USE_THREADS 16 #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" 17 18 #include <memory> 19 #include <string> 20 #include <tuple> 21 22 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" 23 #include "tensorflow/compiler/xla/array2d.h" 24 #include "tensorflow/compiler/xla/client/local_client.h" 25 #include "tensorflow/compiler/xla/ptr_util.h" 26 #include "tensorflow/compiler/xla/service/cpu/runtime_matmul.h" 27 #include "tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.h" 28 #include "tensorflow/compiler/xla/types.h" 29 #include "tensorflow/core/common_runtime/eigen_thread_pool.h" 30 #include "tensorflow/core/lib/strings/stringprintf.h" 31 #include "tensorflow/core/platform/env.h" 32 #include "tensorflow/core/platform/logging.h" 33 #include "tensorflow/core/platform/test.h" 34 35 namespace xla { 36 namespace { 37 38 class CpuRuntimeTest : public ::testing::Test {}; 39 40 template <typename T> 41 std::unique_ptr<Array2D<float>> MaybeTransposeArray2D(const Array2D<T>& array, 42 bool transpose) { 43 int64 output_height = array.height(); 44 int64 output_width = array.width(); 45 if (transpose) { 46 std::swap(output_width, output_height); 47 } 48 auto output = MakeUnique<Array2D<float>>(output_height, output_width); 49 for (int y = 0; y < array.height(); y++) { 50 for (int x = 0; x < array.width(); x++) { 51 if (transpose) { 52 (*output)(x, y) = array(y, x); 53 } else { 54 (*output)(y, x) = array(y, x); 55 } 56 } 57 } 58 return output; 59 } 60 61 // Verifies that matrix 'c' equals the result of matrix 'a' times matrix 'b'. 62 // Each element is compared to within a small error bound. 63 void CheckMatrixMultiply(const Array2D<float>& a, const Array2D<float>& b, 64 const Array2D<float>& c) { 65 for (int i = 0; i < a.height(); ++i) { 66 for (int j = 0; j < b.width(); ++j) { 67 float sum = 0.0; 68 for (int k = 0; k < a.width(); ++k) { 69 sum += a(i, k) * b(k, j); 70 } 71 EXPECT_NEAR(sum, c(i, j), 0.01); 72 } 73 } 74 } 75 76 std::unique_ptr<Array2D<float>> EigenMatrixMultiply(const Array2D<float>& a, 77 const Array2D<float>& b, 78 bool transpose_lhs, 79 bool transpose_rhs, 80 bool single_threaded) { 81 CHECK_EQ(a.width(), b.height()); 82 int64 m = a.height(); 83 int64 n = b.width(); 84 int64 k = a.width(); 85 86 // The Eigen matmul runtime function expects the matrix to be in column major 87 // order and array2d is in row-major order. Create transposes of a and b. The 88 // 'data' buffer in the transposed array is the original array in column major 89 // order. 90 auto a_transpose = MaybeTransposeArray2D(a, !transpose_lhs); 91 auto b_transpose = MaybeTransposeArray2D(b, !transpose_rhs); 92 93 // Since we're going to transpose c before returning it. Swap the order of the 94 // dimension sizes to ensure the returned array is properly dimensioned. 95 auto c_transpose = MakeUnique<Array2D<float>>(n, m); 96 if (single_threaded) { 97 __xla_cpu_runtime_EigenSingleThreadedMatMulF32( 98 nullptr, c_transpose->data(), a_transpose->data(), b_transpose->data(), 99 m, n, k, transpose_lhs, transpose_rhs); 100 } else { 101 tensorflow::thread::ThreadPool pool(tensorflow::Env::Default(), "XLAEigen", 102 2); 103 tensorflow::EigenThreadPoolWrapper tp(&pool); 104 Eigen::ThreadPoolDevice device(&tp, tp.NumThreads()); 105 ExecutableRunOptions run_options; 106 run_options.set_intra_op_thread_pool(&device); 107 108 __xla_cpu_runtime_EigenMatMulF32(&run_options, c_transpose->data(), 109 a_transpose->data(), b_transpose->data(), 110 m, n, k, transpose_lhs, transpose_rhs); 111 } 112 return MaybeTransposeArray2D(*c_transpose, true); 113 } 114 115 struct MatMulShape { 116 int64 m; 117 int64 k; 118 int64 n; 119 }; 120 121 MatMulShape MatMulShapes[] = { 122 MatMulShape{2, 2, 3}, MatMulShape{256, 512, 1024}, 123 MatMulShape{128, 128, 1}, MatMulShape{1, 128, 128}, 124 MatMulShape{1, 32, 128}, MatMulShape{1, 32, 16}, 125 MatMulShape{32, 16, 1}, MatMulShape{32, 128, 1}, 126 }; 127 128 // This takes 4 parameters: 129 // * shape of the matmul 130 // * transpose_lhs 131 // * transpose_rhs 132 // * single_threaded 133 using EigenMatMulTestParam = std::tuple<MatMulShape, bool, bool, bool>; 134 135 class EigenMatMulTest 136 : public CpuRuntimeTest, 137 public ::testing::WithParamInterface<EigenMatMulTestParam> { 138 public: 139 static string Name( 140 const ::testing::TestParamInfo<EigenMatMulTestParam>& info) { 141 MatMulShape shape = std::get<0>(info.param); 142 bool transpose_lhs = std::get<1>(info.param); 143 bool transpose_rhs = std::get<2>(info.param); 144 bool single_threaded = std::get<3>(info.param); 145 146 return tensorflow::strings::Printf( 147 "MatMul_%lld_%lld_%lld_%s%s%s_threaded", shape.m, shape.k, shape.n, 148 transpose_lhs ? "Tlhs_" : "", transpose_rhs ? "Trhs_" : "", 149 single_threaded ? "single" : "multi"); 150 } 151 }; // namespace xla 152 153 TEST_P(EigenMatMulTest, DoIt) { 154 MatMulShape shape = std::get<0>(GetParam()); 155 bool transpose_lhs = std::get<1>(GetParam()); 156 bool transpose_rhs = std::get<2>(GetParam()); 157 bool single_threaded = std::get<3>(GetParam()); 158 159 auto a = MakeLinspaceArray2D(0.0, 1.0, shape.m, shape.k); 160 auto b = MakeLinspaceArray2D(-2.0, 2.0, shape.k, shape.n); 161 auto c = EigenMatrixMultiply(*a, *b, transpose_lhs, transpose_rhs, 162 single_threaded); 163 CheckMatrixMultiply(*a, *b, *c); 164 } 165 166 INSTANTIATE_TEST_CASE_P(EigenMatMulTestInstantiaion, EigenMatMulTest, 167 ::testing::Combine(::testing::ValuesIn(MatMulShapes), 168 ::testing::Bool(), ::testing::Bool(), 169 ::testing::Bool()), 170 EigenMatMulTest::Name); 171 172 } // namespace 173 } // namespace xla 174
