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

[KERNEL_GEN] Move common code for unranked cwise ops to unranked_op_gpu_base.h.

Browse Source 
PiperOrigin-RevId: 335685664
Change-Id: I3bf9af855d2fbd7f71a9e921ca33e717df01ec5a
This commit is contained in:
Alexander Belyaev 2020-10-06 11:42:08 -07:00 committed by TensorFlower Gardener
parent 95776148da
commit 7c769ba200
4 changed files with 101 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -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",

27
tensorflow/core/kernels/mlir_generated/unranked_op_gpu_abs.cc Normal file
View File

@ -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

59
tensorflow/core/kernels/mlir_generated/unranked_op_gpu_base.cc Normal file
View File

@ -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<void*>(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

65
tensorflow/core/kernels/mlir_generated/unranked_gpu_abs.cc → tensorflow/core/kernels/mlir_generated/unranked_op_gpu_base.h
View File

@ -13,25 +13,18 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <array>
#include <cstdint>
#include <string>
#include <vector>
#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 <typename ElemType>
::UnrankedMemRefType<ElemType> ConvertTensorToDescriptor(const Tensor& tensor) {
@ -58,35 +51,6 @@ template <typename ElemType>
return result;
}
// A simple TensorBuffer implementation that allows us to create Tensors that
// take ownership of pre-allocated memory.
template <typename ElemType>
class MlirTensorBuffer : public TensorBuffer {
public:
MlirTensorBuffer(const void* ptr, TensorShape shape, Allocator* allocator)
: TensorBuffer(const_cast<void*>(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 <typename ElemType>
Tensor ConvertDescriptorToTensor(
::UnrankedMemRefType<ElemType> 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<ElemType>(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<data_type>("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_