diff --git a/tensorflow/core/kernels/mlir_generated/BUILD b/tensorflow/core/kernels/mlir_generated/BUILD index f55b7df5165..850432a4419 100644 --- a/tensorflow/core/kernels/mlir_generated/BUILD +++ b/tensorflow/core/kernels/mlir_generated/BUILD @@ -40,7 +40,9 @@ filegroup( name = "kernel_srcs", srcs = if_mlir_unranked_kernels_enabled( [ - "unranked_gpu_abs.cc", + "unranked_op_gpu_abs.cc", + "unranked_op_gpu_base.h", + "unranked_op_gpu_base.cc", ], [ "cwise_op_gpu_abs.cc", diff --git a/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_abs.cc b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_abs.cc new file mode 100644 index 00000000000..19a77cdbed1 --- /dev/null +++ b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_abs.cc @@ -0,0 +1,27 @@ +/* Copyright 2020 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/unranked_op_gpu_base.h" + +namespace tensorflow { + +REGISTER_AND_GENERATE_KERNEL(Abs, f16, Eigen::half, DT_HALF); +REGISTER_AND_GENERATE_KERNEL(Abs, f32, float, DT_FLOAT); +REGISTER_AND_GENERATE_KERNEL(Abs, f64, double, DT_DOUBLE); +REGISTER_AND_GENERATE_KERNEL(Abs, i32, int32, DT_INT32); +REGISTER_AND_GENERATE_KERNEL(Abs, i64, int64, DT_INT64); + +} // namespace tensorflow diff --git a/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_base.cc b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_base.cc new file mode 100644 index 00000000000..bc4a36facd3 --- /dev/null +++ b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_base.cc @@ -0,0 +1,59 @@ +/* Copyright 2020 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/kernels/mlir_generated/unranked_op_gpu_base.h" + +#include "tensorflow/core/framework/allocation_description.pb.h" +#include "tensorflow/core/framework/tensor.h" + +namespace tensorflow { +namespace { + +// A simple TensorBuffer implementation that allows us to create Tensors that +// take ownership of pre-allocated memory. +class MlirTensorBuffer : public TensorBuffer { + public: + MlirTensorBuffer(const void* ptr, size_t size, Allocator* allocator) + : TensorBuffer(const_cast(ptr)), + size_(size), + allocator_(allocator) {} + + ~MlirTensorBuffer() override { + if (data()) { + allocator_->DeallocateRaw(data()); + } + } + + size_t size() const override { return size_; } + + TensorBuffer* root_buffer() override { return this; } + + void FillAllocationDescription(AllocationDescription* proto) const override { + proto->set_allocated_bytes(size_); + } + + private: + size_t size_; + Allocator* allocator_; +}; + +} // namespace + +TensorBuffer* GetMlirTensorBuffer(const void* ptr, size_t size, + Allocator* allocator) { + return new MlirTensorBuffer(ptr, size, allocator); +} + +} // namespace tensorflow diff --git a/tensorflow/core/kernels/mlir_generated/unranked_gpu_abs.cc b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_base.h similarity index 74% rename from tensorflow/core/kernels/mlir_generated/unranked_gpu_abs.cc rename to tensorflow/core/kernels/mlir_generated/unranked_op_gpu_base.h index 24625b75ccf..b9bfab9833b 100644 --- a/tensorflow/core/kernels/mlir_generated/unranked_gpu_abs.cc +++ b/tensorflow/core/kernels/mlir_generated/unranked_op_gpu_base.h @@ -13,25 +13,18 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include -#include -#include -#include +#ifndef TENSORFLOW_CORE_KERNELS_MLIR_GENERATED_UNRANKED_OP_GPU_ABS_H_ +#define TENSORFLOW_CORE_KERNELS_MLIR_GENERATED_UNRANKED_OP_GPU_ABS_H_ -#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "mlir/ExecutionEngine/CRunnerUtils.h" // from @llvm-project -#include "mlir/ExecutionEngine/RunnerUtils.h" // from @llvm-project -#include "tensorflow/core/framework/allocation_description.pb.h" -#include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_kernel.h" -#include "tensorflow/core/framework/tensor_types.h" -#include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/platform/logging.h" -#include "tensorflow/core/platform/stream_executor.h" namespace tensorflow { -namespace { + +// Returns a pointer to an allocated MlirTensorBuffer that takes ownership of +// pre-allocated memory. +TensorBuffer* GetMlirTensorBuffer(const void* ptr, size_t size, + Allocator* allocator); template ::UnrankedMemRefType ConvertTensorToDescriptor(const Tensor& tensor) { @@ -58,35 +51,6 @@ template return result; } -// A simple TensorBuffer implementation that allows us to create Tensors that -// take ownership of pre-allocated memory. -template -class MlirTensorBuffer : public TensorBuffer { - public: - MlirTensorBuffer(const void* ptr, TensorShape shape, Allocator* allocator) - : TensorBuffer(const_cast(ptr)), - size_(sizeof(ElemType) * shape.num_elements()), - allocator_(allocator) {} - - ~MlirTensorBuffer() override { - if (data()) { - allocator_->DeallocateRaw(data()); - } - } - - size_t size() const override { return size_; } - - TensorBuffer* root_buffer() override { return this; } - - void FillAllocationDescription(AllocationDescription* proto) const override { - proto->set_allocated_bytes(size_); - } - - private: - size_t size_; - Allocator* allocator_; -}; - template Tensor ConvertDescriptorToTensor( ::UnrankedMemRefType unranked_descriptor, DataType tf_data_type, @@ -97,8 +61,9 @@ Tensor ConvertDescriptorToTensor( for (int i = 0; i < unranked_descriptor.rank; ++i) { result_shape.AddDim(pointers[3 + i]); } - auto* buffer = - new MlirTensorBuffer(base_ptr, result_shape, allocator); + TensorBuffer* buffer = GetMlirTensorBuffer( + base_ptr, sizeof(ElemType) * result_shape.num_elements(), allocator); + // Tensor takes ownership of the buffer. Tensor tensor{tf_data_type, result_shape, buffer}; // When Tensor is constructed, its ref-counter is incremented. We need to @@ -107,8 +72,6 @@ Tensor ConvertDescriptorToTensor( return tensor; } -} // namespace - #define MLIR_FUNCTION(data_type) _mlir_ciface_abs_##data_type // Generates a class derived from OpKernel with Compute function that converts @@ -149,10 +112,6 @@ Tensor ConvertDescriptorToTensor( Name(#kernel_name).Device(DEVICE_GPU).TypeConstraint("T"), \ MlirUnranked##kernel_name##type_name##Op); -REGISTER_AND_GENERATE_KERNEL(Abs, f16, Eigen::half, DT_HALF); -REGISTER_AND_GENERATE_KERNEL(Abs, f32, float, DT_FLOAT); -REGISTER_AND_GENERATE_KERNEL(Abs, f64, double, DT_DOUBLE); -REGISTER_AND_GENERATE_KERNEL(Abs, i32, int32, DT_INT32); -REGISTER_AND_GENERATE_KERNEL(Abs, i64, int64, DT_INT64); - } // namespace tensorflow + +#endif // TENSORFLOW_CORE_KERNELS_MLIR_GENERATED_UNRANKED_OP_GPU_ABS_H_