experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[KERNEL_GEN][CPU] Add tf.addv2 CPU kernel after MacOs & Win builds were fixed.

Browse Source 
PiperOrigin-RevId: 356280798
Change-Id: I2e2b27c5aea04c68ff1dede4b0a071be07320580
This commit is contained in:
Alexander Belyaev 2021-02-08 09:18:19 -08:00 committed by TensorFlower Gardener
parent 750119433d
commit 4cf728b600
4 changed files with 126 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorflow/core/kernels/cwise_op_add_1.cc
View File

@ -18,8 +18,13 @@ limitations under the License.
namespace tensorflow {
REGISTER6(BinaryOp, CPU, "Add", functor::add, float, Eigen::half, double, int32,
int64, bfloat16);
#if !defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
REGISTER6(BinaryOp, CPU, "AddV2", functor::add, float, Eigen::half, double,
int32, int64, bfloat16);
#else
REGISTER(BinaryOp, CPU, "AddV2", functor::add, bfloat16);
#endif
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
REGISTER3(BinaryOp, GPU, "Add", functor::add, float, Eigen::half, double);

44
tensorflow/core/kernels/mlir_generated/BUILD
View File

@ -260,6 +260,19 @@ tf_kernel_library(
],
)
tf_kernel_library(
name = "cpu_cwise_binary_op",
srcs = [
"cpu_op_add.cc",
],
tags = ["manual"],
deps = [
":base_cpu_op",
":cpu_add_v2_kernels",
"//third_party/eigen3",
],
)
tf_kernel_library(
name = "cwise_op",
srcs = [],
@ -274,7 +287,10 @@ tf_kernel_library(
tf_kernel_library(
name = "experimental_cwise_op",
srcs = [],
deps = [":cpu_cwise_unary_op"] + if_cuda_or_rocm([":gpu_cwise_binary_op"]),
deps = [
":cpu_cwise_unary_op",
":cpu_cwise_binary_op",
] + if_cuda_or_rocm([":gpu_cwise_binary_op"]),
)
cc_library(
@ -379,6 +395,18 @@ tf_cuda_cc_test(
],
)
tf_cuda_cc_test(
name = "cpu_binary_ops_test",
size = "medium",
srcs = if_mlir_generated_gpu_kernels_enabled(["cpu_binary_ops_test.cc"]),
deps = [
":base_binary_ops_test",
":base_ops_test",
"//tensorflow/core/common_runtime:device",
"//tensorflow/core/common_runtime:device_factory",
],
)
# TODO(b/160731748): Re-enable when it works again.
# gpu_kernel_library(
# name = "gpu_bias_add_lib",
@ -991,3 +1019,17 @@ cpu_kernel_library(
],
unroll_factors = "4",
)
cpu_kernel_library(
name = "cpu_add_v2_lib",
op = "add_v2",
tile_size = "1024",
types = [
"f16",
"f32",
"f64",
"i32",
"i64",
],
unroll_factors = "4",
)

52
tensorflow/core/kernels/mlir_generated/cpu_binary_ops_test.cc Normal file
View File

@ -0,0 +1,52 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/kernels/mlir_generated/base_binary_ops_test.h"
#include "tensorflow/core/kernels/mlir_generated/base_ops_test.h"
namespace tensorflow {
namespace {
// Test fixture `BinaryOpsTest` that sets the TF device is expected by the TEST
// macros below.
class BinaryOpsTest : public BinaryOpsTestBase {
protected:
void SetUp() override {
std::unique_ptr<tensorflow::Device> device_cpu(
tensorflow::DeviceFactory::NewDevice("CPU", {},
"/job:a/replica:0/task:0"));
SetDevice(tensorflow::DEVICE_CPU, std::move(device_cpu));
}
};
/// Test `tf.AddV2`.
template <typename T>
T baseline_add(T lhs, T rhs) {
return lhs + rhs;
}
GENERATE_DEFAULT_TESTS(AddV2, /*test_name=*/Half, Eigen::half, Eigen::half,
baseline_add)
GENERATE_DEFAULT_TESTS(AddV2, /*test_name=*/Float, float, float, baseline_add)
GENERATE_DEFAULT_TESTS(AddV2, /*test_name=*/Double, double, double,
baseline_add)
GENERATE_DEFAULT_TESTS(AddV2, /*test_name=*/Int32, int32, int32, baseline_add)
GENERATE_DEFAULT_TESTS(AddV2, /*test_name=*/Int64, int64, int64, baseline_add)
} // namespace
} // namespace tensorflow

26
tensorflow/core/kernels/mlir_generated/cpu_op_add.cc Normal file
View File

@ -0,0 +1,26 @@
/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "tensorflow/core/kernels/mlir_generated/base_cpu_op.h"
namespace tensorflow {
GENERATE_AND_REGISTER_BINARY_CPU_KERNEL(AddV2, f16, DT_HALF, Eigen::half);
GENERATE_AND_REGISTER_BINARY_CPU_KERNEL(AddV2, f32, DT_FLOAT, float);
GENERATE_AND_REGISTER_BINARY_CPU_KERNEL(AddV2, f64, DT_DOUBLE, double);
GENERATE_AND_REGISTER_BINARY_CPU_KERNEL(AddV2, i32, DT_INT32, int32);
GENERATE_AND_REGISTER_BINARY_CPU_KERNEL(AddV2, i64, DT_INT64, int64);
} // namespace tensorflow