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
16 #include "tensorflow/core/kernels/eigen_contraction_kernel.h"
17
18 #include <mutex> // NOLINT(build/c++11)
19
20 // We need a pair of compile time and runtime flags to disable compilation of
21 // custom contraction kernels for unsupported architectures (e.g. Android,
22 // iOS, ARM and PPC CPUs, etc...), and to be able to fallback on default Eigen
23 // matrix multiplication at runtime.
24 //
25 // It's not allowed to use absl flags library in Tensorflow, so we have to pass
26 // the configuration through the environment variable.
27 //
28 // Example:
29 // bazel test --test_env=TENSORFLOW_USE_CUSTOM_CONTRACTION_KERNEL=false //test
30
31 #if defined(TENSORFLOW_USE_CUSTOM_CONTRACTION_KERNEL)
32
33 namespace Eigen {
34 namespace internal {
35
36 // TODO(ezhulenev): This is a temporary workaround for disabling custom kernels
37 // at runtime in tests. We should always rely on compile time flags for that.
38 // Example: ... --test_env=TENSORFLOW_USE_CUSTOM_CONTRACTION_KERNEL=false //test
UseCustomContractionKernels()39 bool UseCustomContractionKernels() {
40 static bool use_custom_contraction_kernel = true;
41
42 static std::once_flag initialized;
43 std::call_once(initialized, [&] {
44 char* flag = std::getenv("TENSORFLOW_USE_CUSTOM_CONTRACTION_KERNEL");
45 if (flag && (strcmp(flag, "false") == 0 || strcmp(flag, "0") == 0)) {
46 use_custom_contraction_kernel = false;
47 }
48 });
49
50 return use_custom_contraction_kernel;
51 }
52
53 } // namespace internal
54 } // namespace Eigen
55 #endif
56