wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/core/kernels/BUILD
+8496
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/kernels/aggregate_ops.cc
+229
View File
@@ -0,0 +1,229 @@
/* Copyright 2015 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.
==============================================================================*/
// See docs in ../ops/math_ops.cc.
#define EIGEN_USE_THREADS
#include "tensorflow/core/kernels/aggregate_ops.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/op_requires.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/variant.h"
#include "tensorflow/core/framework/variant_op_registry.h"
#include "tensorflow/core/kernels/aggregate_ops_cpu.h"
#include "tensorflow/core/kernels/variant_ops_util.h"
#include "tensorflow/core/lib/gtl/inlined_vector.h"
namespace tensorflow {
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
template <typename Device, typename T>
class AddNOp : public OpKernel {
public:
explicit AddNOp(OpKernelConstruction* context) : OpKernel(context) {}
void Compute(OpKernelContext* ctx) override {
if (!ctx->ValidateInputsAreSameShape(this)) return;
const Tensor& input0 = ctx->input(0);
const int num = ctx->num_inputs();
if (num == 1) {
ctx->set_output(0, input0);
return;
}
// Try to forward and accumulate the result in one of the input buffers.
int reused_input = -1;
absl::InlinedVector<int, 8UL> input_indices(num);
std::iota(input_indices.begin(), input_indices.end(), 0);
Tensor* output = nullptr;
for (int input_idx = 0; input_idx < num; ++input_idx) {
if (ctx->forward_input_to_output_with_shape(input_idx, 0, input0.shape(),
&output)) {
reused_input = input_idx;
break;
}
}
if (reused_input == -1) {
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, input0.shape(), &output));
} else if (reused_input > 0) {
// Move the forwarded buffer to the front so we don't double count
// anything if there are more than 8 inputs.
input_indices[0] = reused_input;
input_indices[reused_input] = 0;
}
auto To = output->flat<T>();
#define I(IDX) ctx->input(input_indices[IDX]).template flat<T>()
#if defined(__ANDROID_TYPES_SLIM__)
// On Android by default,we only support additions of two arguments, so we
// can reduce the number of template instantiations.
OP_REQUIRES(ctx, num == 2,
errors::InvalidArgument("Only additions of two arguments "
"supported. Num inputs: ",
num));
functor::Add2Functor<Device, T> functor2;
functor2(ctx->template eigen_device<Device>(), To, I(0), I(1));
#else
static const int kWidth = 8;
int r = num % kWidth;
switch (r) {
case 2: {
functor::Add2Functor<Device, T> functor2;
functor2(ctx->template eigen_device<Device>(), To, I(0), I(1));
break;
}
case 3: {
functor::Add3Functor<Device, T> functor3;
functor3(ctx->template eigen_device<Device>(), To, I(0), I(1), I(2));
break;
}
case 4: {
functor::Add4Functor<Device, T> functor4;
functor4(ctx->template eigen_device<Device>(), To, I(0), I(1), I(2),
I(3));
break;
}
case 5: {
functor::Add5Functor<Device, T> functor5;
functor5(ctx->template eigen_device<Device>(), To, I(0), I(1), I(2),
I(3), I(4));
break;
}
case 6: {
functor::Add6Functor<Device, T> functor6;
functor6(ctx->template eigen_device<Device>(), To, I(0), I(1), I(2),
I(3), I(4), I(5));
break;
}
case 7: {
functor::Add7Functor<Device, T> functor7;
functor7(ctx->template eigen_device<Device>(), To, I(0), I(1), I(2),
I(3), I(4), I(5), I(6));
break;
}
case 0: {
functor::Add8Functor<Device, T> functor8;
functor8(ctx->template eigen_device<Device>(), To, I(0), I(1), I(2),
I(3), I(4), I(5), I(6), I(7));
r = 8;
break;
}
case 1: {
functor::Add9Functor<Device, T> functor9;
functor9(ctx->template eigen_device<Device>(), To, I(0), I(1), I(2),
I(3), I(4), I(5), I(6), I(7), I(8));
r = 9;
break;
}
}
for (; r < num; r += kWidth) {
functor::Add8pFunctor<Device, T> functor8p;
functor8p(ctx->template eigen_device<Device>(), To, I(r), I(r + 1),
I(r + 2), I(r + 3), I(r + 4), I(r + 5), I(r + 6), I(r + 7));
}
#endif // defined(__ANDROID_TYPES_SLIM__)
#undef I
}
};
template <typename Device>
class AddNOp<Device, Variant> : public OpKernel {
public:
explicit AddNOp(OpKernelConstruction* context) : OpKernel(context) {}
void Compute(OpKernelContext* ctx) override {
auto binary_add = [](OpKernelContext* cc_ctx, const Variant& a,
const Variant& b, Variant* out) {
return BinaryOpVariants<Device>(cc_ctx, ADD_VARIANT_BINARY_OP, a, b, out);
};
AddNVariant(ctx, binary_add);
}
private:
// AddVariantTo efficiently performs:
// temp[lhs_ix] <- array(lhs_ix) + array(rhs_ix)
// where array(ix) := (temp_filled[ix]
// ? temp[ix]
// : ctx->input(ix).scalar<Variant>()())
// This reduces (possibly expensive) copying of Variants from
// the inputs into temp at the lowest levels of the summation tree.
static inline absl::Status AddVariantTo(
OpKernelContext* ctx, const int lhs_ix, const int rhs_ix,
absl::InlinedVector<Variant, 4UL>* temp,
absl::InlinedVector<bool, 4UL>* temp_filled) {
Variant tmp;
if (temp_filled->at(lhs_ix)) tmp = std::move(temp->at(lhs_ix));
const Variant& a = temp_filled->at(lhs_ix)
? tmp
: ctx->input(lhs_ix).template scalar<Variant>()();
const Variant& b = temp_filled->at(rhs_ix)
? temp->at(rhs_ix)
: ctx->input(rhs_ix).template scalar<Variant>()();
Variant* c = &temp->at(lhs_ix);
TF_RETURN_IF_ERROR(
BinaryOpVariants<Device>(ctx, ADD_VARIANT_BINARY_OP, a, b, c));
temp_filled->at(lhs_ix) = true;
return absl::OkStatus();
}
};
#define REGISTER_ADDN(type, dev) \
REGISTER_KERNEL_BUILDER( \
Name("AddN").Device(DEVICE_##dev).TypeConstraint<type>("T"), \
AddNOp<dev##Device, type>)
#define REGISTER_ADDN_CPU(type) REGISTER_ADDN(type, CPU)
TF_CALL_NUMBER_TYPES(REGISTER_ADDN_CPU);
REGISTER_ADDN_CPU(Variant);
#undef REGISTER_ADDN_CPU
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
#define REGISTER_ADDN_GPU(type) REGISTER_ADDN(type, GPU)
TF_CALL_int64(REGISTER_ADDN_GPU);
TF_CALL_uint32(REGISTER_ADDN_GPU);
TF_CALL_variant(REGISTER_ADDN_GPU);
TF_CALL_GPU_NUMBER_TYPES(REGISTER_ADDN_GPU);
TF_CALL_COMPLEX_TYPES(REGISTER_ADDN_GPU);
#undef REGISTER_ADDN_GPU
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
// A special DEVICE_DEFAULT kernel for int32.
// TODO(b/25387198): Also enable int32 in device memory. This kernel
// registration requires all int32 inputs and outputs to be in host memory.
REGISTER_KERNEL_BUILDER(Name("AddN")
.Device(DEVICE_DEFAULT)
.TypeConstraint<int32_t>("T")
.HostMemory("inputs")
.HostMemory("sum"),
AddNOp<CPUDevice, int32_t>);
#undef REGISTER_ADDN
} // namespace tensorflow
tensorflow/core/kernels/aggregate_ops.h
+226
View File
@@ -0,0 +1,226 @@
/* Copyright 2015 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_AGGREGATE_OPS_H_
#define TENSORFLOW_CORE_KERNELS_AGGREGATE_OPS_H_
#include <numeric>
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/tensor_types.h"
namespace tensorflow {
namespace functor {
// Functor definitions for Aggregate ops, must be compilable by nvcc.
template <typename Device, typename T>
struct Add2Functor {
void operator()(const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2);
};
template <typename Device, typename T>
struct Add2EigenImpl {
static void Compute(const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2) {
out.device(d) = in1 + in2;
}
};
template <typename Device, typename T>
struct Add3Functor {
void operator()(const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3);
};
template <typename Device, typename T>
struct Add3EigenImpl {
static void Compute(const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3) {
out.device(d) = in1 + in2 + in3;
}
};
template <typename Device, typename T>
struct Add4Functor {
void operator()(const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4);
};
template <typename Device, typename T>
struct Add4EigenImpl {
static void Compute(const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4) {
out.device(d) = in1 + in2 + in3 + in4;
}
};
template <typename Device, typename T>
struct Add5Functor {
void operator()(const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5);
};
template <typename Device, typename T>
struct Add5EigenImpl {
static void Compute(const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5) {
out.device(d) = in1 + in2 + in3 + in4 + in5;
}
};
template <typename Device, typename T>
struct Add6Functor {
void operator()(const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5,
typename TTypes<T>::ConstFlat in6);
};
template <typename Device, typename T>
struct Add6EigenImpl {
static void Compute(const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5,
typename TTypes<T>::ConstFlat in6) {
out.device(d) = in1 + in2 + in3 + in4 + in5 + in6;
}
};
template <typename Device, typename T>
struct Add7Functor {
void operator()(const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5,
typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7);
};
template <typename Device, typename T>
struct Add7EigenImpl {
static void Compute(const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5,
typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7) {
out.device(d) = in1 + in2 + in3 + in4 + in5 + in6 + in7;
}
};
template <typename Device, typename T>
struct Add8Functor {
void operator()(
const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1, typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3, typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5, typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7, typename TTypes<T>::ConstFlat in8);
};
template <typename Device, typename T>
struct Add8EigenImpl {
static void Compute(
const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1, typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3, typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5, typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7, typename TTypes<T>::ConstFlat in8) {
out.device(d) = in1 + in2 + in3 + in4 + in5 + in6 + in7 + in8;
}
};
// Add8p is like Add8 except the underlying implementation should +=
// rather than assign to the output.
template <typename Device, typename T>
struct Add8pFunctor {
void operator()(
const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1, typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3, typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5, typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7, typename TTypes<T>::ConstFlat in8);
};
template <typename Device, typename T>
struct Add8pEigenImpl {
static void Compute(
const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1, typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3, typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5, typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7, typename TTypes<T>::ConstFlat in8) {
out.device(d) += in1 + in2 + in3 + in4 + in5 + in6 + in7 + in8;
}
};
template <typename Device, typename T>
struct Add9Functor {
void operator()(
const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1, typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3, typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5, typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7, typename TTypes<T>::ConstFlat in8,
typename TTypes<T>::ConstFlat in9);
};
template <typename Device, typename T>
struct Add9EigenImpl {
static void Compute(
const Device& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1, typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3, typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5, typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7, typename TTypes<T>::ConstFlat in8,
typename TTypes<T>::ConstFlat in9) {
out.device(d) = in1 + in2 + in3 + in4 + in5 + in6 + in7 + in8 + in9;
}
};
} // namespace functor
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_AGGREGATE_OPS_H_
tensorflow/core/kernels/aggregate_ops_cpu.h
+142
View File
@@ -0,0 +1,142 @@
/* Copyright 2015 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_AGGREGATE_OPS_CPU_H_
#define TENSORFLOW_CORE_KERNELS_AGGREGATE_OPS_CPU_H_
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/kernels/aggregate_ops.h"
typedef Eigen::ThreadPoolDevice CPUDevice;
namespace tensorflow {
// Partial specializations for a CPUDevice, that uses the Eigen implementation
// from AddNEigenImpl.
namespace functor {
template <typename T>
struct Add2Functor<CPUDevice, T> {
void operator()(const CPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2) {
Add2EigenImpl<CPUDevice, T>::Compute(d, out, in1, in2);
}
};
template <typename T>
struct Add3Functor<CPUDevice, T> {
void operator()(const CPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3) {
Add3EigenImpl<CPUDevice, T>::Compute(d, out, in1, in2, in3);
}
};
template <typename T>
struct Add4Functor<CPUDevice, T> {
void operator()(const CPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4) {
Add4EigenImpl<CPUDevice, T>::Compute(d, out, in1, in2, in3, in4);
}
};
template <typename T>
struct Add5Functor<CPUDevice, T> {
void operator()(const CPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5) {
Add5EigenImpl<CPUDevice, T>::Compute(d, out, in1, in2, in3, in4, in5);
}
};
template <typename T>
struct Add6Functor<CPUDevice, T> {
void operator()(const CPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5,
typename TTypes<T>::ConstFlat in6) {
Add6EigenImpl<CPUDevice, T>::Compute(d, out, in1, in2, in3, in4, in5, in6);
}
};
template <typename T>
struct Add7Functor<CPUDevice, T> {
void operator()(const CPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5,
typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7) {
Add7EigenImpl<CPUDevice, T>::Compute(d, out, in1, in2, in3, in4, in5, in6,
in7);
}
};
template <typename T>
struct Add8Functor<CPUDevice, T> {
void operator()(
const CPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1, typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3, typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5, typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7, typename TTypes<T>::ConstFlat in8) {
Add8EigenImpl<CPUDevice, T>::Compute(d, out, in1, in2, in3, in4, in5, in6,
in7, in8);
}
};
template <typename T>
struct Add8pFunctor<CPUDevice, T> {
void operator()(
const CPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1, typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3, typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5, typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7, typename TTypes<T>::ConstFlat in8) {
Add8pEigenImpl<CPUDevice, T>::Compute(d, out, in1, in2, in3, in4, in5, in6,
in7, in8);
}
};
template <typename T>
struct Add9Functor<CPUDevice, T> {
void operator()(
const CPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1, typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3, typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5, typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7, typename TTypes<T>::ConstFlat in8,
typename TTypes<T>::ConstFlat in9) {
Add9EigenImpl<CPUDevice, T>::Compute(d, out, in1, in2, in3, in4, in5, in6,
in7, in8, in9);
}
};
} // namespace functor
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_AGGREGATE_OPS_CPU_H_
tensorflow/core/kernels/aggregate_ops_gpu.cu.cc
+165
View File
@@ -0,0 +1,165 @@
/* Copyright 2015 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.
==============================================================================*/
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
#define EIGEN_USE_GPU
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/kernels/aggregate_ops.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
typedef Eigen::GpuDevice GPUDevice;
// Partial specialization for a GPUDevice, that uses the Eigen implementation.
namespace functor {
template <typename T>
struct Add2Functor<GPUDevice, T> {
void operator()(const GPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2) {
Add2EigenImpl<GPUDevice, T>::Compute(d, out, in1, in2);
}
};
template <typename T>
struct Add3Functor<GPUDevice, T> {
void operator()(const GPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3) {
Add3EigenImpl<GPUDevice, T>::Compute(d, out, in1, in2, in3);
}
};
template <typename T>
struct Add4Functor<GPUDevice, T> {
void operator()(const GPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4) {
Add4EigenImpl<GPUDevice, T>::Compute(d, out, in1, in2, in3, in4);
}
};
template <typename T>
struct Add5Functor<GPUDevice, T> {
void operator()(const GPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5) {
Add5EigenImpl<GPUDevice, T>::Compute(d, out, in1, in2, in3, in4, in5);
}
};
template <typename T>
struct Add6Functor<GPUDevice, T> {
void operator()(const GPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5,
typename TTypes<T>::ConstFlat in6) {
Add6EigenImpl<GPUDevice, T>::Compute(d, out, in1, in2, in3, in4, in5, in6);
}
};
template <typename T>
struct Add7Functor<GPUDevice, T> {
void operator()(const GPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1,
typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3,
typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5,
typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7) {
Add7EigenImpl<GPUDevice, T>::Compute(d, out, in1, in2, in3, in4, in5, in6,
in7);
}
};
template <typename T>
struct Add8Functor<GPUDevice, T> {
void operator()(
const GPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1, typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3, typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5, typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7, typename TTypes<T>::ConstFlat in8) {
Add8EigenImpl<GPUDevice, T>::Compute(d, out, in1, in2, in3, in4, in5, in6,
in7, in8);
}
};
template <typename T>
struct Add8pFunctor<GPUDevice, T> {
void operator()(
const GPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1, typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3, typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5, typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7, typename TTypes<T>::ConstFlat in8) {
Add8pEigenImpl<GPUDevice, T>::Compute(d, out, in1, in2, in3, in4, in5, in6,
in7, in8);
}
};
template <typename T>
struct Add9Functor<GPUDevice, T> {
void operator()(
const GPUDevice& d, typename TTypes<T>::Flat out,
typename TTypes<T>::ConstFlat in1, typename TTypes<T>::ConstFlat in2,
typename TTypes<T>::ConstFlat in3, typename TTypes<T>::ConstFlat in4,
typename TTypes<T>::ConstFlat in5, typename TTypes<T>::ConstFlat in6,
typename TTypes<T>::ConstFlat in7, typename TTypes<T>::ConstFlat in8,
typename TTypes<T>::ConstFlat in9) {
Add9EigenImpl<GPUDevice, T>::Compute(d, out, in1, in2, in3, in4, in5, in6,
in7, in8, in9);
}
};
} // end namespace functor
// Instantiate the GPU implementation for GPU number types.
#define REGISTER_FUNCTORS(type) \
template struct functor::Add2Functor<GPUDevice, type>; \
template struct functor::Add3Functor<GPUDevice, type>; \
template struct functor::Add4Functor<GPUDevice, type>; \
template struct functor::Add5Functor<GPUDevice, type>; \
template struct functor::Add6Functor<GPUDevice, type>; \
template struct functor::Add7Functor<GPUDevice, type>; \
template struct functor::Add8Functor<GPUDevice, type>; \
template struct functor::Add8pFunctor<GPUDevice, type>; \
template struct functor::Add9Functor<GPUDevice, type>;
TF_CALL_int64(REGISTER_FUNCTORS);
TF_CALL_uint32(REGISTER_FUNCTORS);
TF_CALL_GPU_NUMBER_TYPES(REGISTER_FUNCTORS);
TF_CALL_COMPLEX_TYPES(REGISTER_FUNCTORS);
#undef REGISTER_FUNCTORS
} // end namespace tensorflow
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
tensorflow/core/kernels/argmax_op.cc
+261
View File
@@ -0,0 +1,261 @@
/* Copyright 2015 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.
==============================================================================*/
// See docs in ../ops/math_ops.cc.
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#define EIGEN_USE_THREADS
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
#define EIGEN_USE_GPU
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/bounds_check.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/kernels/argmax_op.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/macros.h"
namespace tensorflow {
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
template <typename Device, typename T, typename Tout, typename ArgFunctor>
class ArgOp : public OpKernel {
public:
explicit ArgOp(OpKernelConstruction* context) : OpKernel(context) {}
void Compute(OpKernelContext* context) override {
const Tensor& input = context->input(0);
const Tensor& dimension = context->input(1);
OP_REQUIRES(context, TensorShapeUtils::IsScalar(dimension.shape()),
absl::InvalidArgumentError(absl::StrCat(
"dim must be a scalar, but received tensor of shape: ",
dimension.shape().DebugString())));
const int32_t dim = internal::SubtleMustCopy(dimension.scalar<int32_t>()());
const int input_dims = input.dims();
int axis = dim < 0 ? dim + input_dims : dim;
OP_REQUIRES(context, FastBoundsCheck(axis, input_dims),
absl::InvalidArgumentError(absl::StrCat(
"Expected dimension in the range [", -input_dims, ", ",
input_dims, "), but got ", dim)));
OP_REQUIRES(context, input.dim_size(axis) > 0,
absl::InvalidArgumentError(
absl::StrCat("Reduction axis ", dim, " is empty in shape ",
input.shape().DebugString())));
TensorShape output_shape;
const TensorShape& input_shape = input.shape();
for (int d = 0; d < input_dims - 1; ++d) {
OP_REQUIRES_OK(context,
output_shape.AddDimWithStatus(
input_shape.dim_size((d < axis) ? d : d + 1)));
}
Tensor* output = nullptr;
OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output));
if (output_shape.num_elements() == 0) {
return;
}
#define HANDLE_DIM(NDIM) \
case NDIM: \
ArgFunctor::Reduce##NDIM(context->eigen_device<Device>(), \
input.tensor<T, NDIM>(), axis, \
output->tensor<Tout, NDIM - 1>()); \
break;
switch (input_dims) {
HANDLE_DIM(1);
HANDLE_DIM(2);
HANDLE_DIM(3);
HANDLE_DIM(4);
HANDLE_DIM(5);
HANDLE_DIM(6);
HANDLE_DIM(7);
default:
OP_REQUIRES(
context, false,
absl::InvalidArgumentError(absl::StrCat(
"Argmax and Argmin only support up "
"to 7 input dimensions, but got ",
input_dims, ". Inputs shape: ", input.shape().DebugString())));
}
}
#undef HANDLE_DIM
private:
ArgOp(const ArgOp&) = delete;
void operator=(const ArgOp&) = delete;
};
template <typename Device, typename T, typename Tout>
class ArgMaxOp
: public ArgOp<Device, T, Tout, functor::ArgMax<Device, T, Tout> > {
public:
explicit ArgMaxOp(OpKernelConstruction* context)
: ArgOp<Device, T, Tout, functor::ArgMax<Device, T, Tout> >(context) {}
};
template <typename Device, typename T, typename Tout>
class ArgMinOp
: public ArgOp<Device, T, Tout, functor::ArgMin<Device, T, Tout> > {
public:
explicit ArgMinOp(OpKernelConstruction* context)
: ArgOp<Device, T, Tout, functor::ArgMin<Device, T, Tout> >(context) {}
};
#define REGISTER_ARGMAX(type) \
REGISTER_KERNEL_BUILDER(Name("ArgMax") \
.Device(DEVICE_CPU) \
.TypeConstraint<type>("T") \
.TypeConstraint<int64_t>("output_type") \
.HostMemory("dimension"), \
ArgMaxOp<CPUDevice, type, int64>); \
REGISTER_KERNEL_BUILDER(Name("ArgMin") \
.Device(DEVICE_CPU) \
.TypeConstraint<type>("T") \
.TypeConstraint<int64_t>("output_type") \
.HostMemory("dimension"), \
ArgMinOp<CPUDevice, type, int64>); \
REGISTER_KERNEL_BUILDER(Name("ArgMax") \
.Device(DEVICE_CPU) \
.TypeConstraint<type>("T") \
.TypeConstraint<int32>("output_type") \
.HostMemory("dimension"), \
ArgMaxOp<CPUDevice, type, int32>); \
REGISTER_KERNEL_BUILDER(Name("ArgMin") \
.Device(DEVICE_CPU) \
.TypeConstraint<type>("T") \
.TypeConstraint<int32>("output_type") \
.HostMemory("dimension"), \
ArgMinOp<CPUDevice, type, int32>); \
REGISTER_KERNEL_BUILDER(Name("ArgMax") \
.Device(DEVICE_CPU) \
.TypeConstraint<type>("T") \
.TypeConstraint<int16>("output_type") \
.HostMemory("dimension"), \
ArgMaxOp<CPUDevice, type, int16>); \
REGISTER_KERNEL_BUILDER(Name("ArgMax") \
.Device(DEVICE_CPU) \
.TypeConstraint<type>("T") \
.TypeConstraint<uint16>("output_type") \
.HostMemory("dimension"), \
ArgMaxOp<CPUDevice, type, uint16>);
TF_CALL_REAL_NUMBER_TYPES(REGISTER_ARGMAX);
TF_CALL_bool(REGISTER_ARGMAX);
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
// Forward declarations of the functor specializations for GPU.
namespace functor {
#define DECLARE_GPU_SPEC(T, Tout, Dims) \
template <> \
void ArgMax<GPUDevice, T, Tout>::Reduce##Dims( \
const GPUDevice& d, typename TTypes<T, Dims>::ConstTensor input, \
const int32 dimension, typename TTypes<Tout, Dims - 1>::Tensor output); \
template <> \
void ArgMin<GPUDevice, T, Tout>::Reduce##Dims( \
const GPUDevice& d, typename TTypes<T, Dims>::ConstTensor input, \
const int32 dimension, typename TTypes<Tout, Dims - 1>::Tensor output);
#define DECLARE_GPU_SPECS(T) \
DECLARE_GPU_SPEC(T, int64_t, 1); \
DECLARE_GPU_SPEC(T, int64_t, 2); \
DECLARE_GPU_SPEC(T, int64_t, 3); \
DECLARE_GPU_SPEC(T, int64_t, 4); \
DECLARE_GPU_SPEC(T, int64_t, 5); \
DECLARE_GPU_SPEC(T, int64_t, 6); \
DECLARE_GPU_SPEC(T, int64_t, 7); \
DECLARE_GPU_SPEC(T, int32, 1); \
DECLARE_GPU_SPEC(T, int32, 2); \
DECLARE_GPU_SPEC(T, int32, 3); \
DECLARE_GPU_SPEC(T, int32, 4); \
DECLARE_GPU_SPEC(T, int32, 5); \
DECLARE_GPU_SPEC(T, int32, 6); \
DECLARE_GPU_SPEC(T, int32, 7);
#define DECLARE_GPU_CLASS(T) \
extern template struct ArgMax<GPUDevice, T, int64_t>; \
extern template struct ArgMin<GPUDevice, T, int64_t>; \
extern template struct ArgMax<GPUDevice, T, int32>; \
extern template struct ArgMin<GPUDevice, T, int32>;
TF_CALL_GPU_NUMBER_TYPES(DECLARE_GPU_SPECS);
TF_CALL_bool(DECLARE_GPU_SPECS);
TF_CALL_GPU_NUMBER_TYPES(DECLARE_GPU_CLASS);
TF_CALL_bool(DECLARE_GPU_CLASS);
#undef DECLARE_GPU_SPECS
#undef DECLARE_GPU_CLASS
} // namespace functor
// Registration of the GPU implementations.
#define REGISTER_ARGMAX_GPU(type) \
REGISTER_KERNEL_BUILDER(Name("ArgMax") \
.Device(DEVICE_GPU) \
.TypeConstraint<type>("T") \
.TypeConstraint<int64_t>("output_type") \
.TypeConstraint<int32>("Tidx") \
.HostMemory("dimension"), \
ArgMaxOp<GPUDevice, type, int64>); \
REGISTER_KERNEL_BUILDER(Name("ArgMin") \
.Device(DEVICE_GPU) \
.TypeConstraint<type>("T") \
.TypeConstraint<int64_t>("output_type") \
.TypeConstraint<int32>("Tidx") \
.HostMemory("dimension"), \
ArgMinOp<GPUDevice, type, int64>); \
REGISTER_KERNEL_BUILDER(Name("ArgMax") \
.Device(DEVICE_GPU) \
.TypeConstraint<type>("T") \
.TypeConstraint<int32>("output_type") \
.TypeConstraint<int32>("Tidx") \
.HostMemory("dimension"), \
ArgMaxOp<GPUDevice, type, int32>); \
REGISTER_KERNEL_BUILDER(Name("ArgMin") \
.Device(DEVICE_GPU) \
.TypeConstraint<type>("T") \
.TypeConstraint<int32>("output_type") \
.TypeConstraint<int32>("Tidx") \
.HostMemory("dimension"), \
ArgMinOp<GPUDevice, type, int32>);
TF_CALL_GPU_NUMBER_TYPES(REGISTER_ARGMAX_GPU);
TF_CALL_bool(REGISTER_ARGMAX_GPU);
#undef REGISTER_ARGMAX_GPU
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow
tensorflow/core/kernels/argmax_op.h
+72
View File
@@ -0,0 +1,72 @@
/* Copyright 2015 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_ARGMAX_OP_H_
#define TENSORFLOW_CORE_KERNELS_ARGMAX_OP_H_
// Generator definition for ArgMaxOp, must be compilable by nvcc.
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace functor {
template <typename Device, typename T, typename Tout>
struct ArgMax {
#define DECLARE_COMPUTE_SPEC(Dims) \
EIGEN_ALWAYS_INLINE static void Reduce##Dims( \
const Device& d, typename TTypes<T, Dims>::ConstTensor input, \
const int32 dimension, typename TTypes<Tout, Dims - 1>::Tensor output) { \
output.device(d) = input.argmax(dimension).template cast<Tout>(); \
}
DECLARE_COMPUTE_SPEC(1);
DECLARE_COMPUTE_SPEC(2);
DECLARE_COMPUTE_SPEC(3);
DECLARE_COMPUTE_SPEC(4);
DECLARE_COMPUTE_SPEC(5);
DECLARE_COMPUTE_SPEC(6);
DECLARE_COMPUTE_SPEC(7);
#undef DECLARE_COMPUTE_SPEC
};
template <typename Device, typename T, typename Tout>
struct ArgMin {
#define DECLARE_COMPUTE_SPEC(Dims) \
EIGEN_ALWAYS_INLINE static void Reduce##Dims( \
const Device& d, typename TTypes<T, Dims>::ConstTensor input, \
const int32 dimension, typename TTypes<Tout, Dims - 1>::Tensor output) { \
output.device(d) = input.argmin(dimension).template cast<Tout>(); \
}
DECLARE_COMPUTE_SPEC(1);
DECLARE_COMPUTE_SPEC(2);
DECLARE_COMPUTE_SPEC(3);
DECLARE_COMPUTE_SPEC(4);
DECLARE_COMPUTE_SPEC(5);
DECLARE_COMPUTE_SPEC(6);
DECLARE_COMPUTE_SPEC(7);
#undef DECLARE_COMPUTE_SPEC
};
} // namespace functor
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_ARGMAX_OP_H_
tensorflow/core/kernels/argmax_op_gpu.cu.cc
+39
View File
@@ -0,0 +1,39 @@
/* Copyright 2015 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.
==============================================================================*/
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
#define EIGEN_USE_GPU
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/kernels/argmax_op.h"
namespace tensorflow {
typedef Eigen::GpuDevice GPUDevice;
#define DEFINE_GPU_SPEC(T) \
template struct functor::ArgMax<GPUDevice, T, int64>; \
template struct functor::ArgMin<GPUDevice, T, int64>; \
template struct functor::ArgMax<GPUDevice, T, int32>; \
template struct functor::ArgMin<GPUDevice, T, int32>;
TF_CALL_GPU_NUMBER_TYPES(DEFINE_GPU_SPEC);
TF_CALL_bool(DEFINE_GPU_SPEC);
} // end namespace tensorflow
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
tensorflow/core/kernels/as_string_op.cc
+293
View File
@@ -0,0 +1,293 @@
/* Copyright 2016 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.
==============================================================================*/
// See docs in ../ops/string_ops.cc.
#include <cstdint>
#include <string>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "Eigen/Core" // from @eigen_archive
#include "tensorflow/core/framework/kernel_def_builder.h"
#include "tensorflow/core/framework/numeric_types.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/op_requires.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/framework/variant.h"
#include "tensorflow/core/framework/variant_encode_decode.h"
#include "tensorflow/core/framework/variant_tensor_data.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/bfloat16.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/platform/tstring.h"
namespace tensorflow {
class AsStringOp : public OpKernel {
public:
using OpKernel::OpKernel;
explicit AsStringOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
bool scientific;
bool shortest;
std::string fill_string;
DataType dtype;
OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype));
OP_REQUIRES_OK(ctx, ctx->GetAttr("precision", &precision_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("scientific", &scientific));
OP_REQUIRES_OK(ctx, ctx->GetAttr("shortest", &shortest));
OP_REQUIRES_OK(ctx, ctx->GetAttr("width", &width_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("fill", &fill_string));
if (dtype != DT_STRING && !DataTypeIsFloating(dtype) &&
!DataTypeIsComplex(dtype)) {
OP_REQUIRES(ctx, !(scientific || shortest),
absl::InvalidArgumentError(
absl::StrCat("scientific and shortest format "
"not supported for datatype ",
DataTypeString(dtype))));
OP_REQUIRES(
ctx, precision_ < 0,
absl::InvalidArgumentError(absl::StrCat("precision not supported "
"for datatype ",
DataTypeString(dtype))));
}
OP_REQUIRES(ctx, fill_string.size() <= 1,
absl::InvalidArgumentError(
"Fill string must be one or fewer characters"));
OP_REQUIRES(ctx, !(scientific && shortest),
absl::InvalidArgumentError(
"Cannot select both scientific and shortest notation"));
if (!fill_string.empty()) {
switch (fill_string[0]) {
case ' ':
case '+':
case '-':
case '0':
case '#':
break;
default:
bool fill_not_supported = true;
OP_REQUIRES(
ctx, !fill_not_supported,
absl::InvalidArgumentError(absl::StrCat(
"Fill argument not supported: \"", fill_string, "\"")));
}
}
if (width_ <= -1) {
width_ = 0;
}
// If input is string and width unspecified, simply forward to output.
if (dtype == DT_STRING && width_ <= 0) {
return;
}
char format_char;
if (dtype == DT_STRING) {
format_char = 's';
} else if (DataTypeIsUnsigned(dtype)) {
format_char = 'u';
} else if (DataTypeIsSigned(dtype)) {
format_char = 'd';
} else if (DataTypeIsFloating(dtype) || DataTypeIsComplex(dtype)) {
if (shortest) {
format_char = 'g';
} else if (scientific) {
format_char = 'e';
} else {
format_char = 'f';
}
} else if (dtype == DT_BOOL) {
return;
} else if (dtype == DT_VARIANT) {
return;
} else {
bool type_not_supported = true;
OP_REQUIRES(ctx, !type_not_supported,
absl::InvalidArgumentError(absl::StrCat(
"Type not supported: ", DataTypeString(dtype))));
}
format_ = absl::StrCat("%", fill_string, "*.*",
absl::string_view(&format_char, 1));
if (format_char == 's') {
string_format_ = StringFormat::New(format_);
OP_REQUIRES(ctx, string_format_ != nullptr,
absl::InvalidArgumentError(
absl::StrCat("Invalid format: ", format_)));
} else if (format_char == 'u' || format_char == 'd') {
integral_format_ = IntegralFormat::New(format_);
OP_REQUIRES(ctx, integral_format_ != nullptr,
absl::InvalidArgumentError(
absl::StrCat("Invalid format: ", format_)));
} else {
floating_format_ = FloatingFormat::New(format_);
OP_REQUIRES(ctx, floating_format_ != nullptr,
absl::InvalidArgumentError(
absl::StrCat("Invalid format: ", format_)));
}
}
void Compute(OpKernelContext* context) override {
const Tensor* input_tensor;
OP_REQUIRES_OK(context, context->input("input", &input_tensor));
const DataType& dtype = input_tensor->dtype();
// If input is string and width unspecified, simply forward to output.
if (dtype == DT_STRING && width_ <= 0) {
context->set_output(0, context->input(0));
return;
}
Tensor* output_tensor = nullptr;
OP_REQUIRES_OK(context,
context->allocate_output("output", input_tensor->shape(),
&output_tensor));
auto output_flat = output_tensor->flat<tstring>();
if (dtype == DT_BOOL) {
const auto& input_flat = input_tensor->flat<bool>();
for (int64_t i = 0; i < input_flat.size(); ++i) {
output_flat(i) = (input_flat(i)) ? "true" : "false";
}
return;
}
if (dtype == DT_VARIANT) {
const auto& input_flat = input_tensor->flat<Variant>();
for (int64_t i = 0; i < input_flat.size(); ++i) {
output_flat(i) = input_flat(i).DebugString();
}
return;
}
// All other cases use the format string.
#define ENCODE_TYPE(type, T, enc_fmt) \
case (type): { \
const auto& input_flat = input_tensor->flat<T>(); \
for (int64_t i = 0; i < input_flat.size(); ++i) { \
output_flat(i) = \
absl::StrFormat(*enc_fmt, width_, precision_, input_flat(i)); \
} \
} break
switch (dtype) {
ENCODE_TYPE(DT_UINT8, uint8_t, integral_format_);
ENCODE_TYPE(DT_UINT16, uint16_t, integral_format_);
ENCODE_TYPE(DT_UINT32, uint32_t, integral_format_);
ENCODE_TYPE(DT_UINT64, uint64_t, integral_format_);
ENCODE_TYPE(DT_INT8, int8_t, integral_format_);
ENCODE_TYPE(DT_INT16, int16_t, integral_format_);
ENCODE_TYPE(DT_INT32, int32_t, integral_format_);
ENCODE_TYPE(DT_INT64, int64_t, integral_format_);
ENCODE_TYPE(DT_FLOAT, float, floating_format_);
ENCODE_TYPE(DT_DOUBLE, double, floating_format_);
case (DT_STRING): {
const auto& input_flat = input_tensor->flat<tstring>();
for (int64_t i = 0; i < input_flat.size(); ++i) {
output_flat(i) = absl::StrFormat(*string_format_, width_, precision_,
absl::string_view(input_flat(i)));
}
} break;
case (DT_HALF): {
const auto& input_flat = input_tensor->flat<Eigen::half>();
for (int64_t i = 0; i < input_flat.size(); ++i) {
output_flat(i) =
absl::StrFormat(*floating_format_, width_, precision_,
static_cast<float>(input_flat(i)));
}
} break;
case (DT_BFLOAT16): {
const auto& input_flat = input_tensor->flat<bfloat16>();
for (int64_t i = 0; i < input_flat.size(); ++i) {
output_flat(i) =
absl::StrFormat(*floating_format_, width_, precision_,
static_cast<float>(input_flat(i)));
}
} break;
case (DT_COMPLEX64): {
const auto& input_flat = input_tensor->flat<complex64>();
for (int64_t i = 0; i < input_flat.size(); ++i) {
output_flat(i) =
absl::StrCat("(",
absl::StrFormat(*floating_format_, width_,
precision_, input_flat(i).real()),
",",
absl::StrFormat(*floating_format_, width_,
precision_, input_flat(i).imag()),
")");
}
} break;
case (DT_COMPLEX128): {
const auto& input_flat = input_tensor->flat<complex128>();
for (int64_t i = 0; i < input_flat.size(); ++i) {
output_flat(i) =
absl::StrCat("(",
absl::StrFormat(*floating_format_, width_,
precision_, input_flat(i).real()),
",",
absl::StrFormat(*floating_format_, width_,
precision_, input_flat(i).imag()),
")");
}
} break;
default:
bool can_encode_type = false;
OP_REQUIRES(
context, can_encode_type,
absl::InvalidArgumentError(absl::StrCat(
"Cannot encode input of type ", DataTypeString(dtype))));
}
#undef ENCODE_TYPE
}
private:
// Used to parse "%*.*g", etc.
using FloatingFormat =
absl::ParsedFormat<absl::FormatConversionCharSet::kStar,
absl::FormatConversionCharSet::kStar,
absl::FormatConversionCharSet::g |
absl::FormatConversionCharSet::e |
absl::FormatConversionCharSet::f>;
// Used to parse "%*.*u", etc.
using IntegralFormat =
absl::ParsedFormat<absl::FormatConversionCharSet::kStar,
absl::FormatConversionCharSet::kStar,
absl::FormatConversionCharSet::u |
absl::FormatConversionCharSet::d>;
// Used to parse "%*.*s", etc.
using StringFormat = absl::ParsedFormat<absl::FormatConversionCharSet::kStar,
absl::FormatConversionCharSet::kStar,
absl::FormatConversionCharSet::s>;
int precision_ = -1;
int width_ = -1;
decltype(StringFormat::New("%*.*s")) string_format_;
decltype(IntegralFormat::New("%*.*u")) integral_format_;
decltype(FloatingFormat::New("%*.*g")) floating_format_;
std::string format_;
};
REGISTER_KERNEL_BUILDER(Name("AsString").Device(DEVICE_CPU), AsStringOp);
} // namespace tensorflow
tensorflow/core/kernels/as_string_op_test.cc
+370
View File
@@ -0,0 +1,370 @@
/* 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 <cstdint>
#include "Eigen/Core" // from @eigen_archive // IWYU pragma: keep
#include "tensorflow/core/framework/fake_input.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/numeric_types.h" // IWYU pragma: keep
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/variant.h"
#include "tensorflow/core/framework/variant_encode_decode.h"
#include "tensorflow/core/framework/variant_tensor_data.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/bfloat16.h" // IWYU pragma: keep
#include "tsl/platform/bfloat16.h" // IWYU pragma: keep
namespace tensorflow {
namespace {
class AsStringGraphTest : public OpsTestBase {
protected:
absl::Status Init(DataType input_type, const std::string& fill = "",
int width = -1, int precision = -1, bool scientific = false,
bool shortest = false) {
TF_CHECK_OK(NodeDefBuilder("op", "AsString")
.Input(FakeInput(input_type))
.Attr("fill", fill)
.Attr("precision", precision)
.Attr("scientific", scientific)
.Attr("shortest", shortest)
.Attr("width", width)
.Finalize(node_def()));
return InitOp();
}
};
TEST_F(AsStringGraphTest, Int8) {
TF_ASSERT_OK(Init(DT_INT8));
AddInputFromList<int8_t, int8_t>(TensorShape({3}), {-42, 0, 42});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({3}));
test::FillValues<tstring>(&expected, {"-42", "0", "42"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, Int64) {
TF_ASSERT_OK(Init(DT_INT64));
AddInputFromList<int64_t, int64_t>(TensorShape({3}), {-42, 0, 42});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({3}));
test::FillValues<tstring>(&expected, {"-42", "0", "42"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, FloatDefault) {
TF_ASSERT_OK(Init(DT_FLOAT));
AddInputFromList<float, float>(TensorShape({4}), {-42, 0, 3.14159, 42});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({4}));
test::FillValues<tstring>(
&expected, {"-42.000000", "0.000000", "3.141590", "42.000000"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, FloatScientific) {
TF_ASSERT_OK(Init(DT_FLOAT, /*fill=*/"", /*width=*/-1, /*precision=*/-1,
/*scientific=*/true));
AddInputFromList<float, float>(TensorShape({4}), {-42, 0, 3.14159, 42});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({4}));
test::FillValues<tstring>(&expected, {"-4.200000e+01", "0.000000e+00",
"3.141590e+00", "4.200000e+01"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, FloatShortest) {
TF_ASSERT_OK(Init(DT_FLOAT, /*fill=*/"", /*width=*/-1, /*precision=*/-1,
/*scientific=*/false, /*shortest=*/true));
AddInputFromList<float, float>(TensorShape({4}), {-42, 0, 3.14159, 42});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({4}));
test::FillValues<tstring>(&expected, {"-42", "0", "3.14159", "42"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, FloatPrecisionOnly) {
TF_ASSERT_OK(Init(DT_FLOAT, /*fill=*/"", /*width=*/-1, /*precision=*/2));
AddInputFromList<float, float>(TensorShape({4}), {-42, 0, 3.14159, 42});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({4}));
test::FillValues<tstring>(&expected, {"-42.00", "0.00", "3.14", "42.00"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, FloatWidthOnly) {
TF_ASSERT_OK(Init(DT_FLOAT, /*fill=*/"", /*width=*/5));
AddInputFromList<float, float>(TensorShape({4}), {-42, 0, 3.14159, 42});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({4}));
test::FillValues<tstring>(
&expected, {"-42.000000", "0.000000", "3.141590", "42.000000"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, Float_5_2_Format) {
TF_ASSERT_OK(Init(DT_FLOAT, /*fill=*/"", /*width=*/5, /*precision=*/2));
AddInputFromList<float, float>(TensorShape({4}), {-42, 0, 3.14159, 42});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({4}));
test::FillValues<tstring>(&expected, {"-42.00", " 0.00", " 3.14", "42.00"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, Complex) {
TF_ASSERT_OK(Init(DT_COMPLEX64, /*fill=*/"", /*width=*/5, /*precision=*/2));
AddInputFromList<complex64, complex64>(TensorShape({3}),
{{-4, 2}, {0}, {3.14159, -1}});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({3}));
test::FillValues<tstring>(
&expected, {"(-4.00, 2.00)", "( 0.00, 0.00)", "( 3.14,-1.00)"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, Bool) {
TF_ASSERT_OK(Init(DT_BOOL));
AddInputFromList<bool, bool>(TensorShape({2}), {true, false});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({2}));
test::FillValues<tstring>(&expected, {"true", "false"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, Variant) {
TF_ASSERT_OK(Init(DT_VARIANT));
AddInput(DT_VARIANT, TensorShape({4}));
auto inputs = mutable_input(0)->flat<Variant>();
inputs(0) = 2;
inputs(1) = 3;
inputs(2) = true;
inputs(3) = Tensor("hi");
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({4}));
test::FillValues<tstring>(
&expected, {"Variant<type: int value: 2>", "Variant<type: int value: 3>",
"Variant<type: bool value: 1>",
("Variant<type: tensorflow::Tensor value: Tensor<type: string"
" shape: [] values: hi>>")});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, OnlyOneOfScientificAndShortest) {
absl::Status s = Init(DT_FLOAT, /*fill=*/"", /*width=*/-1, /*precision=*/-1,
/*scientific=*/true, /*shortest=*/true);
ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
ASSERT_TRUE(absl::StrContains(
s.message(), "Cannot select both scientific and shortest notation"));
}
TEST_F(AsStringGraphTest, NoShortestForNonFloat) {
absl::Status s = Init(DT_INT32, /*fill=*/"", /*width=*/-1, /*precision=*/-1,
/*scientific=*/false, /*shortest=*/true);
ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
ASSERT_TRUE(absl::StrContains(
s.message(),
"scientific and shortest format not supported for datatype"));
}
TEST_F(AsStringGraphTest, NoScientificForNonFloat) {
absl::Status s = Init(DT_INT32, /*fill=*/"", /*width=*/-1, /*precision=*/-1,
/*scientific=*/true);
ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
ASSERT_TRUE(absl::StrContains(
s.message(),
"scientific and shortest format not supported for datatype"));
}
TEST_F(AsStringGraphTest, NoPrecisionForNonFloat) {
absl::Status s = Init(DT_INT32, /*fill=*/"", /*width=*/-1, /*precision=*/5);
ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
ASSERT_TRUE(
absl::StrContains(s.message(), "precision not supported for datatype"));
}
TEST_F(AsStringGraphTest, LongFill) {
absl::Status s = Init(DT_INT32, /*fill=*/"asdf");
ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
ASSERT_TRUE(absl::StrContains(s.message(),
"Fill string must be one or fewer characters"));
}
TEST_F(AsStringGraphTest, FillWithZero) {
TF_ASSERT_OK(Init(DT_INT64, /*fill=*/"0", /*width=*/4));
AddInputFromList<int64_t, int64_t>(TensorShape({3}), {-42, 0, 42});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({3}));
test::FillValues<tstring>(&expected, {"-042", "0000", "0042"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, FillWithSpace) {
TF_ASSERT_OK(Init(DT_INT64, /*fill=*/" ", /*width=*/4));
AddInputFromList<int64_t, int64_t>(TensorShape({3}), {-42, 0, 42});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({3}));
test::FillValues<tstring>(&expected, {" -42", " 0", " 42"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, FillWithChar1) {
TF_ASSERT_OK(Init(DT_INT64, /*fill=*/"-", /*width=*/4));
AddInputFromList<int64_t, int64_t>(TensorShape({3}), {-42, 0, 42});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({3}));
test::FillValues<tstring>(&expected, {"-42 ", "0 ", "42 "});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, FillWithChar3) {
absl::Status s = Init(DT_INT32, /*fill=*/"s");
ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
ASSERT_TRUE(absl::StrContains(s.message(), "Fill argument not supported"));
}
TEST_F(AsStringGraphTest, FillWithChar4) {
absl::Status s = Init(DT_INT32, /*fill=*/"n");
ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
ASSERT_TRUE(absl::StrContains(s.message(), "Fill argument not supported"));
}
TEST_F(AsStringGraphTest, Int16) {
TF_ASSERT_OK(Init(DT_INT16));
AddInputFromList<int16_t, int16_t>(TensorShape({2}), {-10, 10});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({2}));
test::FillValues<tstring>(&expected, {"-10", "10"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, Int32) {
TF_ASSERT_OK(Init(DT_INT32));
AddInputFromList<int32_t, int32_t>(TensorShape({2}), {-20, 20});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({2}));
test::FillValues<tstring>(&expected, {"-20", "20"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, Uint8) {
TF_ASSERT_OK(Init(DT_UINT8));
AddInputFromList<uint8_t, uint8_t>(TensorShape({2}), {0, 255});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({2}));
test::FillValues<tstring>(&expected, {"0", "255"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, Uint16) {
TF_ASSERT_OK(Init(DT_UINT16));
AddInputFromList<uint16_t, uint16_t>(TensorShape({2}), {0, 65535});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({2}));
test::FillValues<tstring>(&expected, {"0", "65535"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, Uint32) {
TF_ASSERT_OK(Init(DT_UINT32));
AddInputFromList<uint32_t, uint32_t>(TensorShape({2}), {0, 4294967295U});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({2}));
test::FillValues<tstring>(&expected, {"0", "4294967295"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, Uint64) {
TF_ASSERT_OK(Init(DT_UINT64));
AddInputFromList<uint64_t, uint64_t>(TensorShape({2}),
{0, 18446744073709551615ULL});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({2}));
test::FillValues<tstring>(&expected, {"0", "18446744073709551615"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, Double) {
TF_ASSERT_OK(Init(DT_DOUBLE, /*fill=*/"", /*width=*/-1, /*precision=*/2));
AddInputFromList<double, double>(TensorShape({2}), {-3.14159, 2.71828});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({2}));
test::FillValues<tstring>(&expected, {"-3.14", "2.72"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, Half) {
TF_ASSERT_OK(Init(DT_HALF, /*fill=*/"", /*width=*/-1, /*precision=*/2));
AddInputFromList<Eigen::half, Eigen::half>(
TensorShape({2}),
{static_cast<Eigen::half>(-1.5f), static_cast<Eigen::half>(1.5f)});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({2}));
test::FillValues<tstring>(&expected, {"-1.50", "1.50"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, Bfloat16) {
TF_ASSERT_OK(Init(DT_BFLOAT16, /*fill=*/"", /*width=*/-1, /*precision=*/2));
AddInputFromList<tsl::bfloat16, tsl::bfloat16>(
TensorShape({2}),
{static_cast<tsl::bfloat16>(-2.5f), // NOLINT(misc-include-cleaner)
static_cast<tsl::bfloat16>(2.5f)}); // NOLINT(misc-include-cleaner)
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({2}));
test::FillValues<tstring>(&expected, {"-2.50", "2.50"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, Complex128) {
TF_ASSERT_OK(Init(DT_COMPLEX128, /*fill=*/"", /*width=*/5, /*precision=*/2));
AddInputFromList<complex128, complex128>(TensorShape({2}),
{{-4, 2}, {3.14159, -1}});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({2}));
test::FillValues<tstring>(&expected, {"(-4.00, 2.00)", "( 3.14,-1.00)"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
TEST_F(AsStringGraphTest, StringWidth) {
TF_ASSERT_OK(Init(DT_STRING, /*fill=*/"", /*width=*/5));
AddInputFromList<tstring, tstring>(TensorShape({2}), {"abc", "de"});
TF_ASSERT_OK(RunOpKernel());
Tensor expected(allocator(), DT_STRING, TensorShape({2}));
test::FillValues<tstring>(&expected, {" abc", " de"});
test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}
} // end namespace
} // end namespace tensorflow
tensorflow/core/kernels/assign_op.h
+71
View File
@@ -0,0 +1,71 @@
/* Copyright 2015 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_ASSIGN_OP_H_
#define TENSORFLOW_CORE_KERNELS_ASSIGN_OP_H_
#define EIGEN_USE_THREADS
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/ref_var.h"
#include "tensorflow/core/framework/tensor_types.h"
namespace tensorflow {
// TODO(jeff): Get rid of use_exclusive_lock_ option
// Computes *input[0] = input[1]
class AssignOp : public OpKernel {
public:
explicit AssignOp(OpKernelConstruction* context) : OpKernel(context) {
OP_REQUIRES_OK(context,
context->GetAttr("use_locking", &use_exclusive_lock_));
OP_REQUIRES_OK(context,
context->GetAttr("validate_shape", &validate_shape_));
OP_REQUIRES(context, IsRefType(context->input_type(0)),
absl::InvalidArgumentError("lhs input needs to be a ref type"));
if (!context
->GetAttr("_grappler_relax_allocator_constraints",
&relax_constraints_)
.ok()) {
relax_constraints_ = false;
}
}
void Compute(OpKernelContext* context) override {
constexpr int input_ref_index = 0;
constexpr int output_ref_index = 0;
constexpr int value_index = 1;
auto copy = [this](OpKernelContext* cc_ctx, Tensor* lhs,
const Tensor& rhs) { Copy(cc_ctx, lhs, rhs); };
AssignRefVariable(context, input_ref_index, output_ref_index, value_index,
use_exclusive_lock_, validate_shape_, relax_constraints_,
copy);
}
virtual void Copy(OpKernelContext* context, Tensor* lhs,
const Tensor& rhs) = 0;
bool use_exclusive_lock_;
bool validate_shape_;
bool relax_constraints_;
};
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_ASSIGN_OP_H_
tensorflow/core/kernels/autotune_conv_impl.h
+97
View File
@@ -0,0 +1,97 @@
/* Copyright 2022 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.
------------------------------------------------------------------------------*/
#ifndef TENSORFLOW_CORE_KERNELS_AUTOTUNE_CONV_IMPL_H_
#define TENSORFLOW_CORE_KERNELS_AUTOTUNE_CONV_IMPL_H_
#if GOOGLE_CUDA
#define EIGEN_USE_THREADS
#include "xla/stream_executor/gpu/redzone_allocator.h"
#include "xla/stream_executor/integrations/tf_allocator_adapter.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/kernels/conv_ops_gpu.h"
#include "tensorflow/core/util/proto/proto_utils.h"
namespace tensorflow::internal {
template <typename LaunchFunc, typename Sig>
absl::StatusOr<std::vector<xla::AutotuneResult>> AutotuneConvImpl(
OpKernelContext* ctx,
std::vector<std::unique_ptr<const se::dnn::OpRunner<Sig>>>& runners,
bool actually_do_autotune, const LaunchFunc& launch_func,
size_t scratch_size_limit, const se::RedzoneAllocator& rz_allocator) {
auto* stream = ctx->op_device_context()->stream();
se::TfAllocatorAdapter tf_allocator_adapter(ctx->device()->GetAllocator({}),
stream);
std::vector<xla::AutotuneResult> results;
// TODO(reedwm): Warn if determinism is enabled after autotune is run
for (auto& runner : runners) {
// TODO(zhengxq): profile each algorithm multiple times to better
// accuracy.
se::RedzoneAllocator rz_scratch_allocator(
stream, &tf_allocator_adapter,
/*memory_limit=*/scratch_size_limit);
DnnScratchAllocator scratch_allocator(scratch_size_limit, ctx);
se::ScratchAllocator* allocator_used =
!RedzoneCheckDisabled()
? static_cast<se::ScratchAllocator*>(&rz_scratch_allocator)
: static_cast<se::ScratchAllocator*>(&scratch_allocator);
TF_ASSIGN_OR_RETURN(auto desc, runner->ToAlgorithmDesc());
se::dnn::ProfileResult profile_result;
absl::Status cudnn_launch_status =
actually_do_autotune
? launch_func(allocator_used, runner, &profile_result)
: absl::OkStatus();
if (!actually_do_autotune) {
// Make the result valid according to `is_valid`.
profile_result.set_algorithm(desc);
profile_result.set_elapsed_time_in_ms(0);
}
// We need to make sure the profiling results are one-to-one with the
// "runners". So, we insert dummy results when the execution fails.
results.emplace_back();
auto& result = results.back();
*result.mutable_algorithm() = desc.ToProto();
if (cudnn_launch_status.ok() && profile_result.is_valid()) {
result.set_scratch_bytes(
!RedzoneCheckDisabled()
? rz_scratch_allocator.TotalAllocatedBytesExcludingRedzones()
: scratch_allocator.TotalByteSize());
*result.mutable_run_time() = proto_utils::ToDurationProto(
absl::Milliseconds(profile_result.elapsed_time_in_ms()));
CheckRedzones(rz_scratch_allocator, &result);
CheckRedzones(rz_allocator, &result);
} else {
result.mutable_failure()->set_kind(xla::AutotuneResult::UNKNOWN);
result.mutable_failure()->set_msg(
absl::StrCat("Profiling failure on CUDNN engine ", desc.ToString(),
": ", cudnn_launch_status.ToString()));
}
}
return results;
}
} // namespace tensorflow::internal
#endif // GOOGLE_CUDA
#endif // TENSORFLOW_CORE_KERNELS_AUTOTUNE_CONV_IMPL_H_
tensorflow/core/kernels/avgpooling_op.cc
+671
View File
@@ -0,0 +1,671 @@
/* Copyright 2015 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.
==============================================================================*/
// See docs in ../ops/nn_ops.cc.
#define EIGEN_USE_THREADS
#include "tensorflow/core/kernels/avgpooling_op.h"
#include <vector>
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/kernel_shape_util.h"
#include "tensorflow/core/framework/numeric_op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_slice.h"
#include "tensorflow/core/kernels/eigen_pooling.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/kernels/pooling_ops_common.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/util/overflow.h"
#include "tensorflow/core/util/padding.h"
#include "tensorflow/core/util/tensor_format.h"
#if GOOGLE_CUDA
#include "third_party/gpus/cudnn/cudnn.h"
#endif // GOOGLE_CUDA
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#include "tensorflow/core/kernels/maxpooling_op_gpu.h"
#include "tensorflow/core/kernels/pooling_ops_common_gpu.h"
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
namespace tensorflow {
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
template <typename Device, typename T>
class AvgPoolingOp : public UnaryOp<T> {
public:
explicit AvgPoolingOp(OpKernelConstruction* context) : UnaryOp<T>(context) {
std::string data_format;
OP_REQUIRES_OK(context, context->GetAttr("data_format", &data_format));
OP_REQUIRES(context, FormatFromString(data_format, &data_format_),
absl::InvalidArgumentError("Invalid data format"));
OP_REQUIRES(
context, data_format_ == FORMAT_NHWC,
absl::InvalidArgumentError(absl::StrCat(
"Default AvgPoolingOp only supports NHWC ", "on device type ",
DeviceTypeString(context->device_type()))));
OP_REQUIRES_OK(context, context->GetAttr("ksize", &ksize_));
OP_REQUIRES(context, ksize_.size() == 4,
absl::InvalidArgumentError("Sliding window ksize field must "
"specify 4 dimensions"));
OP_REQUIRES_OK(context, context->GetAttr("strides", &stride_));
OP_REQUIRES(context, stride_.size() == 4,
absl::InvalidArgumentError("Sliding window stride field must "
"specify 4 dimensions"));
OP_REQUIRES_OK(context, context->GetAttr("padding", &padding_));
OP_REQUIRES(context, ksize_[0] == 1 && stride_[0] == 1,
absl::UnimplementedError(
"Pooling is not yet supported on the batch dimension."));
for (int i = 0; i < ksize_.size(); ++i) {
OP_REQUIRES(
context, ksize_[i] > 0,
absl::InvalidArgumentError(absl::StrCat(
"ksize must be a positive int32 value, got:", ksize_[i])));
}
}
void Compute(OpKernelContext* context) override {
const Tensor& tensor_in = context->input(0);
PoolParameters params{context,
ksize_,
stride_,
padding_,
/*explicit_paddings=*/{},
data_format_,
tensor_in.shape()};
if (!context->status().ok()) {
return;
}
OP_REQUIRES(context, params.depth_window == 1,
absl::UnimplementedError("Non-spatial pooling is not "
"yet supported. Volunteers? :)"));
// For avgpooling, tensor_in should have 4 dimensions.
OP_REQUIRES(context, tensor_in.dims() == 4,
absl::InvalidArgumentError("tensor_in must be 4-dimensional"));
Tensor* output = nullptr;
TensorShape params_forward_output_shape;
OP_REQUIRES_OK(context,
params.forward_output_shape(&params_forward_output_shape));
OP_REQUIRES_OK(context, context->allocate_output(
0, params_forward_output_shape, &output));
SpatialAvgPool<Device, T>(context, output, tensor_in, params, padding_);
}
private:
std::vector<int32_t> ksize_;
std::vector<int32_t> stride_;
Padding padding_;
TensorFormat data_format_;
};
#define REGISTER_CPU_KERNEL_AVGPOOL(T) \
REGISTER_KERNEL_BUILDER( \
Name("AvgPool").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
AvgPoolingOp<CPUDevice, T>);
TF_CALL_FLOAT_TYPES(REGISTER_CPU_KERNEL_AVGPOOL);
#undef REGISTER_CPU_KERNEL_AVGPOOL
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
template <typename T>
class AvgPoolingOp<GPUDevice, T> : public UnaryOp<T> {
public:
typedef GPUDevice Device;
explicit AvgPoolingOp(OpKernelConstruction* context) : UnaryOp<T>(context) {
std::string data_format;
OP_REQUIRES_OK(context, context->GetAttr("data_format", &data_format));
OP_REQUIRES(context, FormatFromString(data_format, &data_format_),
absl::InvalidArgumentError("Invalid data format"));
OP_REQUIRES_OK(context, context->GetAttr("ksize", &ksize_));
OP_REQUIRES(context, ksize_.size() == 4,
absl::InvalidArgumentError("Sliding window ksize field must "
"specify 4 dimensions"));
for (int i = 0; i < ksize_.size(); ++i) {
OP_REQUIRES(
context, ksize_[i] > 0,
absl::InvalidArgumentError(absl::StrCat(
"ksize must be a positive int32 value, got:", ksize_[i])));
}
OP_REQUIRES_OK(context, context->GetAttr("strides", &stride_));
OP_REQUIRES(context, stride_.size() == 4,
absl::InvalidArgumentError("Sliding window stride field must "
"specify 4 dimensions"));
OP_REQUIRES_OK(context, context->GetAttr("padding", &padding_));
const int32_t ksize_n = GetTensorDim(ksize_, data_format_, 'N');
const int32_t stride_n = GetTensorDim(stride_, data_format_, 'N');
OP_REQUIRES(context, ksize_n == 1 && stride_n == 1,
absl::UnimplementedError(
"Pooling is not yet supported on the batch dimension."));
for (int i = 0; i < ksize_.size(); ++i) {
OP_REQUIRES(context, ksize_[i] != 0,
absl::InvalidArgumentError("ksize cannot be zero"));
}
}
void Compute(OpKernelContext* context) override {
const Tensor& tensor_in = context->input(0);
PoolParameters params{context,
ksize_,
stride_,
padding_,
/*explicit_paddings=*/{},
data_format_,
tensor_in.shape()};
if (!context->status().ok()) {
return;
}
OP_REQUIRES(context, params.depth_window == 1,
absl::UnimplementedError("Non-spatial pooling is not "
"yet supported. Volunteers? :)"));
// For avgpooling, tensor_in should have 4 dimensions.
OP_REQUIRES(context, tensor_in.dims() == 4,
absl::InvalidArgumentError("tensor_in must be 4-dimensional"));
TensorShape output_shape;
OP_REQUIRES_OK(context, params.forward_output_shape(&output_shape));
if (output_shape.num_elements() == 0) {
Tensor* output = nullptr;
OP_REQUIRES_OK(context,
context->allocate_output(0, output_shape, &output));
return;
}
#if CUDNN_VERSION >= 7300
DnnPoolingOp<T>::Compute(context, se::dnn::PoolingMode::kAverage, ksize_,
stride_, padding_, /*explicit_paddings=*/{},
data_format_, tensor_in, output_shape,
/*propagate_nans=*/false);
#else
if (data_format_ == FORMAT_NCHW) {
DnnPoolingOp<T>::Compute(context, se::dnn::PoolingMode::kAverage, ksize_,
stride_, padding_, /*explicit_paddings=*/{},
data_format_, tensor_in, output_shape,
/*propagate_nans=*/false);
} else {
Tensor* output = nullptr;
OP_REQUIRES_OK(context,
context->allocate_output(0, output_shape, &output));
Eigen::PaddingType pt = BrainPadding2EigenPadding(padding_);
functor::SpatialAvgPooling<Device, T>()(
context->eigen_device<Device>(), output->tensor<T, 4>(),
tensor_in.tensor<T, 4>(), params.window_rows, params.window_cols,
params.row_stride, params.col_stride, pt);
}
#endif // CUDNN_VERSION >= 7300
}
private:
std::vector<int32_t> ksize_;
std::vector<int32_t> stride_;
Padding padding_;
TensorFormat data_format_;
};
// Forward declarations of the functor specializations for GPU.
namespace functor {
#define DECLARE_GPU_SPEC(T) \
template <> \
void SpatialAvgPooling<GPUDevice, T>::operator()( \
const GPUDevice& d, typename TTypes<T, 4>::Tensor output, \
typename TTypes<T, 4>::ConstTensor input, int window_rows, \
int window_cols, int row_stride, int col_stride, \
const Eigen::PaddingType& padding); \
extern template struct SpatialAvgPooling<GPUDevice, T>;
TF_CALL_GPU_NUMBER_TYPES(DECLARE_GPU_SPEC)
#undef DECLARE_GPU_SPEC
} // namespace functor
#define REGISTER_GPU_KERNEL_AVGPOOL(T) \
REGISTER_KERNEL_BUILDER( \
Name("AvgPool").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
AvgPoolingOp<GPUDevice, T>);
TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNEL_AVGPOOL)
#undef REGISTER_GPU_KERNEL_AVGPOOL
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
// The operation to compute AvgPool gradients.
// It takes two inputs:
// - The original input tensor shape
// - Backprop tensor for output
// It produces one output: backprop tensor for input.
template <typename Device, class T>
class AvgPoolingGradOp : public OpKernel {
public:
explicit AvgPoolingGradOp(OpKernelConstruction* context) : OpKernel(context) {
std::string data_format;
OP_REQUIRES_OK(context, context->GetAttr("data_format", &data_format));
OP_REQUIRES(context, FormatFromString(data_format, &data_format_),
absl::InvalidArgumentError("Invalid data format"));
OP_REQUIRES(
context, data_format_ == FORMAT_NHWC,
absl::InvalidArgumentError(absl::StrCat(
"Default AvgPoolingGradOp only supports NHWC ", "on device type ",
DeviceTypeString(context->device_type()))));
OP_REQUIRES_OK(context, context->GetAttr("ksize", &ksize_));
OP_REQUIRES(context, ksize_.size() == 4,
absl::InvalidArgumentError("Sliding window ksize field must "
"specify 4 dimensions"));
OP_REQUIRES_OK(context, context->GetAttr("strides", &stride_));
OP_REQUIRES(context, stride_.size() == 4,
absl::InvalidArgumentError("Sliding window strides field must "
"specify 4 dimensions"));
OP_REQUIRES_OK(context, context->GetAttr("padding", &padding_));
OP_REQUIRES(context, ksize_[0] == 1 && stride_[0] == 1,
absl::UnimplementedError(
"Pooling is not yet supported on the batch dimension."));
}
void Compute(OpKernelContext* context) override {
const Tensor& tensor_in_shape = context->input(0);
const Tensor& out_backprop = context->input(1);
// For avgpooling, tensor_in_shape should have 1 dimension, and 4 elements.
OP_REQUIRES(
context,
tensor_in_shape.dims() == 1 && tensor_in_shape.NumElements() == 4,
absl::InvalidArgumentError("out_backprop must be 1-dimensional and 4 "
"elements"));
// For avgpooling, out_backprop should have 4 dimensions.
OP_REQUIRES(
context, out_backprop.dims() == 4,
absl::InvalidArgumentError("out_backprop must be 4-dimensional"));
const int64_t out_backprop_batch = out_backprop.dim_size(0);
const int64_t out_backprop_rows = out_backprop.dim_size(1);
const int64_t out_backprop_cols = out_backprop.dim_size(2);
const int64_t out_backprop_depth = out_backprop.dim_size(3);
TensorShape output_shape;
auto shape_vec = tensor_in_shape.vec<int32_t>();
for (int64_t i = 0; i < tensor_in_shape.NumElements(); ++i) {
OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(shape_vec(i)));
}
const int64_t in_rows = output_shape.dim_size(1);
const int64_t in_cols = output_shape.dim_size(2);
Tensor* output = nullptr;
OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output));
output->flat<T>().setZero();
if (output_shape.num_elements() == 0) {
return;
}
const int window_rows = ksize_[1];
const int window_cols = ksize_[2];
const int depth_window = ksize_[3];
const int row_stride = stride_[1];
const int col_stride = stride_[2];
// We (will) use different code for spatial pooling and
// non-spatial pooling.
//
// Spatial pooling is when depth_window = 1
OP_REQUIRES(context, depth_window == 1,
absl::UnimplementedError("Non-spatial pooling is not "
"yet supported. Volunteers? :)"));
int64_t out_height, out_width, pad_rows, pad_cols;
OP_REQUIRES_OK(context, GetWindowedOutputSize(
in_rows, window_rows, /*dilation_rate=*/1,
row_stride, padding_, &out_height, &pad_rows));
OP_REQUIRES_OK(context, GetWindowedOutputSize(
in_cols, window_cols, /*dilation_rate=*/1,
col_stride, padding_, &out_width, &pad_cols));
const T* out_backprop_ptr = out_backprop.flat<T>().data();
T* input_backprop_ptr = output->flat<T>().data();
for (int64_t r = 0; r < out_backprop_rows; ++r) {
int rindex, rsize;
OP_REQUIRES_OK(context,
GetBroadcastSize(r, in_rows, window_rows, row_stride,
pad_rows, &rindex, &rsize));
for (int64_t c = 0; c < out_backprop_cols; ++c) {
int cindex, csize;
OP_REQUIRES_OK(context,
GetBroadcastSize(c, in_cols, window_cols, col_stride,
pad_cols, &cindex, &csize));
int64_t input_max =
((out_backprop_batch - 1) * in_rows + rindex + rsize - 1) *
in_cols +
cindex + csize - 1;
OP_REQUIRES(context, input_max < output->NumElements(),
absl::InvalidArgumentError(absl::StrCat(
"Output only has ", output->NumElements(),
" elements but computation requested would "
"use element with index=",
input_max)));
}
}
auto shard = [context, out_backprop_ptr, input_backprop_ptr,
out_backprop_rows, out_backprop_cols, out_backprop_depth,
in_rows, in_cols, window_rows, window_cols, row_stride,
col_stride, pad_rows,
pad_cols](int64_t start, int64_t limit) {
for (int64_t b = start; b < limit; ++b) {
for (int64_t r = 0; r < out_backprop_rows; ++r) {
// Calculates row broadcast size. For SAME padding, current
// index could be in the padding area, and r*row_stride +
// window_rows could be beyond the input tensor's boundary. In
// such cases, change the starting index and reduce the
// broadcast size.
int rindex, rsize;
OP_REQUIRES_OK(context,
GetBroadcastSize(r, in_rows, window_rows, row_stride,
pad_rows, &rindex, &rsize));
for (int64_t c = 0; c < out_backprop_cols; ++c) {
// Calculates col broadcast size. For SAME padding, current
// index could be in the padding area, and c*col_stride +
// window_cols could be beyond the input tensor's boundary. In
// such cases, change the starting index and reduce the
// broadcast size.
int cindex, csize;
OP_REQUIRES_OK(context,
GetBroadcastSize(c, in_cols, window_cols, col_stride,
pad_cols, &cindex, &csize));
T divide_coeff(1.0 / (rsize * csize));
int64_t output_index =
(b * out_backprop_rows + r) * out_backprop_cols + c;
for (int64_t r_dst = rindex; r_dst < rindex + rsize; ++r_dst) {
for (int64_t c_dst = cindex; c_dst < cindex + csize; ++c_dst) {
int64_t input_index = (b * in_rows + r_dst) * in_cols + c_dst;
const T* output_offset =
out_backprop_ptr + output_index * out_backprop_depth;
T* input_offset =
input_backprop_ptr + input_index * out_backprop_depth;
for (int64_t d = 0; d < out_backprop_depth; ++d) {
*input_offset += *output_offset * divide_coeff;
++output_offset;
++input_offset;
}
}
}
}
}
}
};
const DeviceBase::CpuWorkerThreads& worker_threads =
*(context->device()->tensorflow_cpu_worker_threads());
const int64_t shard_cost =
window_rows * window_cols * depth_window * in_rows * in_rows * in_cols;
Shard(worker_threads.num_threads, worker_threads.workers,
out_backprop_batch, shard_cost, shard);
}
private:
std::vector<int32_t> ksize_;
std::vector<int32_t> stride_;
Padding padding_;
TensorFormat data_format_;
};
#define REGISTER_CPU_KERNEL_AVGPOOLGRAD(T) \
REGISTER_KERNEL_BUILDER(Name("AvgPoolGrad") \
.Device(DEVICE_CPU) \
.TypeConstraint<T>("T") \
.HostMemory("orig_input_shape"), \
AvgPoolingGradOp<CPUDevice, T>);
TF_CALL_FLOAT_TYPES(REGISTER_CPU_KERNEL_AVGPOOLGRAD);
#undef REGISTER_CPU_KERNEL_AVGPOOLGRAD
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
// A CUDNN based AvgPoolingGrad implementation. It includes the padding as the
// candidates for the pooling operation.
template <class T>
class AvgPoolingGradOp<GPUDevice, T> : public OpKernel {
public:
typedef GPUDevice Device;
explicit AvgPoolingGradOp(OpKernelConstruction* context) : OpKernel(context) {
std::string data_format;
OP_REQUIRES_OK(context, context->GetAttr("data_format", &data_format));
OP_REQUIRES(context, FormatFromString(data_format, &data_format_),
absl::InvalidArgumentError("Invalid data format"));
OP_REQUIRES_OK(context, context->GetAttr("ksize", &ksize_));
OP_REQUIRES(context, ksize_.size() == 4,
absl::InvalidArgumentError("Sliding window ksize field must "
"specify 4 dimensions"));
OP_REQUIRES_OK(context, context->GetAttr("strides", &stride_));
OP_REQUIRES(context, stride_.size() == 4,
absl::InvalidArgumentError("Sliding window strides field must "
"specify 4 dimensions"));
OP_REQUIRES_OK(context, context->GetAttr("padding", &padding_));
const int32_t ksize_n = GetTensorDim(ksize_, data_format_, 'N');
const int32_t stride_n = GetTensorDim(stride_, data_format_, 'N');
OP_REQUIRES(context, ksize_n == 1 && stride_n == 1,
absl::UnimplementedError(
"Pooling is not yet supported on the batch dimension."));
}
void Compute(OpKernelContext* context) override {
const Tensor& tensor_in_shape = context->input(0);
const Tensor& out_backprop = context->input(1);
// For avgpooling, tensor_in_shape should have 1 dimension, and 4 elements.
OP_REQUIRES(
context,
tensor_in_shape.dims() == 1 && tensor_in_shape.NumElements() == 4,
absl::InvalidArgumentError("out_backprop must be 1-dimensional and 4 "
"elements"));
// For avgpooling, out_backprop should have 4 dimensions.
OP_REQUIRES(
context, out_backprop.dims() == 4,
absl::InvalidArgumentError("out_backprop must be 4-dimensional"));
TensorShape output_shape;
auto shape_vec = tensor_in_shape.vec<int32_t>();
for (int64_t i = 0; i < tensor_in_shape.NumElements(); ++i) {
OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(shape_vec(i)));
}
if (output_shape.num_elements() == 0) {
Tensor* output = nullptr;
OP_REQUIRES_OK(context,
context->allocate_output(0, output_shape, &output));
return;
}
DnnPoolingGradOp<T>::Compute(
context, se::dnn::PoolingMode::kAverage, ksize_, stride_, padding_,
/*explicit_paddings=*/{}, data_format_, nullptr, nullptr, out_backprop,
output_shape, /*propagate_nans=*/false);
}
private:
std::vector<int32_t> ksize_;
std::vector<int32_t> stride_;
Padding padding_;
TensorFormat data_format_;
};
REGISTER_KERNEL_BUILDER(Name("AvgPoolGrad")
.Device(DEVICE_GPU)
.TypeConstraint<double>("T")
.HostMemory("orig_input_shape")
.Label("cudnn"),
AvgPoolingGradOp<GPUDevice, double>);
REGISTER_KERNEL_BUILDER(Name("AvgPoolGrad")
.Device(DEVICE_GPU)
.TypeConstraint<float>("T")
.HostMemory("orig_input_shape")
.Label("cudnn"),
AvgPoolingGradOp<GPUDevice, float>);
REGISTER_KERNEL_BUILDER(Name("AvgPoolGrad")
.Device(DEVICE_GPU)
.TypeConstraint<Eigen::half>("T")
.HostMemory("orig_input_shape")
.Label("cudnn"),
AvgPoolingGradOp<GPUDevice, Eigen::half>);
// A custom GPU kernel based AvgPoolingGrad implementation. It includes the
// padding as the candidates for the pooling operation.
template <class T>
class AvgPoolingGradOpCustomGPUKernel : public OpKernel {
public:
typedef GPUDevice Device;
explicit AvgPoolingGradOpCustomGPUKernel(OpKernelConstruction* context)
: OpKernel(context) {
std::string data_format;
OP_REQUIRES_OK(context, context->GetAttr("data_format", &data_format));
OP_REQUIRES(context, FormatFromString(data_format, &data_format_),
absl::InvalidArgumentError("Invalid data format"));
OP_REQUIRES_OK(context, context->GetAttr("ksize", &ksize_));
OP_REQUIRES(context, ksize_.size() == 4,
absl::InvalidArgumentError("Sliding window ksize field must "
"specify 4 dimensions"));
OP_REQUIRES_OK(context, context->GetAttr("strides", &stride_));
OP_REQUIRES(context, stride_.size() == 4,
absl::InvalidArgumentError("Sliding window strides field must "
"specify 4 dimensions"));
OP_REQUIRES_OK(context, context->GetAttr("padding", &padding_));
const int32_t ksize_n = GetTensorDim(ksize_, data_format_, 'N');
const int32_t stride_n = GetTensorDim(stride_, data_format_, 'N');
OP_REQUIRES(context, ksize_n == 1 && stride_n == 1,
absl::UnimplementedError(
"Pooling is not yet supported on the batch dimension."));
}
void Compute(OpKernelContext* context) override {
const Tensor& tensor_in_shape = context->input(0);
const Tensor& out_backprop = context->input(1);
// For avgpooling, tensor_in_shape should have 1 dimension, and 4 elements.
OP_REQUIRES(
context,
tensor_in_shape.dims() == 1 && tensor_in_shape.NumElements() == 4,
absl::InvalidArgumentError("out_backprop must be 1-dimensional and 4 "
"elements"));
// For avgpooling, out_backprop should have 4 dimensions.
OP_REQUIRES(
context, out_backprop.dims() == 4,
absl::InvalidArgumentError("out_backprop must be 4-dimensional"));
TensorShape output_shape;
auto shape_vec = tensor_in_shape.vec<int32_t>();
for (int64_t i = 0; i < tensor_in_shape.NumElements(); ++i) {
OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(shape_vec(i)));
}
if (output_shape.num_elements() == 0) {
Tensor* output = nullptr;
OP_REQUIRES_OK(context,
context->allocate_output(0, output_shape, &output));
return;
}
#if CUDNN_VERSION >= 7300
DnnPoolingGradOp<T>::Compute(context, se::dnn::PoolingMode::kAverage,
ksize_, stride_, padding_,
/*explicit_paddings=*/{}, data_format_,
nullptr, nullptr, out_backprop, output_shape,
/*propagate_nans=*/false);
#else
if (data_format_ == FORMAT_NHWC) {
const int64 out_backprop_batch = out_backprop.dim_size(0);
const int64 out_backprop_rows = out_backprop.dim_size(1);
const int64 out_backprop_cols = out_backprop.dim_size(2);
const int64 out_backprop_depth = out_backprop.dim_size(3);
const int64 in_rows = output_shape.dim_size(1);
const int64 in_cols = output_shape.dim_size(2);
Tensor* output = nullptr;
OP_REQUIRES_OK(context,
context->allocate_output(0, output_shape, &output));
const int window_rows = ksize_[1];
const int window_cols = ksize_[2];
const int depth_window = ksize_[3];
const int row_stride = stride_[1];
const int col_stride = stride_[2];
// We (will) use different code for spatial pooling and
// non-spatial pooling.
//
// Spatial pooling is when depth_window = 1
OP_REQUIRES(context, depth_window == 1,
errors::Unimplemented("Non-spatial pooling is not "
"yet supported. Volunteers? :)"));
int64 out_height, out_width, pad_rows, pad_cols;
OP_REQUIRES_OK(
context,
GetWindowedOutputSize(in_rows, window_rows, /*dilation_rate=*/1,
row_stride, padding_, &out_height, &pad_rows));
OP_REQUIRES_OK(context, GetWindowedOutputSize(
in_cols, window_cols, /*dilation_rate=*/1,
col_stride, padding_, &out_width, &pad_cols));
OP_REQUIRES_OK(
context,
RunAvePoolBackwardNHWC<T>(out_backprop.flat<T>().data(), // top_diff
out_backprop_batch, // num
in_rows, // height
in_cols, // width
out_backprop_depth, // channels
out_backprop_rows, // pooled_height
out_backprop_cols, // pooled_width
window_rows, // kernel_h
window_cols, // kernel_w
row_stride, // stride_h
col_stride, // stride_w
pad_rows, // pad_t
pad_cols, // pad_l
output->flat<T>().data(), // bottom_diff
context->eigen_gpu_device())); // d
} else {
DnnPoolingGradOp<T>::Compute(context, se::dnn::PoolingMode::kAverage,
ksize_, stride_, padding_,
/*explicit_paddings=*/{}, data_format_,
nullptr, nullptr, out_backprop, output_shape,
/*propagate_nans=*/false);
}
#endif // CUDNN_VERSION >= 7300
}
private:
std::vector<int32_t> ksize_;
std::vector<int32_t> stride_;
Padding padding_;
TensorFormat data_format_;
};
#define REGISTER_GPU_KERNEL_AVGPOOLGRAD(T) \
REGISTER_KERNEL_BUILDER(Name("AvgPoolGrad") \
.Device(DEVICE_GPU) \
.TypeConstraint<T>("T") \
.HostMemory("orig_input_shape"), \
AvgPoolingGradOpCustomGPUKernel<T>);
TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNEL_AVGPOOLGRAD);
#undef REGISTER_GPU_KERNEL_AVGPOOLGRAD
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow
tensorflow/core/kernels/avgpooling_op.h
+75
View File
@@ -0,0 +1,75 @@
/* Copyright 2015 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_AVGPOOLING_OP_H_
#define TENSORFLOW_CORE_KERNELS_AVGPOOLING_OP_H_
// Functor definition for AvgPoolingOp, must be compilable by nvcc.
#include "absl/status/status.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/kernels/eigen_pooling.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace functor {
template <typename Device, typename T>
struct SpatialAvgPooling {
void operator()(const Device& d, typename TTypes<T, 4>::Tensor output,
typename TTypes<T, 4>::ConstTensor input, int window_rows,
int window_cols, int row_stride, int col_stride,
const Eigen::PaddingType& padding) {
MaybeWith32BitIndexing<Device>(
[&](auto output32, auto input32) {
// Because we swap the layout, we swap the row/cols as well.
output32.swap_layout().device(d) = Eigen::SpatialAvgPooling(
input32.swap_layout(), window_cols, window_rows, col_stride,
row_stride, padding);
},
output, input);
}
};
} // namespace functor
typedef Eigen::GpuDevice GPUDevice;
// Launch a custom GPU kernels from Yanqing for the avgpooling backward
// operation that works NHWC data formats. Arguments:
// top_diff: backprop to the output of the pooling layer
// num: number of input batches
// height: input height
// width: input width
// channels: number of input channels
// pooled_height: the height of the output to the pooling layer
// pooled_width: the width of the output to the pooling layer
// kernel_h: the height of the pooling kernel
// kernel_w: the width of the pooling kernel
// stride_h: the height of the vertical stride
// stride_w: the width of the horizontal stride
// pad_t: padding size to the top side
// pad_l: padding size to the left side
// bottom_diff: backprop to the input of the pooling layer.
template <typename T>
absl::Status RunAvePoolBackwardNHWC(const T* top_diff, int num, int height,
int width, int channels, int pooled_height,
int pooled_width, int kernel_h,
int kernel_w, int stride_h, int stride_w,
int pad_t, int pad_l, T* bottom_diff,
const GPUDevice& d);
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_AVGPOOLING_OP_H_
tensorflow/core/kernels/avgpooling_op_gpu.cu.cc
+129
View File
@@ -0,0 +1,129 @@
/* Copyright 2015 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.
==============================================================================*/
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define EIGEN_USE_GPU
#include <stdio.h>
#include <iostream>
#include <limits>
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/kernels/avgpooling_op.h"
#include "tensorflow/core/util/gpu_kernel_helper.h"
#include "tensorflow/core/util/overflow.h"
namespace tensorflow {
typedef Eigen::GpuDevice GPUDevice;
#define DEFINE_GPU_KERNELS(T) \
template struct functor::SpatialAvgPooling<GPUDevice, T>;
DEFINE_GPU_KERNELS(Eigen::half)
DEFINE_GPU_KERNELS(Eigen::bfloat16)
DEFINE_GPU_KERNELS(float)
DEFINE_GPU_KERNELS(double)
#undef DEFINE_GPU_KERNELS
template <typename dtype>
__global__ void AvePoolBackwardNHWC(
const int nthreads, const dtype* const __restrict__ top_diff, const int num,
const int height, const int width, const int channels,
const int pooled_height, const int pooled_width, const int kernel_h,
const int kernel_w, const int stride_h, const int stride_w, const int pad_t,
const int pad_l, dtype* const __restrict__ bottom_diff) {
GPU_1D_KERNEL_LOOP(index, nthreads) {
// find out the local index
// find out the local offset
const int c = index % channels;
const int w = index / channels % width + pad_l;
const int h = (index / channels / width) % height + pad_t;
const int n = index / channels / width / height;
const int phstart = (h < kernel_h) ? 0 : (h - kernel_h) / stride_h + 1;
const int phend = min(h / stride_h + 1, pooled_height);
const int pwstart = (w < kernel_w) ? 0 : (w - kernel_w) / stride_w + 1;
const int pwend = min(w / stride_w + 1, pooled_width);
dtype gradient(0);
const dtype* const top_diff_slice =
top_diff + n * pooled_height * pooled_width * channels + c;
for (int ph = phstart; ph < phend; ++ph) {
for (int pw = pwstart; pw < pwend; ++pw) {
// figure out the pooling size
int hstart = ph * stride_h - pad_t;
int wstart = pw * stride_w - pad_l;
int hend = min(hstart + kernel_h, height);
int wend = min(wstart + kernel_w, width);
hstart = max(hstart, 0);
wstart = max(wstart, 0);
int pool_size = (hend - hstart) * (wend - wstart);
gradient += top_diff_slice[(ph * pooled_width + pw) * channels] /
dtype(pool_size);
}
}
bottom_diff[index] = gradient;
}
}
template <typename T>
absl::Status RunAvePoolBackwardNHWC(const T* const top_diff, const int num,
const int height, const int width,
const int channels, const int pooled_height,
const int pooled_width, const int kernel_h,
const int kernel_w, const int stride_h,
const int stride_w, const int pad_t,
const int pad_l, T* const bottom_diff,
const GPUDevice& d) {
int64_t size_1 = MultiplyWithoutOverflow(num, height);
int64_t size_2 = MultiplyWithoutOverflow(size_1, width);
int64_t x_size = MultiplyWithoutOverflow(size_2, channels);
if (x_size < 0 || x_size > std::numeric_limits<int32_t>::max()) {
return absl::InternalError(
"RunAvePoolBackwardNHWC: num * height * width * channels exceeds "
"int32 bounds");
}
GpuLaunchConfig config = GetGpuLaunchConfig(x_size, d);
TF_CHECK_OK(GpuLaunchKernel(
AvePoolBackwardNHWC<T>, config.block_count, config.thread_per_block, 0,
d.stream(), config.virtual_thread_count, top_diff, num, height, width,
channels, pooled_height, pooled_width, kernel_h, kernel_w, stride_h,
stride_w, pad_t, pad_t, bottom_diff));
return d.ok() ? absl::OkStatus()
: absl::InternalError("GPU execution failed");
}
#define DECLARE_GPU_SPEC(T) \
template absl::Status RunAvePoolBackwardNHWC( \
const T* const top_diff, const int num, const int height, \
const int width, const int channels, const int pooled_height, \
const int pooled_width, const int kernel_h, const int kernel_w, \
const int stride_h, const int stride_w, const int pad_t, \
const int pad_l, T* const bottom_diff, const GPUDevice& d);
DECLARE_GPU_SPEC(Eigen::half);
DECLARE_GPU_SPEC(Eigen::bfloat16);
DECLARE_GPU_SPEC(float);
DECLARE_GPU_SPEC(double);
#undef DECLARE_GPU_SPEC
} // end namespace tensorflow
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
tensorflow/core/kernels/barrier_ops.cc
+695
View File
@@ -0,0 +1,695 @@
/* Copyright 2016 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.
==============================================================================*/
// See docs in ../ops/data_flow_ops.cc.
#include <limits.h>
#include <unordered_map>
#include <vector>
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/resource_op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/kernels/priority_queue.h"
#include "tensorflow/core/kernels/queue_base.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/notification.h"
#include "tensorflow/core/lib/gtl/map_util.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace barrier {
class Barrier : public ResourceBase {
public:
typedef std::vector<Tensor> Tuple;
typedef std::function<void()> DoneCallback;
typedef std::function<void(const Tensor&, const Tensor&, const Tuple&)>
IndicesKeysValuesCallback;
Barrier(const DataTypeVector& value_component_types,
const std::vector<TensorShape>& value_component_shapes,
const std::string& name)
: closed_(false),
queue_closed_(false),
queue_cancelled_(false),
cancel_pending_enqueues_(false),
value_component_types_(value_component_types),
value_component_shapes_(value_component_shapes),
name_(name),
input_index_(std::numeric_limits<int64_t>::min()) {
DataTypeVector queue_component_types;
std::vector<TensorShape> queue_component_shapes;
// First queue component is for the input index;
// Second queue component is for the key;
// remaining queue components are for the value.
queue_component_types.push_back(DT_INT64);
queue_component_types.push_back(DT_STRING);
for (DataType dt : value_component_types) {
queue_component_types.push_back(dt);
}
// NOTE(mrry): PriorityQueue expects all shapes specified because
// we'll be issuing TakeMany.
queue_component_shapes.push_back(TensorShape({}));
queue_component_shapes.push_back(TensorShape({}));
queue_component_shapes.insert(queue_component_shapes.end(),
value_component_shapes.begin(),
value_component_shapes.end());
ready_queue_ = new PriorityQueue(
QueueBase::kUnbounded /* capacity */, queue_component_types,
queue_component_shapes, absl::StrCat(name_, "_queue"));
}
absl::Status Initialize() { return ready_queue_->Initialize(); }
template <typename T>
void TryInsertMany(const Tensor& keys, int component_index,
const Tensor& values, OpKernelContext* ctx,
const DoneCallback& callback) {
TensorShape element_shape = values.shape();
OP_REQUIRES_ASYNC(
ctx, keys.NumElements() == 0 || element_shape.num_elements() > 0,
absl::InvalidArgumentError(
absl::StrCat("Tensors with no elements are not supported ", name_,
": received shape ", element_shape.DebugString())),
callback);
if (element_shape.dims() > 0) element_shape.RemoveDim(0);
const std::size_t num_inserted = keys.NumElements();
// For each key, update the corresponding incomplete tuple with the
// the corresponding given value at component_index.
// This will be passed to the final callback at the very end.
bool new_elements = false;
// Will be used for the final insert into the queue.
Tuple insert_tuple;
{
mutex_lock lock(mu_);
if (closed_) {
OP_REQUIRES_ASYNC(
ctx,
!cancel_pending_enqueues_ &&
(num_inserted == 0 || !incomplete_.empty()),
errors::Cancelled(
"Barrier ", name_, " is closed. Pending enqueues cancelled: ",
cancel_pending_enqueues_,
". Number of new insertions: ", num_inserted,
". Number of incomplete keys: ", incomplete_.size(), "."),
callback);
}
// Step 1: insert into the incomplete map and identify which
// entries are, in fact, complete and ready for enqueueing. Store
// them in a vector
std::vector<Tuple> ready_tuples;
for (int i = 0; i < num_inserted; ++i) {
OP_REQUIRES_OK_ASYNC(
ctx,
InsertOneLocked<T>(ctx, keys, values, element_shape,
component_index, i, &ready_tuples,
&new_elements),
callback);
}
if (new_elements) ++input_index_;
// This probably won't happen before the heat death of the
// universe, but who knows? Moore's law FTW.
OP_REQUIRES_ASYNC(
ctx, input_index_ != std::numeric_limits<int64_t>::max(),
absl::InternalError(absl::StrCat(
"Barrier has had ", input_index_,
" insertions and can no longer keep track of new ones.")),
callback);
if (ready_tuples.empty()) {
// Nothing to insert into the queue - so return early.
callback();
return;
}
// We have something to Enqueue. Convert the Tuples into a single
// tuple by slicing entries into new Tensors. This part is slow
// but seems the cleanest solution for now.
insert_tuple.reserve(2 + num_components()); // indices, keys, rest
int insertion_size = ready_tuples.size();
for (int i = 0; i < 2 + num_components(); ++i) {
TensorShape component_shape(ready_tuples[0][i].shape());
component_shape.InsertDim(0, insertion_size);
Tensor component(ready_tuples[0][i].dtype(), component_shape);
for (int b = 0; b < insertion_size; ++b) {
OP_REQUIRES_OK_ASYNC(
ctx,
batch_util::CopyElementToSlice(std::move(ready_tuples[b][i]),
&component, b),
callback);
}
insert_tuple.push_back(component);
}
}
// Update the input index for the next batch.
ready_queue_->TryEnqueueMany(
insert_tuple, ctx,
// To avoid early closing of the queue, only close it if the
// SQSS is closed, nothing is left in the incomplete set,
// the queue is not already marked as closed, and (most
// importantly), the queue has entries in it.
[this, ctx, callback]() {
if (!ctx->status().ok()) {
callback();
return;
}
{
mutex_lock lock(mu_);
int32_t ready = ready_size();
if (closed_ && incomplete_.empty() && queue_closed_ && ready > 0) {
CloseQueueLocked(ctx, false, callback);
} else {
callback();
}
return;
}
});
}
void TryTakeMany(int num_elements, bool allow_small_batch, int64_t timeout,
OpKernelContext* ctx,
const IndicesKeysValuesCallback& callback) {
int num_elements_to_deliver = num_elements;
{
mutex_lock lock(mu_);
if (closed_) {
int available_elements = ready_size();
if (allow_small_batch) {
// We want to deliver a maximum of num_elements, if there are less
// elements available, we deliver at most the available_elements. If
// there are no
// elements available, a call to TryTakeMany should fail with
// OutOfRange. We trigger this error by setting the request here to 1.
num_elements_to_deliver = std::min(num_elements, available_elements);
} else {
// We're happy to wait for additional elements to be completed.
available_elements += incomplete_.size();
}
// If there are 0 available elements or less elements than the
// number we can deliver, then we are done.
if (available_elements < std::max(num_elements_to_deliver, 1)) {
ctx->SetStatus(absl::OutOfRangeError(absl::StrCat(
"Barrier '", name_, "' is closed and has ",
"insufficient elements (requested ", num_elements_to_deliver,
", total size ", available_elements, ")")));
callback(Tensor(DT_INT64), Tensor(DT_STRING), Tuple());
return;
}
}
}
ready_queue_->TryDequeueMany(
num_elements_to_deliver, ctx, allow_small_batch,
[this, ctx, callback](const Tuple& t) {
Tensor indices(DT_INT64);
Tensor keys(DT_STRING);
Tuple values;
if (!ctx->status().ok()) {
callback(indices, keys, values);
return;
}
CHECK_EQ(t.size(), 2 + num_components());
indices = t[0];
keys = t[1];
values.insert(values.begin(), t.begin() + 2, t.end());
callback(indices, keys, values);
});
}
void Close(OpKernelContext* ctx, bool cancel_pending_enqueues,
const DoneCallback& callback) {
mutex_lock lock(mu_);
// We're allowed to close twice if the first close wasn't a
// cancel but the second one is.
if (closed_ && (cancel_pending_enqueues_ || !cancel_pending_enqueues)) {
ctx->SetStatus(absl::CancelledError(
absl::StrCat("Barrier '", name_, "' is already closed.")));
callback();
return;
}
cancel_pending_enqueues_ = cancel_pending_enqueues;
closed_ = true;
if (cancel_pending_enqueues_ || incomplete_.empty()) {
incomplete_.clear();
// CloseQueueLocked runs the callback
CloseQueueLocked(ctx, cancel_pending_enqueues_, callback);
return;
}
callback();
}
int32_t ready_size() { return ready_queue_->size(); }
int32_t incomplete_size() {
mutex_lock lock(mu_);
return incomplete_.size();
}
const std::string& name() const { return name_; }
int num_components() const { return value_component_types_.size(); }
DataType component_type(int i) const {
CHECK_GE(i, 0);
CHECK_LT(static_cast<size_t>(i), value_component_types_.size());
return value_component_types_[i];
}
const DataTypeVector component_types() const {
return value_component_types_;
}
const absl::Span<const TensorShape> component_shapes() const {
return value_component_shapes_;
}
~Barrier() override TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
mutex_lock lock(mu_);
incomplete_.clear();
ready_queue_->Unref();
}
std::string DebugString() const override { return "A barrier"; }
protected:
template <typename T>
absl::Status InsertOneLocked(OpKernelContext* ctx, const Tensor& keys,
const Tensor& values,
const TensorShape& element_shape,
int component_index, int i,
std::vector<Tuple>* ready_tuples,
bool* new_elements)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
auto keys_vec = keys.flat<tstring>();
auto values_matrix = values.flat_outer_dims<T>();
TensorTuple* element_ptr;
if (closed_) {
element_ptr = gtl::FindOrNull(incomplete_, keys_vec(i));
if (element_ptr == nullptr) {
return absl::CancelledError(absl::StrCat(
"Barrier ", name_,
" is closed, but attempted to insert a brand new key: ",
keys_vec(i), ". Pending enqueues cancelled: ",
cancel_pending_enqueues_, ". Insertion index: ", i,
". Number of incomplete keys: ", incomplete_.size(), "."));
}
} else {
element_ptr =
&gtl::LookupOrInsert(&incomplete_, keys_vec(i), TensorTuple());
}
TensorTuple& element = *element_ptr;
if (element.empty()) { // Never seen before key
// Added a new element, for keeping track of the insertion index
*new_elements = true;
// Initialize the incomplete tuple for a new key.
element.reserve(1 + num_components());
// The first entry in element is the priority: the
// input_index_, so that tensors that entered the Barrier
// earlier have higher priority in the queue.
Tensor allocate_index_tensor;
TF_RETURN_IF_ERROR(ctx->allocate_temp(DT_INT64, TensorShape({}),
&allocate_index_tensor));
Tensor index_tensor(DT_INT64, TensorShape({}));
allocate_index_tensor.scalar<int64_t>()() = input_index_;
element.push_back(allocate_index_tensor);
// The rest of the element stores uninitialized Tensors with
// the appropriate dtype.
for (int j = 0; j < num_components(); ++j) {
Tensor uninitialized(component_type(j));
element.push_back(Tensor(uninitialized));
}
}
const Tensor& component = element[1 + component_index];
if (component.IsInitialized() && component.NumElements() > 0) {
return absl::InvalidArgumentError(absl::StrCat(
"Key ", keys_vec(i), " already has a value for component ",
component_index, " in barrier ", name()));
}
// Extract the slice corresponding to the value from the value Tensor,
// and store it in the incomplete tuple at component_index.
Tensor next_element;
TF_RETURN_IF_ERROR(
ctx->allocate_temp(values.dtype(), element_shape, &next_element));
element[1 + component_index] = next_element;
next_element.flat<T>() = values_matrix.template chip<0>(i);
// Check the components of the tuple to see if it has become complete
// (i.e. all of its components are initialized). If so, add it to the
// ready queue.
bool is_complete = true;
for (int j = 0; is_complete && j < element.size(); ++j) {
is_complete = element[j].IsInitialized() && element[j].NumElements() > 0;
}
if (is_complete) {
// Add tuple to the ready queue. A queue tuple has the index
// as the first element and the key as the second element,
// followed by the value components.
Tuple ready_tuple;
ready_tuple.reserve(2 + num_components()); // index, key, rest
// Build a tensor for the key. TODO(mrry): Something more efficient.
Tensor key;
TF_RETURN_IF_ERROR(ctx->allocate_temp(DT_STRING, TensorShape({}), &key));
ready_tuple.push_back(element[0]); // index
ready_tuple.push_back(key); // key
ready_tuple[1].scalar<tstring>()() = keys_vec(i); // set the key
for (int j = 1; j < num_components() + 1; ++j) {
ready_tuple.push_back(element[j]);
}
incomplete_.erase(incomplete_.find(keys_vec(i)));
TF_RETURN_IF_ERROR(ready_queue_->ValidateTuple(ready_tuple));
ready_tuples->push_back(ready_tuple);
}
return absl::OkStatus();
}
void CloseQueueLocked(OpKernelContext* ctx, bool cancel_pending_enqueues,
const DoneCallback& callback)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
// CloseQueueLocked may only be called with mu_ held.
if (!cancel_pending_enqueues && queue_closed_) {
callback();
return;
}
if (cancel_pending_enqueues && queue_cancelled_) {
callback();
return;
}
queue_closed_ = true;
if (cancel_pending_enqueues) queue_cancelled_ = true;
if (!ready_queue_->is_closed()) {
ready_queue_->Close(ctx, cancel_pending_enqueues, callback);
}
}
private:
typedef std::vector<Tensor> TensorTuple;
mutex mu_;
bool closed_ TF_GUARDED_BY(mu_);
bool queue_closed_ TF_GUARDED_BY(mu_);
bool queue_cancelled_ TF_GUARDED_BY(mu_);
bool cancel_pending_enqueues_ TF_GUARDED_BY(mu_);
const DataTypeVector value_component_types_;
const std::vector<TensorShape>& value_component_shapes_;
const std::string name_;
int64_t input_index_ TF_GUARDED_BY(mu_);
std::unordered_map<std::string, TensorTuple> incomplete_ TF_GUARDED_BY(mu_);
PriorityQueue* ready_queue_;
Barrier(const Barrier&) = delete;
void operator=(const Barrier&) = delete;
};
class BarrierOp : public ResourceOpKernel<Barrier> {
public:
explicit BarrierOp(OpKernelConstruction* context)
: ResourceOpKernel(context) {
OP_REQUIRES_OK(
context, context->GetAttr("component_types", &value_component_types_));
OP_REQUIRES_OK(context,
context->GetAttr("shapes", &value_component_shapes_));
OP_REQUIRES(context,
value_component_shapes_.size() == value_component_types_.size(),
absl::InvalidArgumentError(
"All of the component shapes must be specified"));
int32_t value_capacity;
OP_REQUIRES_OK(context, context->GetAttr("capacity", &value_capacity));
OP_REQUIRES(context, value_capacity == -1,
absl::InvalidArgumentError(
"Barrier only accepts capacity=-1. Feed the "
"inputs to your Barrier through a queue to enforce a "
"limited capacity."));
}
private:
absl::Status CreateResource(Barrier** barrier) override
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
*barrier = new Barrier(value_component_types_, value_component_shapes_,
cinfo_.name());
if (*barrier == nullptr) {
return absl::ResourceExhaustedError("Failed to allocate barrier");
}
return (*barrier)->Initialize();
}
absl::Status VerifyResource(Barrier* barrier) override
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
if (barrier->component_types() != value_component_types_) {
return absl::InvalidArgumentError(absl::StrCat(
"Shared barrier '", cinfo_.name(), "' has component types ",
DataTypeSliceString(barrier->component_types()),
" but requested component types were ",
DataTypeSliceString(value_component_types_)));
}
if (barrier->component_shapes() != value_component_shapes_) {
return absl::InvalidArgumentError(absl::StrCat(
"Shared barrier '", cinfo_.name(), "' has component shapes ",
TensorShapeUtils::ShapeListString(barrier->component_shapes()),
" but requested component shapes were ",
TensorShapeUtils::ShapeListString(value_component_shapes_)));
}
return absl::OkStatus();
}
DataTypeVector value_component_types_;
std::vector<TensorShape> value_component_shapes_;
BarrierOp(const BarrierOp&) = delete;
void operator=(const BarrierOp&) = delete;
};
REGISTER_KERNEL_BUILDER(Name("Barrier").Device(DEVICE_CPU), BarrierOp);
class BarrierOpKernel : public AsyncOpKernel {
public:
explicit BarrierOpKernel(OpKernelConstruction* context)
: AsyncOpKernel(context) {}
void ComputeAsync(OpKernelContext* ctx, DoneCallback callback) final {
Barrier* barrier = nullptr;
OP_REQUIRES_OK_ASYNC(ctx, GetResourceFromContext(ctx, "handle", &barrier),
callback);
ComputeAsync(ctx, barrier, [callback, barrier]() {
barrier->Unref();
callback();
});
}
protected:
virtual void ComputeAsync(OpKernelContext* ctx, Barrier* barrier,
DoneCallback callback) = 0;
};
template <typename T>
class InsertManyOp : public BarrierOpKernel {
public:
explicit InsertManyOp(OpKernelConstruction* context)
: BarrierOpKernel(context) {
OP_REQUIRES_OK(context,
context->GetAttr("component_index", &component_index_));
}
protected:
void ComputeAsync(OpKernelContext* ctx, Barrier* barrier,
DoneCallback callback) override {
OP_REQUIRES_ASYNC(
ctx, component_index_ < barrier->num_components(),
absl::InvalidArgumentError(absl::StrCat(
"The component ID is out of range ", component_index_,
" > num_components", " (= ", barrier->num_components(), ")")),
callback);
OP_REQUIRES_OK_ASYNC(
ctx,
ctx->MatchSignature({DT_STRING_REF, DT_STRING,
barrier->component_type(component_index_)},
{}),
callback);
const Tensor* keys;
const Tensor* values;
OP_REQUIRES_OK_ASYNC(ctx, ctx->input("keys", &keys), callback);
OP_REQUIRES_OK_ASYNC(ctx, ctx->input("values", &values), callback);
barrier->TryInsertMany<T>(*keys, component_index_, *values, ctx, callback);
}
private:
int component_index_;
InsertManyOp(const InsertManyOp&) = delete;
void operator=(const InsertManyOp&) = delete;
};
#define REGISTER_INSERTMANY(T) \
REGISTER_KERNEL_BUILDER( \
Name("BarrierInsertMany").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
InsertManyOp<T>);
TF_CALL_ALL_TYPES(REGISTER_INSERTMANY);
#undef REGISTER_INSERTMANY
class TakeManyOp : public BarrierOpKernel {
public:
explicit TakeManyOp(OpKernelConstruction* context)
: BarrierOpKernel(context) {
OP_REQUIRES_OK(context, context->GetAttr("timeout_ms", &timeout_));
// TODO(keveman): Enable timeout.
OP_REQUIRES(context, timeout_ == -1,
absl::InvalidArgumentError("Timeout not supported yet."));
OP_REQUIRES_OK(context,
context->GetAttr("allow_small_batch", &allow_small_batch_));
}
protected:
void ComputeAsync(OpKernelContext* ctx, Barrier* barrier,
DoneCallback callback) override {
const Tensor* Tnum_elements;
OP_REQUIRES_OK_ASYNC(ctx, ctx->input("num_elements", &Tnum_elements),
callback);
OP_REQUIRES_ASYNC(
ctx, TensorShapeUtils::IsScalar(Tnum_elements->shape()),
absl::InvalidArgumentError("num_elements must be a scalar."), callback);
const int32_t num_elements = Tnum_elements->scalar<int32_t>()();
DataTypeVector expected_inputs = {DT_STRING_REF, DT_INT32};
// The first output is the insertion index, the second output is the key.
DataTypeVector expected_outputs = {DT_INT64, DT_STRING};
for (DataType dt : barrier->component_types()) {
expected_outputs.push_back(dt);
}
OP_REQUIRES_OK_ASYNC(
ctx, ctx->MatchSignature(expected_inputs, expected_outputs), callback);
barrier->TryTakeMany(
num_elements, allow_small_batch_, timeout_, ctx,
[ctx, callback](const Tensor& indices, const Tensor& keys,
const Barrier::Tuple& values) {
if (!ctx->status().ok()) {
callback();
return;
}
// At this point, indices, keys, and values
// have all been written to successfully.
OP_REQUIRES_OK_ASYNC(ctx, ctx->set_output("indices", indices),
callback);
OP_REQUIRES_OK_ASYNC(ctx, ctx->set_output("keys", keys), callback);
OpOutputList values_output;
OP_REQUIRES_OK_ASYNC(ctx, ctx->output_list("values", &values_output),
callback);
for (size_t i = 0; i < values.size(); ++i) {
values_output.set(i, values[i]);
}
callback();
});
}
private:
int64_t timeout_;
bool allow_small_batch_;
TakeManyOp(const TakeManyOp&) = delete;
void operator=(const TakeManyOp&) = delete;
};
REGISTER_KERNEL_BUILDER(Name("BarrierTakeMany").Device(DEVICE_CPU), TakeManyOp);
class BarrierCloseOp : public BarrierOpKernel {
public:
explicit BarrierCloseOp(OpKernelConstruction* context)
: BarrierOpKernel(context) {
OP_REQUIRES_OK(context, context->GetAttr("cancel_pending_enqueues",
&cancel_pending_enqueues_));
}
protected:
void ComputeAsync(OpKernelContext* ctx, Barrier* barrier,
DoneCallback callback) override {
barrier->Close(ctx, cancel_pending_enqueues_, callback);
}
private:
bool cancel_pending_enqueues_;
BarrierCloseOp(const BarrierCloseOp&) = delete;
void operator=(const BarrierCloseOp&) = delete;
};
REGISTER_KERNEL_BUILDER(Name("BarrierClose").Device(DEVICE_CPU),
BarrierCloseOp);
class BarrierIncompleteSizeOp : public BarrierOpKernel {
public:
explicit BarrierIncompleteSizeOp(OpKernelConstruction* context)
: BarrierOpKernel(context) {}
protected:
void ComputeAsync(OpKernelContext* ctx, Barrier* barrier,
DoneCallback callback) override {
Tensor* Tsize = nullptr;
OP_REQUIRES_OK_ASYNC(ctx, ctx->allocate_output(0, TensorShape({}), &Tsize),
callback);
Tsize->scalar<int32_t>().setConstant(barrier->incomplete_size());
callback();
}
};
REGISTER_KERNEL_BUILDER(Name("BarrierIncompleteSize").Device(DEVICE_CPU),
BarrierIncompleteSizeOp);
class BarrierReadySizeOp : public BarrierOpKernel {
public:
explicit BarrierReadySizeOp(OpKernelConstruction* context)
: BarrierOpKernel(context) {}
protected:
void ComputeAsync(OpKernelContext* ctx, Barrier* barrier,
DoneCallback callback) override {
Tensor* Tsize = nullptr;
OP_REQUIRES_OK_ASYNC(ctx, ctx->allocate_output(0, TensorShape({}), &Tsize),
callback);
Tsize->scalar<int32_t>().setConstant(barrier->ready_size());
callback();
}
};
REGISTER_KERNEL_BUILDER(Name("BarrierReadySize").Device(DEVICE_CPU),
BarrierReadySizeOp);
} // namespace barrier
} // namespace tensorflow
tensorflow/core/kernels/base64_ops.cc
+77
View File
@@ -0,0 +1,77 @@
/* Copyright 2016 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.
==============================================================================*/
// See docs in ../ops/string_ops.cc.
#include <cstdint>
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/lib/strings/base64.h"
namespace tensorflow {
namespace {
class EncodeBase64Op : public OpKernel {
public:
explicit EncodeBase64Op(OpKernelConstruction* context) : OpKernel(context) {
OP_REQUIRES_OK(context, context->GetAttr("pad", &pad_));
}
void Compute(OpKernelContext* context) override {
const Tensor& input_tensor = context->input(0);
Tensor* output_tensor = nullptr;
OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
&output_tensor));
auto input = input_tensor.flat<tstring>();
auto output = output_tensor->flat<tstring>();
for (int64_t i = 0; i < input.dimension(0); ++i) {
OP_REQUIRES_OK(context, Base64Encode(input(i), pad_, &output(i)));
}
}
private:
bool pad_;
};
REGISTER_KERNEL_BUILDER(Name("EncodeBase64").Device(DEVICE_CPU),
EncodeBase64Op);
class DecodeBase64Op : public OpKernel {
public:
using OpKernel::OpKernel;
void Compute(OpKernelContext* context) override {
const Tensor& input_tensor = context->input(0);
Tensor* output_tensor = nullptr;
OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
&output_tensor));
auto input = input_tensor.flat<tstring>();
auto output = output_tensor->flat<tstring>();
for (int64_t i = 0; i < input.dimension(0); ++i) {
OP_REQUIRES_OK(context, Base64Decode(input(i), &output(i)));
}
}
};
REGISTER_KERNEL_BUILDER(Name("DecodeBase64").Device(DEVICE_CPU),
DecodeBase64Op);
} // namespace
} // namespace tensorflow
tensorflow/core/kernels/basic_ops_benchmark_test.cc
+80
View File
@@ -0,0 +1,80 @@
/* Copyright 2016 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/kernel_benchmark_testlib.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session_options.h"
namespace tensorflow {
// We focus on the single thread performance of running ops.
static SessionOptions InitOptions() {
SessionOptions opts;
opts.config.set_intra_op_parallelism_threads(1);
opts.config.set_inter_op_parallelism_threads(1);
return opts;
}
static SessionOptions* GetOptions() {
static SessionOptions opts = InitOptions();
return &opts;
}
static Node* Var(Graph* g, int n) {
return test::graph::Var(g, DT_FLOAT, TensorShape({n}));
}
static Node* Zeros(Graph* g, int n) {
Tensor data(DT_FLOAT, TensorShape({n}));
data.flat<float>().setZero();
return test::graph::Constant(g, data);
}
static void MulChain(int chain_length, Graph** init_g, Graph** run_g) {
{
Graph* g = new Graph(OpRegistry::Global());
auto var = Var(g, 1);
test::graph::Assign(g, var, Zeros(g, 1));
*init_g = g;
}
{
Graph* g = new Graph(OpRegistry::Global());
Node* cur = Var(g, 1);
for (int i = 0; i < chain_length - 1; ++i) {
cur = test::graph::Multi(g, i % 2 == 0 ? "Div" : "Mul", {cur, cur});
}
*run_g = g;
}
}
// Benchmark a chain of simple multiplications.
// This emphasizes per-op overhead.
static void BM_MulChain(::testing::benchmark::State& state) {
const int chain_length = state.range(0);
Graph* init;
Graph* run;
MulChain(chain_length, &init, &run);
test::Benchmark("cpu", run, GetOptions(), init, nullptr, "",
/*old_benchmark_api=*/false)
.Run(state);
state.SetItemsProcessed(state.iterations());
}
BENCHMARK(BM_MulChain)->Arg(1 << 10);
} // end namespace tensorflow
tensorflow/core/kernels/batch_kernel_test_util.cc
+67
View File
@@ -0,0 +1,67 @@
/* 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/kernels/batch_kernel_test_util.h"
#include <vector>
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/kernels/batch_kernels.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace test_util {
BatchFunctionKernelTestAccess::BatchFunctionKernelTestAccess(
const BatchFunctionKernel* kernel)
: kernel_(kernel) {}
bool BatchFunctionKernelTestAccess::enable_adaptive_batch_threads() const {
return kernel_->enable_adaptive_batch_threads_;
}
absl::Status BatchFunctionKernelTestBase::Init(bool enable_adaptive_scheduler) {
std::vector<DataType> input_dtypes({DataType::DT_INT64, DataType::DT_INT64});
std::vector<NodeDefBuilder::NodeOut> inputs(
{NodeDefBuilder::NodeOut({"n1", 0, DataType::DT_INT64}),
NodeDefBuilder::NodeOut({"n2", 1, DataType::DT_INT64})});
NameAttrList f;
f.set_name("func_to_batch");
TF_CHECK_OK(NodeDefBuilder("BatchTPUInput", "BatchFunction")
.Attr("max_batch_size", 32)
.Attr("num_batch_threads", enable_adaptive_scheduler ? 0 : 8)
.Attr("allowed_batch_sizes", {2, 4, 8})
.Attr("batch_timeout_micros", 1000)
.Attr("max_enqueued_batches", 100)
.Attr("enable_large_batch_splitting", true)
.Attr("low_priority_max_batch_size", 64)
.Attr("low_priority_batch_timeout_micros", 8000)
.Attr("low_priority_allowed_batch_sizes", {32, 64})
.Attr("low_priority_max_enqueued_batches", 1000)
.Attr("Tcaptured", std::vector<DataType>{DataType::DT_INT64})
.Attr("Tin", input_dtypes)
.Input(inputs)
.Attr("Tcaptured", std::vector<DataType>{DataType::DT_INT64})
.Input(std::vector<NodeDefBuilder::NodeOut>{
NodeDefBuilder::NodeOut({"n3", 1, DataType::DT_INT64})})
.Attr("Tout", std::vector<DataType>(4, DataType::DT_INT64))
.Attr("f", f)
.Finalize(node_def()));
return InitOp();
}
} // namespace test_util
} // namespace tensorflow
tensorflow/core/kernels/batch_kernel_test_util.h
+48
View File
@@ -0,0 +1,48 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BATCH_KERNEL_TEST_UTIL_H_
#define TENSORFLOW_CORE_KERNELS_BATCH_KERNEL_TEST_UTIL_H_
#include <gtest/gtest.h>
#include "tensorflow/core/kernels/batch_kernels.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace test_util {
// A test util for accessing private members of `BatchFunctionKernel`.
class BatchFunctionKernelTestAccess {
public:
explicit BatchFunctionKernelTestAccess(const BatchFunctionKernel* kernel);
bool enable_adaptive_batch_threads() const;
private:
const BatchFunctionKernel* const kernel_;
};
class BatchFunctionKernelTestBase : public OpsTestBase,
public ::testing::WithParamInterface<bool> {
public:
// Init test fixture with a batch kernel instance.
absl::Status Init(bool enable_adaptive_scheduler);
};
} // namespace test_util
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCH_KERNEL_TEST_UTIL_H_
tensorflow/core/kernels/batch_kernels.cc
+1231
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/kernels/batch_kernels.h
+146
View File
@@ -0,0 +1,146 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BATCH_KERNELS_H_
#define TENSORFLOW_CORE_KERNELS_BATCH_KERNELS_H_
#include <cstdint>
#include <string>
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "xla/tsl/platform/types.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
// Per-model inflight batches parameters.
ABSL_CONST_INIT extern const int64_t kMinInflightBatches;
ABSL_CONST_INIT extern const int64_t kInitialInflightBatches;
ABSL_CONST_INIT extern const int64_t kBatchesToAverageOver;
ABSL_CONST_INIT extern const int64_t kMaxInflightBatches;
namespace test_util {
class BatchFunctionKernelTestAccess;
} // namespace test_util
// Records the usage of attribute `enable_large_batch_splitting`.
void RecordBatchSplitUsage(
std::optional<bool> maybe_enable_large_batch_splitting,
absl::string_view model_name);
// Records the number of batch threads of a model.
void RecordBatchParamNumBatchThreads(int64_t num_batch_threads,
absl::string_view model_name);
// Returns the model name from the context.
absl::string_view GetModelName(OpKernelContext* ctx);
// `BatchFunctionKernel` is the implementation of op `BatchFunction`.
//
// `BatchFunctionKernel` will batch (tensor) inputs by concatenating them
// along the 0-th dimension, schedule a user-defined computation, and then
// splits the returned tensors as batch output.
//
// In particular, an instance of `BatchFunctionKernel` creates or re-uses a
// a batch scheduler instance based on op attributes, pre-processes and enqueues
// concatenated inputs to the scheduler which invokes user-defined function,
// and then splits function output as op output.
//
// User defined function is named by attribute `f` and defined in the graph.
class BatchFunctionKernel : public AsyncOpKernel {
public:
explicit BatchFunctionKernel(OpKernelConstruction* c);
bool IsExpensive() override;
void ComputeAsync(OpKernelContext* c, DoneCallback done) final;
private:
friend class test_util::BatchFunctionKernelTestAccess;
// Validates 'allowed_batch_sizes_'. The entries must increase monotonically.
// If large batch split is not enabled, the last one must equal
// `max_batch_size_`. otherwise the last element must be smaller than or equal
// to `max_batch_size_`.
absl::Status ValidateAllowedBatchSizes() const;
// Creates the function handle if it isn't initialized yet; and re-use it
// afterwards.
absl::Status GetOrCreateFunctionHandle(
OpKernelContext* c, FunctionLibraryRuntime::Handle* handle);
// Instantiate the user-defined function and emits `handle`.
absl::Status InstantiateFunction(
OpKernelContext* c, FunctionLibraryRuntime::Handle* handle) const;
// Initialize vars by reading from op-kernel-construction.
// Vars
// - enable_adaptive_batch_threads_
// true if value of attribute `kEnableAdaptiveSchedulerAttr` is true, or
// if `num_batch_threads` is not positive.
// - adaptive_batch_scheduler_options_
// Read from corresponding attributes as long as they are set.
void SetAdaptiveBatchSchedulerOptions(OpKernelConstruction* c,
int32_t num_batch_threads);
std::string container_;
std::string shared_name_;
std::string batcher_queue_;
int32_t num_batch_threads_;
int32_t max_batch_size_;
int32_t batch_timeout_micros_;
int32_t max_enqueued_batches_;
std::vector<int32_t> allowed_batch_sizes_;
int32_t low_priority_max_batch_size_;
int32_t low_priority_batch_timeout_micros_;
int32_t low_priority_max_enqueued_batches_;
std::vector<int32_t> low_priority_allowed_batch_sizes_;
std::string mixed_priority_policy_;
std::string batch_padding_policy_;
int32_t num_warmup_batch_threads_ = 0;
NameAttrList func_;
absl::optional<FunctionLibraryRuntime::Handle> fhandle_ TF_GUARDED_BY(mu_);
bool enable_large_batch_splitting_ = false;
bool has_attribute_enable_large_batch_splitting_ = false;
bool enable_priority_aware_batch_scheduler_ = false;
bool enable_priority_aware_batch_scheduler_resplit_ = false;
// If true, the priority-aware batch scheduler will lazily filter out and
// cancel tasks that have been cancelled or have exceeded their deadline
// before batch formation.
bool enable_batching_task_lazy_cancellation_ = false;
bool enable_adaptive_batch_threads_ = false;
mutex mu_;
// Parameters for adaptive batch scheduler only.
// Note 'num_batch_threads_' above is shared by two implementations of batch
// scheduler.
struct AdaptiveBatchSchedulerOptions {
int32_t min_in_flight_batches_limit = kMinInflightBatches;
int32_t initial_in_flight_batches_limit = kInitialInflightBatches;
int32_t max_in_flight_batches_limit = kMaxInflightBatches;
int32_t batches_to_average_over = kBatchesToAverageOver;
int64_t full_batch_scheduling_boost_micros = -1;
};
absl::optional<AdaptiveBatchSchedulerOptions>
adaptive_batch_scheduler_options_ = std::nullopt;
};
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCH_KERNELS_H_
tensorflow/core/kernels/batch_kernels_auto_warmup_test.cc
+283
View File
@@ -0,0 +1,283 @@
/* 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 <cstdint>
#include <memory>
#include <utility>
#include <vector>
#include <gtest/gtest.h>
#include "absl/strings/match.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/common_runtime/rendezvous_mgr.h"
#include "tensorflow/core/framework/device_factory.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/kernels/batch_kernel_test_util.h"
#include "tensorflow/core/kernels/batch_kernels.h"
#include "tensorflow/core/kernels/batching_util/warmup.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/public/version.h"
#include "tsl/platform/blocking_counter.h"
#include "tsl/platform/refcount.h"
namespace tensorflow {
namespace {
using PerModelData = serving::WarmupStateRegistry::PerModelData;
class BatchFunctionKernelTest : public test_util::BatchFunctionKernelTestBase {
};
class BatchFunctionKernelParallelWarmupTestState : public OpsTestBase {
public:
// Init test fixture with a batch kernel instance.
absl::Status Init(bool enable_splitting, bool check_output_shape) {
static auto *const cpu_device = []() {
auto device =
DeviceFactory::NewDevice("CPU", {}, "/job:a/replica:0/task:0");
return device.release();
}();
// Override the per-test/per-op device with a global device so that it can
// be shared between ops.
device_ = cpu_device;
NameAttrList f;
f.set_name("BatchFunctionKernelParallelWarmupTestStateFunc");
FunctionDef func;
if (check_output_shape) {
func = FunctionDefHelper::Create(
// function_name
f.name(),
// in_def
{"x:int64"},
// out_def
{"o:int64"},
// attr_def
{},
// node_def
{{{"o"},
"EnsureShape",
{"x"},
{{"T", DataType::DT_INT64}, {"shape", TensorShape({2})}}}},
// ret_def
{{"o", "o:output"}});
} else {
func = FunctionDefHelper::Create(
// function_name
f.name(),
// in_def
{"x:int64"},
// out_def
{"o:int64"},
// attr_def
{},
// node_def
{{{"o"}, "Identity", {"x"}, {{"T", DataType::DT_INT64}}}},
// ret_def
{{"o", "o:output"}});
}
TF_RETURN_IF_ERROR(flib_def_->AddFunctionDef(func));
pflr_ = std::make_unique<ProcessFunctionLibraryRuntime>(
device_mgr_.get(), Env::Default(), /*config=*/nullptr,
TF_GRAPH_DEF_VERSION, flib_def_.get(), OptimizerOptions(),
/*thread_pool=*/nullptr, /*parent=*/nullptr,
/*session_metadata=*/nullptr,
Rendezvous::Factory{[](const int64_t, const DeviceMgr *device_mgr,
tsl::core::RefCountPtr<Rendezvous> *r) {
*r = tsl::core::RefCountPtr<Rendezvous>(
new IntraProcessRendezvous(device_mgr));
return absl::OkStatus();
}});
std::vector<NodeDefBuilder::NodeOut> inputs(
{NodeDefBuilder::NodeOut({"n1", 0, DataType::DT_INT64})});
TF_CHECK_OK(NodeDefBuilder("BatchTPUInput", "BatchFunction")
.Attr("max_batch_size", enable_splitting ? 16 : 8)
.Attr("num_batch_threads", 8)
.Attr("allowed_batch_sizes", {2, 4, 8})
.Attr("batch_timeout_micros", 1000)
.Attr("max_enqueued_batches", 10)
.Attr("enable_large_batch_splitting", true)
.Attr("low_priority_max_batch_size", 64)
.Attr("low_priority_batch_timeout_micros", 8000)
.Attr("low_priority_allowed_batch_sizes", {32, 64})
.Attr("low_priority_max_enqueued_batches", 1000)
.Attr("Tin", {DataType::DT_INT64})
.Input(inputs)
.Attr("Tcaptured", std::vector<DataType>{})
.Input(std::vector<NodeDefBuilder::NodeOut>{})
.Attr("Tout", std::vector<DataType>{DT_INT64})
.Attr("f", f)
.Finalize(node_def()));
return InitOp();
}
void TestBody() override {}
};
class BatchFunctionKernelParallelWarmupTest
: public ::testing::TestWithParam<bool> {};
TEST_P(BatchFunctionKernelParallelWarmupTest, HandlesLargeBatchSplitting) {
// This test fails if it does not come before the others in the suite,
// because `SharedBatchScheduler::QueueOptions::input_batch_size_limit`
// does not get reset.
SessionMetadata session_metadata;
session_metadata.set_name("test_model");
session_metadata.set_version(123);
serving::WarmupStateRegistry::Key key(session_metadata.name(),
session_metadata.version());
int num_requests = 16;
{
auto per_model_data = std::make_unique<PerModelData>();
per_model_data->warmup_all_batch_sizes = true;
auto handle = serving::GetGlobalWarmupStateRegistry().Register(
key, std::move(per_model_data));
tsl::BlockingCounter blocking_counter(num_requests);
for (int i = 0; i < num_requests; ++i) {
Env::Default()->SchedClosure([&]() {
BatchFunctionKernelParallelWarmupTestState test;
test.set_session_metadata(session_metadata);
TF_CHECK_OK(test.Init(/*enable_splitting=*/true,
/*check_output_shape=*/true));
test.AddInputFromList<int64_t>(
TensorShape({16}),
{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15});
auto status = test.RunOpKernel();
ASSERT_FALSE(status.ok());
// This proves the kernel is executed with batch sizes other than 2.
EXPECT_TRUE(absl::StrContains(status.message(),
"is not compatible with expected shape"));
blocking_counter.DecrementCount();
});
}
blocking_counter.Wait();
}
{
EXPECT_FALSE(serving::GetGlobalWarmupStateRegistry().Lookup(key));
auto per_model_data = std::make_unique<PerModelData>();
per_model_data->warmup_all_batch_sizes = true;
auto handle = serving::GetGlobalWarmupStateRegistry().Register(
key, std::move(per_model_data));
tsl::BlockingCounter blocking_counter(num_requests);
for (int i = 0; i < num_requests; ++i) {
Env::Default()->SchedClosure([&]() {
BatchFunctionKernelParallelWarmupTestState test;
test.set_session_metadata(session_metadata);
// Error free when the EnsureShapeOp is replaced with an Identity op.
TF_CHECK_OK(
test.Init(/*enable_splitting=*/true, /*check_output_shape=*/false));
test.AddInputFromList<int64_t>(
TensorShape({16}),
{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15});
TF_CHECK_OK(test.RunOpKernel());
test::ExpectTensorEqual<int64_t>(
*test.GetOutput(0),
test::AsTensor<int64_t>(
{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}));
blocking_counter.DecrementCount();
});
}
blocking_counter.Wait();
}
}
TEST_P(BatchFunctionKernelParallelWarmupTest, AutoBatchWorks) {
SessionMetadata session_metadata;
session_metadata.set_name("test_model");
session_metadata.set_version(123);
serving::WarmupStateRegistry::Key key(session_metadata.name(),
session_metadata.version());
int num_requests = 16;
bool enable_splitting = GetParam();
{
auto per_model_data = std::make_unique<PerModelData>();
per_model_data->warmup_all_batch_sizes = true;
auto handle = serving::GetGlobalWarmupStateRegistry().Register(
key, std::move(per_model_data));
tsl::BlockingCounter blocking_counter(num_requests);
for (int i = 0; i < num_requests; ++i) {
Env::Default()->SchedClosure([&]() {
BatchFunctionKernelParallelWarmupTestState test;
test.set_session_metadata(session_metadata);
TF_CHECK_OK(test.Init(enable_splitting, /*check_output_shape=*/true));
test.AddInputFromList<int64_t>(TensorShape({2}), {123, 456});
auto status = test.RunOpKernel();
ASSERT_FALSE(status.ok());
// This proves the kernel is executed with batch sizes other than 2.
EXPECT_TRUE(absl::StrContains(status.message(),
"is not compatible with expected shape"));
blocking_counter.DecrementCount();
});
}
blocking_counter.Wait();
}
{
EXPECT_FALSE(serving::GetGlobalWarmupStateRegistry().Lookup(key));
auto per_model_data = std::make_unique<PerModelData>();
per_model_data->warmup_all_batch_sizes = true;
auto handle = serving::GetGlobalWarmupStateRegistry().Register(
key, std::move(per_model_data));
tsl::BlockingCounter blocking_counter(num_requests);
for (int i = 0; i < num_requests; ++i) {
Env::Default()->SchedClosure([&]() {
BatchFunctionKernelParallelWarmupTestState test;
test.set_session_metadata(session_metadata);
// Error free when the EnsureShapeOp is replaced with an Identity op.
TF_CHECK_OK(test.Init(enable_splitting, /*check_output_shape=*/false));
test.AddInputFromList<int64_t>(TensorShape({2}), {123, 456});
auto status = test.RunOpKernel();
TF_CHECK_OK(test.RunOpKernel());
test::ExpectTensorEqual<int64_t>(*test.GetOutput(0),
test::AsTensor<int64_t>({123, 456}));
blocking_counter.DecrementCount();
});
}
blocking_counter.Wait();
}
}
INSTANTIATE_TEST_SUITE_P(BatchFunctionKernelParallelWarmupTestSuite,
BatchFunctionKernelParallelWarmupTest,
::testing::Bool());
} // namespace
} // namespace tensorflow
tensorflow/core/kernels/batch_kernels_env_test.cc
+49
View File
@@ -0,0 +1,49 @@
/* 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 <gmock/gmock.h>
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/kernels/batch_kernel_test_util.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
namespace tensorflow {
namespace {
// Tests that batch kernel initialization returns error when it's configured to
// use adaptive scheduling yet batching thread pool creation fails.
class BatchFunctionKernelEnvTest
: public test_util::BatchFunctionKernelTestBase {};
TEST_P(BatchFunctionKernelEnvTest, Basic) {
tensorflow::setenv("TF_NUM_BATCH_THREADS", "0", 1 /* overwrite */);
const bool adaptive_scheduler_enabled = GetParam();
absl::Status status = Init(adaptive_scheduler_enabled);
if (adaptive_scheduler_enabled) {
EXPECT_THAT(status,
absl_testing::StatusIs(error::FAILED_PRECONDITION,
"Failed to create batch threads pool"));
} else {
// Initialization is ok since batch kernel doesn't use adaptive
// scheduler.
TF_EXPECT_OK(status);
}
}
INSTANTIATE_TEST_SUITE_P(Params, BatchFunctionKernelEnvTest, ::testing::Bool());
} // namespace
} // namespace tensorflow
tensorflow/core/kernels/batch_kernels_test.cc
+750
View File
@@ -0,0 +1,750 @@
/* 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/kernels/batch_kernels.h"
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/criticality.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/status.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/common_runtime/rendezvous_mgr.h"
#include "tensorflow/core/framework/device_factory.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/kernels/batch_kernel_test_util.h"
#include "tensorflow/core/kernels/batching_util/batch_scheduler.h"
#include "tensorflow/core/kernels/batching_util/warmup.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/public/version.h"
#include "tsl/platform/blocking_counter.h"
#include "tsl/platform/refcount.h"
namespace tensorflow {
namespace {
using PerModelData = serving::WarmupStateRegistry::PerModelData;
class BatchFunctionKernelTest : public test_util::BatchFunctionKernelTestBase {
};
TEST_P(BatchFunctionKernelTest, EnableAdaptiveScheduler) {
const bool adaptive_scheduler_enabled = GetParam();
TF_EXPECT_OK(Init(adaptive_scheduler_enabled));
BatchFunctionKernel *batch_kernel =
dynamic_cast<BatchFunctionKernel *>(op_kernel());
EXPECT_EQ(adaptive_scheduler_enabled,
test_util::BatchFunctionKernelTestAccess(batch_kernel)
.enable_adaptive_batch_threads());
}
INSTANTIATE_TEST_SUITE_P(Params, BatchFunctionKernelTest, ::testing::Bool());
class SharedBatchFunctionTestState : public OpsTestBase {
public:
// Init test fixture with a batch kernel instance.
void CreateFunctionLibraryRuntime() {
pflr_ = std::make_unique<ProcessFunctionLibraryRuntime>(
device_mgr_.get(), Env::Default(), /*config=*/nullptr,
TF_GRAPH_DEF_VERSION, flib_def_.get(), OptimizerOptions(),
/*thread_pool=*/nullptr, /*parent=*/nullptr,
/*session_metadata=*/nullptr,
Rendezvous::Factory{[](const int64_t, const DeviceMgr *device_mgr,
tsl::core::RefCountPtr<Rendezvous> *r) {
*r = tsl::core::RefCountPtr<Rendezvous>(
new IntraProcessRendezvous(device_mgr));
return absl::OkStatus();
}});
}
protected:
// Create common batch function op for testing.
absl::StatusOr<NodeDefBuilder> CreateBatchFunctionBuilder(
const std::vector<int> &allowed_batch_sizes, int max_batch_size,
absl::string_view padding_policy,
const TensorShape &expected_output_shape) {
NameAttrList f;
f.set_name("ShapeEnforcingFunction");
FunctionDef func = FunctionDefHelper::Create(
// function_name
f.name(),
// in_def
{"x:int64"},
// out_def
{"o:int64"},
// attr_def
{},
// node_def
{{{"o"},
"EnsureShape",
{"x"},
{{"T", DataType::DT_INT64}, {"shape", expected_output_shape}}}},
// ret_def
{{"o", "o:output"}});
TF_RETURN_IF_ERROR(flib_def_->AddFunctionDef(func));
SharedBatchFunctionTestState::CreateFunctionLibraryRuntime();
std::vector<NodeDefBuilder::NodeOut> inputs(
{NodeDefBuilder::NodeOut({"n1", 0, DataType::DT_INT64})});
return NodeDefBuilder(absl::StrCat("BatchTPUInput", padding_policy),
"BatchFunction")
.Attr("max_batch_size", max_batch_size)
.Attr("num_batch_threads", 8)
.Attr("allowed_batch_sizes", allowed_batch_sizes)
.Attr("batch_timeout_micros", 1000000)
.Attr("max_enqueued_batches", 10)
.Attr("enable_large_batch_splitting", true)
.Attr("batch_padding_policy", padding_policy)
.Attr("Tin", {DataType::DT_INT64})
.Input(inputs)
.Attr("Tcaptured", std::vector<DataType>{})
.Input(std::vector<NodeDefBuilder::NodeOut>{})
.Attr("Tout", std::vector<DataType>{DT_INT64})
.Attr("f", f);
}
};
class BatchFunctionTestState : public SharedBatchFunctionTestState {
public:
// Init test fixture with a batch kernel instance. The caller guarantees that
// the device pointer is valid throughout the life of this class.
absl::Status Init(Device *device, bool enable_low_priority_queue,
absl::string_view mixed_priority_policy,
int64_t expected_batch_size) {
// Override the per-test/per-op device with a given device so that it can
// be shared between ops.
device_ = device;
const TensorShape expected_output_shape({expected_batch_size, 2});
TF_ASSIGN_OR_RETURN(
NodeDefBuilder builder,
CreateBatchFunctionBuilder({4, 8}, 8, "PAD_UP", expected_output_shape));
TF_RETURN_IF_ERROR(builder
.Attr("low_priority_max_batch_size",
enable_low_priority_queue ? 8 : 0)
.Attr("low_priority_batch_timeout_micros",
enable_low_priority_queue ? 2000000 : 0)
.Attr("low_priority_allowed_batch_sizes",
enable_low_priority_queue
? std::vector<int>{4, 8}
: std::vector<int>())
.Attr("low_priority_max_enqueued_batches",
enable_low_priority_queue ? 2 : 0)
.Attr("mixed_priority_policy", mixed_priority_policy)
.Finalize(node_def()));
return OpsTestBase::InitOp();
}
void TestBody() override {}
};
class BatchFunctionTest : public ::testing::TestWithParam<bool> {
protected:
void SetUp() override {
// The device needs to be shared in each test case and within each test case
// only.
cpu_device_ =
DeviceFactory::NewDevice("CPU", {}, "/job:a/replica:0/task:0");
}
std::unique_ptr<Device> cpu_device_;
};
TEST_P(BatchFunctionTest, BatchingWorksWithoutCriticality) {
SessionMetadata session_metadata;
session_metadata.set_name("test_model");
session_metadata.set_version(123);
bool enable_low_priority_queue = GetParam();
{
tsl::BlockingCounter blocking_counter(8);
// 8 threads run the batch op with no explicit criticality set. They are
// eventually batched to form a tensor with [8, 2] shape which is verified
// within the function.
for (int i = 0; i < 8; ++i) {
Env::Default()->SchedClosure([&]() {
ASSERT_EQ(tsl::criticality::GetCriticality(),
tsl::criticality::Criticality::kCritical);
BatchFunctionTestState test_state;
test_state.set_session_metadata(session_metadata);
TF_ASSERT_OK(test_state.Init(
cpu_device_.get(), enable_low_priority_queue,
serving::kLowPriorityPaddingWithMaxBatchSizeAttrValue,
/*expected_batch_size=*/8));
test_state.AddInputFromList<int64_t>(TensorShape({1, 2}), {123, 456});
TF_EXPECT_OK(test_state.RunOpKernel());
test::ExpectTensorEqual<int64_t>(
*test_state.GetOutput(0),
test::AsTensor<int64_t>({123, 456}, TensorShape({1, 2})));
blocking_counter.DecrementCount();
});
}
blocking_counter.Wait();
}
}
TEST_P(BatchFunctionTest, PaddingWorksWithoutCriticality) {
SessionMetadata session_metadata;
session_metadata.set_name("test_model");
session_metadata.set_version(123);
bool enable_low_priority_queue = GetParam();
{
tsl::BlockingCounter blocking_counter(2);
// 2 threads run the batch op with no explicit criticality set. They are
// eventually batched and padded to form a tensor with [4, 2] shape which is
// verified within the function.
for (int i = 0; i < 2; ++i) {
Env::Default()->SchedClosure([&]() {
ASSERT_EQ(tsl::criticality::GetCriticality(),
tsl::criticality::Criticality::kCritical);
BatchFunctionTestState test_state;
test_state.set_session_metadata(session_metadata);
TF_ASSERT_OK(test_state.Init(
cpu_device_.get(), enable_low_priority_queue,
serving::kLowPriorityPaddingWithMaxBatchSizeAttrValue,
/*expected_batch_size=*/4));
test_state.AddInputFromList<int64_t>(TensorShape({1, 2}), {123, 456});
TF_EXPECT_OK(test_state.RunOpKernel());
test::ExpectTensorEqual<int64_t>(
*test_state.GetOutput(0),
test::AsTensor<int64_t>({123, 456}, TensorShape({1, 2})));
blocking_counter.DecrementCount();
});
}
blocking_counter.Wait();
}
}
#if defined(PLATFORM_GOOGLE)
TEST_P(BatchFunctionTest,
LowPriorityTaskPaddingHighPriorityBatchUptoMaxBatchSize) {
SessionMetadata session_metadata;
session_metadata.set_name("test_model");
session_metadata.set_version(123);
bool enable_low_priority_queue = GetParam();
{
tsl::BlockingCounter blocking_counter(8);
// 4 threads run the batch op with critical plus and 4 threads run the batch
// op with sheddable. They are eventually batched to form a tensor with [8,
// 2] shape which is verified within the function.
for (int i = 0; i < 4; ++i) {
Env::Default()->SchedClosure([&]() {
tsl::criticality::ScopedCriticality scoped_criticality(
tsl::criticality::Criticality::kCriticalPlus);
ASSERT_EQ(tsl::criticality::GetCriticality(),
tsl::criticality::Criticality::kCriticalPlus);
BatchFunctionTestState test_state;
test_state.set_session_metadata(session_metadata);
TF_ASSERT_OK(test_state.Init(
cpu_device_.get(), enable_low_priority_queue,
serving::kLowPriorityPaddingWithMaxBatchSizeAttrValue,
/*expected_batch_size=*/8));
test_state.AddInputFromList<int64_t>(TensorShape({1, 2}), {123, 456});
TF_EXPECT_OK(test_state.RunOpKernel());
test::ExpectTensorEqual<int64_t>(
*test_state.GetOutput(0),
test::AsTensor<int64_t>({123, 456}, TensorShape({1, 2})));
blocking_counter.DecrementCount();
});
}
for (int i = 0; i < 4; ++i) {
Env::Default()->SchedClosure([&]() {
tsl::criticality::ScopedCriticality scoped_criticality(
tsl::criticality::Criticality::kSheddable);
ASSERT_EQ(tsl::criticality::GetCriticality(),
tsl::criticality::Criticality::kSheddable);
BatchFunctionTestState test_state;
test_state.set_session_metadata(session_metadata);
TF_ASSERT_OK(test_state.Init(
cpu_device_.get(), enable_low_priority_queue,
serving::kLowPriorityPaddingWithMaxBatchSizeAttrValue,
/*expected_batch_size=*/8));
test_state.AddInputFromList<int64_t>(TensorShape({1, 2}), {234, 567});
TF_EXPECT_OK(test_state.RunOpKernel());
test::ExpectTensorEqual<int64_t>(
*test_state.GetOutput(0),
test::AsTensor<int64_t>({234, 567}, TensorShape({1, 2})));
blocking_counter.DecrementCount();
});
}
blocking_counter.Wait();
}
}
TEST_P(BatchFunctionTest,
LowPriorityTaskPaddingHighPriorityBatchWithExtraPadding) {
SessionMetadata session_metadata;
session_metadata.set_name("test_model");
session_metadata.set_version(123);
bool enable_low_priority_queue = GetParam();
{
tsl::BlockingCounter blocking_counter(2);
// 1 thread run the batch op with critical plus and 1 threads run the batch
// op with sheddable. They are eventually batched and padded to form a
// tensor with [4, 2] shape which is verified within the function.
Env::Default()->SchedClosure([&]() {
tsl::criticality::ScopedCriticality scoped_criticality(
tsl::criticality::Criticality::kCriticalPlus);
ASSERT_EQ(tsl::criticality::GetCriticality(),
tsl::criticality::Criticality::kCriticalPlus);
BatchFunctionTestState test_state;
test_state.set_session_metadata(session_metadata);
TF_ASSERT_OK(
test_state.Init(cpu_device_.get(), enable_low_priority_queue,
serving::kLowPriorityPaddingWithMaxBatchSizeAttrValue,
/*expected_batch_size=*/4));
test_state.AddInputFromList<int64_t>(TensorShape({1, 2}), {123, 456});
TF_EXPECT_OK(test_state.RunOpKernel());
test::ExpectTensorEqual<int64_t>(
*test_state.GetOutput(0),
test::AsTensor<int64_t>({123, 456}, TensorShape({1, 2})));
blocking_counter.DecrementCount();
});
Env::Default()->SchedClosure([&]() {
tsl::criticality::ScopedCriticality scoped_criticality(
tsl::criticality::Criticality::kSheddable);
ASSERT_EQ(tsl::criticality::GetCriticality(),
tsl::criticality::Criticality::kSheddable);
BatchFunctionTestState test_state;
test_state.set_session_metadata(session_metadata);
TF_ASSERT_OK(
test_state.Init(cpu_device_.get(), enable_low_priority_queue,
serving::kLowPriorityPaddingWithMaxBatchSizeAttrValue,
/*expected_batch_size=*/4));
test_state.AddInputFromList<int64_t>(TensorShape({1, 2}), {234, 567});
TF_EXPECT_OK(test_state.RunOpKernel());
test::ExpectTensorEqual<int64_t>(
*test_state.GetOutput(0),
test::AsTensor<int64_t>({234, 567}, TensorShape({1, 2})));
blocking_counter.DecrementCount();
});
blocking_counter.Wait();
}
}
TEST_P(BatchFunctionTest,
LowPriorityTaskPaddingHighPriorityBatchUptoNextAllowedBatchSize) {
SessionMetadata session_metadata;
session_metadata.set_name("test_model");
session_metadata.set_version(123);
bool enable_low_priority_queue = GetParam();
{
tsl::BlockingCounter blocking_counter(4);
// 2 threads run the batch op with critical plus and 2 threads run the batch
// op with sheddable. They are eventually batched to form a tensor with [4,
// 2] shape which is verified within the function.
for (int i = 0; i < 2; ++i) {
Env::Default()->SchedClosure([&]() {
tsl::criticality::ScopedCriticality scoped_criticality(
tsl::criticality::Criticality::kCriticalPlus);
ASSERT_EQ(tsl::criticality::GetCriticality(),
tsl::criticality::Criticality::kCriticalPlus);
BatchFunctionTestState test_state;
test_state.set_session_metadata(session_metadata);
TF_ASSERT_OK(test_state.Init(
cpu_device_.get(), enable_low_priority_queue,
serving::kLowPriorityPaddingWithNextAllowedBatchSizeAttrValue,
/*expected_batch_size=*/4));
test_state.AddInputFromList<int64_t>(TensorShape({1, 2}), {123, 456});
TF_EXPECT_OK(test_state.RunOpKernel());
test::ExpectTensorEqual<int64_t>(
*test_state.GetOutput(0),
test::AsTensor<int64_t>({123, 456}, TensorShape({1, 2})));
blocking_counter.DecrementCount();
});
}
for (int i = 0; i < 2; ++i) {
Env::Default()->SchedClosure([&]() {
tsl::criticality::ScopedCriticality scoped_criticality(
tsl::criticality::Criticality::kSheddable);
ASSERT_EQ(tsl::criticality::GetCriticality(),
tsl::criticality::Criticality::kSheddable);
BatchFunctionTestState test_state;
test_state.set_session_metadata(session_metadata);
TF_ASSERT_OK(test_state.Init(
cpu_device_.get(), enable_low_priority_queue,
serving::kLowPriorityPaddingWithNextAllowedBatchSizeAttrValue,
/*expected_batch_size=*/4));
test_state.AddInputFromList<int64_t>(TensorShape({1, 2}), {234, 567});
TF_EXPECT_OK(test_state.RunOpKernel());
test::ExpectTensorEqual<int64_t>(
*test_state.GetOutput(0),
test::AsTensor<int64_t>({234, 567}, TensorShape({1, 2})));
blocking_counter.DecrementCount();
});
}
blocking_counter.Wait();
}
}
#endif
INSTANTIATE_TEST_SUITE_P(BatchFunctionTest, BatchFunctionTest,
::testing::Bool());
#if defined(PLATFORM_GOOGLE)
TEST_F(BatchFunctionTest, HighPriorityBatchNotPaddedWithLowPriorityTasks) {
SessionMetadata session_metadata;
session_metadata.set_name("test_model");
session_metadata.set_version(123);
{
tsl::BlockingCounter blocking_counter(8);
// 4 threads run the batch op with critical plus and 4 threads run the batch
// op with sheddable. They each get batched separately to form tensors with
// [4, 2] shape which is verified within the function.
for (int i = 0; i < 4; ++i) {
Env::Default()->SchedClosure([&]() {
tsl::criticality::ScopedCriticality scoped_criticality(
tsl::criticality::Criticality::kCriticalPlus);
ASSERT_EQ(tsl::criticality::GetCriticality(),
tsl::criticality::Criticality::kCriticalPlus);
BatchFunctionTestState test_state;
test_state.set_session_metadata(session_metadata);
TF_ASSERT_OK(test_state.Init(cpu_device_.get(),
/*enable_low_priority_queue=*/true,
serving::kPriorityIsolationAttrValue,
/*expected_batch_size=*/4));
test_state.AddInputFromList<int64_t>(TensorShape({1, 2}), {123, 456});
TF_EXPECT_OK(test_state.RunOpKernel());
test::ExpectTensorEqual<int64_t>(
*test_state.GetOutput(0),
test::AsTensor<int64_t>({123, 456}, TensorShape({1, 2})));
blocking_counter.DecrementCount();
});
}
for (int i = 0; i < 4; ++i) {
Env::Default()->SchedClosure([&]() {
tsl::criticality::ScopedCriticality scoped_criticality(
tsl::criticality::Criticality::kSheddable);
ASSERT_EQ(tsl::criticality::GetCriticality(),
tsl::criticality::Criticality::kSheddable);
BatchFunctionTestState test_state;
test_state.set_session_metadata(session_metadata);
TF_ASSERT_OK(test_state.Init(cpu_device_.get(),
/*enable_low_priority_queue=*/true,
serving::kPriorityIsolationAttrValue,
/*expected_batch_size=*/4));
test_state.AddInputFromList<int64_t>(TensorShape({1, 2}), {234, 567});
TF_EXPECT_OK(test_state.RunOpKernel());
test::ExpectTensorEqual<int64_t>(
*test_state.GetOutput(0),
test::AsTensor<int64_t>({234, 567}, TensorShape({1, 2})));
blocking_counter.DecrementCount();
});
}
blocking_counter.Wait();
}
}
TEST_F(BatchFunctionTest, LowPriorityOnlyBatchAtMaxLowPriorityBatchSize) {
SessionMetadata session_metadata;
session_metadata.set_name("test_model");
session_metadata.set_version(123);
{
tsl::BlockingCounter blocking_counter(8);
// 8 threads run the batch op with sheddable. They are eventually batched to
// form a tensor with [8, 2] shape, which is verified within the function,
// since the low priority max batch size is set to 8.
for (int i = 0; i < 8; ++i) {
Env::Default()->SchedClosure([&]() {
tsl::criticality::ScopedCriticality scoped_criticality(
tsl::criticality::Criticality::kSheddable);
ASSERT_EQ(tsl::criticality::GetCriticality(),
tsl::criticality::Criticality::kSheddable);
BatchFunctionTestState test_state;
test_state.set_session_metadata(session_metadata);
TF_ASSERT_OK(test_state.Init(
cpu_device_.get(),
/*enable_low_priority_queue=*/true,
serving::kLowPriorityPaddingWithMaxBatchSizeAttrValue,
/*expected_batch_size=*/8));
test_state.AddInputFromList<int64_t>(TensorShape({1, 2}), {234, 567});
TF_EXPECT_OK(test_state.RunOpKernel());
test::ExpectTensorEqual<int64_t>(
*test_state.GetOutput(0),
test::AsTensor<int64_t>({234, 567}, TensorShape({1, 2})));
blocking_counter.DecrementCount();
});
}
blocking_counter.Wait();
}
}
TEST_F(BatchFunctionTest, LowPriorityBatchPaddedToLowPriorityAllowedBatchSize) {
SessionMetadata session_metadata;
session_metadata.set_name("test_model");
session_metadata.set_version(123);
{
tsl::BlockingCounter blocking_counter(2);
// 2 threads run the batch op with sheddable. They are eventually batched
// and padded to form a tensor with [4, 2] shape, which is verified within
// the function, since the low priority allowed batch size is set to [4, 8].
for (int i = 0; i < 2; ++i) {
Env::Default()->SchedClosure([&]() {
tsl::criticality::ScopedCriticality scoped_criticality(
tsl::criticality::Criticality::kSheddable);
ASSERT_EQ(tsl::criticality::GetCriticality(),
tsl::criticality::Criticality::kSheddable);
BatchFunctionTestState test_state;
test_state.set_session_metadata(session_metadata);
TF_ASSERT_OK(test_state.Init(
cpu_device_.get(),
/*enable_low_priority_queue=*/true,
serving::kLowPriorityPaddingWithMaxBatchSizeAttrValue,
/*expected_batch_size=*/4));
test_state.AddInputFromList<int64_t>(TensorShape({1, 2}), {234, 567});
TF_EXPECT_OK(test_state.RunOpKernel());
test::ExpectTensorEqual<int64_t>(
*test_state.GetOutput(0),
test::AsTensor<int64_t>({234, 567}, TensorShape({1, 2})));
blocking_counter.DecrementCount();
});
}
blocking_counter.Wait();
}
}
#endif
class BatchFunctionKernelParallelWarmupTestState
: public SharedBatchFunctionTestState {
public:
// Init test fixture with a batch kernel instance.
absl::Status Init(bool enable_splitting) {
static auto *const cpu_device = []() {
auto device =
DeviceFactory::NewDevice("CPU", {}, "/job:a/replica:0/task:0");
return device.release();
}();
// Override the per-test/per-op device with a global device so that it can
// be shared between ops.
device_ = cpu_device;
const TensorShape expected_output_shape({2});
TF_ASSIGN_OR_RETURN(
NodeDefBuilder builder,
CreateBatchFunctionBuilder({2, 4, 8}, enable_splitting ? 16 : 8,
"PAD_UP", expected_output_shape));
TF_RETURN_IF_ERROR(builder.Finalize(node_def()));
return OpsTestBase::InitOp();
}
void TestBody() override {}
};
class BatchFunctionKernelParallelWarmupTest
: public ::testing::TestWithParam<bool> {};
TEST_P(BatchFunctionKernelParallelWarmupTest, ParallelWarmup) {
SessionMetadata session_metadata;
session_metadata.set_name("test_model");
session_metadata.set_version(123);
serving::WarmupStateRegistry::Key key(session_metadata.name(),
session_metadata.version());
int num_requests = 16;
bool enable_splitting = GetParam();
{
// Setting the state to warmup disables batching in the BatchFunction op. We
// are checking this behavior by checking the tensor shape inside batch
// function is the same as the input tensor shape using EnsureShape op.
auto per_model_data = std::make_unique<PerModelData>();
auto handle = serving::GetGlobalWarmupStateRegistry().Register(
key, std::move(per_model_data));
tsl::BlockingCounter blocking_counter(num_requests);
for (int i = 0; i < num_requests; ++i) {
Env::Default()->SchedClosure([&]() {
BatchFunctionKernelParallelWarmupTestState test;
test.set_session_metadata(session_metadata);
TF_CHECK_OK(test.Init(enable_splitting));
test.AddInputFromList<int64_t>(TensorShape({2}), {123, 456});
TF_CHECK_OK(test.RunOpKernel());
test::ExpectTensorEqual<int64_t>(*test.GetOutput(0),
test::AsTensor<int64_t>({123, 456}));
blocking_counter.DecrementCount();
});
}
// Note this times out after 60s, so `batch_timeout_micros` and `batch_size`
// need to be set accordingly.
blocking_counter.Wait();
}
EXPECT_FALSE(serving::GetGlobalWarmupStateRegistry().Lookup(key));
{
tsl::BlockingCounter blocking_counter(num_requests);
for (int i = 0; i < num_requests; ++i) {
Env::Default()->SchedClosure([&]() {
BatchFunctionKernelParallelWarmupTestState test;
test.set_session_metadata(session_metadata);
TF_CHECK_OK(test.Init(enable_splitting));
test.AddInputFromList<int64_t>(TensorShape({2}), {123, 456});
// We expect requests to be batched together when the warm-up mode is
// turned off, which will make the execution fail at `EnsureShape`.
EXPECT_FALSE(test.RunOpKernel().ok());
blocking_counter.DecrementCount();
});
}
blocking_counter.Wait();
}
}
INSTANTIATE_TEST_SUITE_P(BatchFunctionKernelParallelWarmupTestSuite,
BatchFunctionKernelParallelWarmupTest,
::testing::Bool());
class BatchFunctionKernelPaddingTestState
: public SharedBatchFunctionTestState {
public:
// Init test fixture with a batch kernel instance.
absl::Status Init(absl::string_view padding_policy, int expected_batch_size) {
static auto *const cpu_device = []() {
auto device =
DeviceFactory::NewDevice("CPU", {}, "/job:a/replica:0/task:0");
return device.release();
}();
// Override the per-test/per-op device with a global device so that it can
// be shared between ops.
device_ = cpu_device;
const TensorShape expected_output_shape({expected_batch_size, 2});
TF_RETURN_IF_ERROR(CreateBatchFunctionBuilder({4, 8}, 8, padding_policy,
expected_output_shape)
->Finalize(node_def()));
return OpsTestBase::InitOp();
}
void TestBody() override {}
};
class BatchFunctionKernelPaddingTest
: public ::testing::TestWithParam<std::string> {};
TEST_P(BatchFunctionKernelPaddingTest, PadUp) {
SessionMetadata session_metadata;
session_metadata.set_name("test_model");
session_metadata.set_version(123);
// Send 5 requests in parallel and check that the given batch padding
// policy behaves as expected.
int64_t num_requests = 5;
int64_t expected_batch_size = 0;
std::string padding_policy = GetParam();
if (padding_policy == "PAD_UP") {
expected_batch_size = 8;
} else if (padding_policy == "BATCH_DOWN") {
expected_batch_size = 4;
} else if (padding_policy == "MINIMIZE_TPU_COST_PER_REQUEST") {
expected_batch_size = 8;
} else {
FAIL() << "Unsupported padding policy: " << padding_policy;
}
{
tsl::BlockingCounter blocking_counter(num_requests);
for (int i = 0; i < num_requests; ++i) {
Env::Default()->SchedClosure([&]() {
BatchFunctionKernelPaddingTestState test_state;
test_state.set_session_metadata(session_metadata);
TF_CHECK_OK(test_state.Init(padding_policy, expected_batch_size));
test_state.AddInputFromList<int64_t>(TensorShape({1, 2}), {123, 456});
TF_EXPECT_OK(test_state.RunOpKernel());
test::ExpectTensorEqual<int64_t>(
*test_state.GetOutput(0),
test::AsTensor<int64_t>({123, 456}, TensorShape({1, 2})));
blocking_counter.DecrementCount();
});
}
blocking_counter.Wait();
}
}
INSTANTIATE_TEST_SUITE_P(BatchFunctionKernelPaddingTestSuite,
BatchFunctionKernelPaddingTest,
::testing::Values("PAD_UP", "BATCH_DOWN",
"MINIMIZE_TPU_COST_PER_REQUEST"));
} // namespace
} // namespace tensorflow
tensorflow/core/kernels/batch_norm_op.cc
+264
View File
@@ -0,0 +1,264 @@
/* Copyright 2015 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.
==============================================================================*/
// See docs in ../ops/nn_ops.cc.
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#define EIGEN_USE_THREADS
#include "tensorflow/core/kernels/batch_norm_op.h"
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/numeric_op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
namespace tensorflow {
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
template <typename Device, typename T>
class BatchNormOp : public OpKernel {
public:
explicit BatchNormOp(OpKernelConstruction* context) : OpKernel(context) {
float variance_epsilon;
OP_REQUIRES_OK(context,
context->GetAttr("variance_epsilon", &variance_epsilon));
variance_epsilon_ = T(variance_epsilon);
OP_REQUIRES_OK(context, context->GetAttr("scale_after_normalization",
&scale_after_normalization_));
}
void Compute(OpKernelContext* context) override {
const Tensor& input = context->input(0);
const Tensor& mean = context->input(1);
const Tensor& var = context->input(2);
const Tensor& beta = context->input(3);
const Tensor& gamma = context->input(4);
OP_REQUIRES(
context, input.dims() == 4,
absl::InvalidArgumentError(absl::StrCat("input must be 4-dimensional",
input.shape().DebugString())));
OP_REQUIRES(context, mean.dims() == 1,
absl::InvalidArgumentError(absl::StrCat(
"mean must be 1-dimensional", mean.shape().DebugString())));
OP_REQUIRES(context, var.dims() == 1,
absl::InvalidArgumentError(absl::StrCat(
"var must be 1-dimensional", var.shape().DebugString())));
OP_REQUIRES(context, beta.dims() == 1,
absl::InvalidArgumentError(absl::StrCat(
"beta must be 1-dimensional", beta.shape().DebugString())));
OP_REQUIRES(
context, gamma.dims() == 1,
absl::InvalidArgumentError(absl::StrCat("gamma must be 1-dimensional",
gamma.shape().DebugString())));
Tensor* output = nullptr;
OP_REQUIRES_OK(context,
context->allocate_output(0, input.shape(), &output));
functor::BatchNorm<Device, T>()(
context->eigen_device<Device>(), input.tensor<T, 4>(), mean.vec<T>(),
var.vec<T>(), beta.vec<T>(), gamma.vec<T>(), variance_epsilon_,
scale_after_normalization_, output->tensor<T, 4>());
}
private:
T variance_epsilon_;
bool scale_after_normalization_;
};
template <typename Device, typename T>
class BatchNormGradOp : public OpKernel {
public:
explicit BatchNormGradOp(OpKernelConstruction* context) : OpKernel(context) {
float variance_epsilon;
OP_REQUIRES_OK(context,
context->GetAttr("variance_epsilon", &variance_epsilon));
variance_epsilon_ = T(variance_epsilon);
OP_REQUIRES_OK(context, context->GetAttr("scale_after_normalization",
&scale_after_normalization_));
}
void Compute(OpKernelContext* context) override {
const Tensor& input = context->input(0);
const Tensor& mean = context->input(1);
const Tensor& var = context->input(2);
const Tensor& gamma = context->input(3);
const Tensor& out_backprop = context->input(4);
OP_REQUIRES(
context, input.dims() == 4,
absl::InvalidArgumentError(absl::StrCat("input must be 4-dimensional",
input.shape().DebugString())));
OP_REQUIRES(context, mean.dims() == 1,
absl::InvalidArgumentError(absl::StrCat(
"mean must be 1-dimensional", mean.shape().DebugString())));
OP_REQUIRES(context, var.dims() == 1,
absl::InvalidArgumentError(absl::StrCat(
"var must be 1-dimensional", var.shape().DebugString())));
OP_REQUIRES(
context, gamma.dims() == 1,
absl::InvalidArgumentError(absl::StrCat("gamma must be 1-dimensional",
gamma.shape().DebugString())));
OP_REQUIRES(context, out_backprop.dims() == 4,
absl::InvalidArgumentError(
absl::StrCat("out_backprop must be 4-dimensional",
out_backprop.shape().DebugString())));
Tensor* dx = nullptr;
OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
{0, 4}, 0, input.shape(), &dx));
Tensor* dm = nullptr;
OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
{1}, 1, mean.shape(), &dm));
Tensor* dv = nullptr;
OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
{2}, 2, var.shape(), &dv));
Tensor* db = nullptr;
if (scale_after_normalization_) {
OP_REQUIRES_OK(context, context->allocate_output(3, mean.shape(), &db));
} else {
OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
{3}, 3, mean.shape(), &db));
}
Tensor* dg = nullptr;
OP_REQUIRES_OK(context, context->allocate_output(4, gamma.shape(), &dg));
// Scratch buffer of [depth] dimension, aka the 4th dimension of input,
// which is dim_size(3), for calculating various combinations of
// (var + epsilon).
Tensor scratch1;
OP_REQUIRES_OK(context, context->allocate_temp(
DataTypeToEnum<T>::value,
TensorShape({input.dim_size(3)}), &scratch1));
// Scratch buffer of [depth] dimension for saving intermediate calculation
// values.
Tensor scratch2;
OP_REQUIRES_OK(context, context->allocate_temp(
DataTypeToEnum<T>::value,
TensorShape({input.dim_size(3)}), &scratch2));
functor::BatchNormGrad<Device, T>()(
context->eigen_device<Device>(), input.tensor<T, 4>(), mean.vec<T>(),
var.vec<T>(), gamma.vec<T>(), out_backprop.tensor<T, 4>(),
variance_epsilon_, scale_after_normalization_, dx->tensor<T, 4>(),
dm->vec<T>(), dv->vec<T>(), db->vec<T>(), dg->vec<T>(),
scratch1.vec<T>(), scratch2.vec<T>());
}
private:
T variance_epsilon_;
bool scale_after_normalization_;
};
#define REGISTER_KERNEL(T) \
REGISTER_KERNEL_BUILDER(Name("BatchNormWithGlobalNormalization") \
.Device(DEVICE_CPU) \
.TypeConstraint<T>("T"), \
BatchNormOp<CPUDevice, T>);
TF_CALL_half(REGISTER_KERNEL);
TF_CALL_float(REGISTER_KERNEL);
TF_CALL_double(REGISTER_KERNEL);
#undef REGISTER_KERNEL
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
// Forward declarations of the functor specializations for GPU.
namespace functor {
#define DECLARE_GPU_SPEC(T) \
template <> \
void BatchNorm<GPUDevice, T>::operator()( \
const GPUDevice& d, typename TTypes<T, 4>::ConstTensor input, \
typename TTypes<T>::ConstVec mean, typename TTypes<T>::ConstVec var, \
typename TTypes<T>::ConstVec beta, typename TTypes<T>::ConstVec gamma, \
T variance_epsilon, bool scale_after_normalization, \
typename TTypes<T, 4>::Tensor output); \
extern template struct BatchNorm<GPUDevice, T>;
#define DECLARE_GPU_SPECS(T) DECLARE_GPU_SPEC(T);
TF_CALL_half(DECLARE_GPU_SPECS);
TF_CALL_float(DECLARE_GPU_SPECS);
#undef DECLARE_GPU_SPEC
} // namespace functor
// Registration of the GPU implementations.
#define REGISTER_GPU_KERNEL(T) \
REGISTER_KERNEL_BUILDER(Name("BatchNormWithGlobalNormalization") \
.Device(DEVICE_GPU) \
.TypeConstraint<T>("T"), \
BatchNormOp<GPUDevice, T>);
TF_CALL_half(REGISTER_GPU_KERNEL);
TF_CALL_float(REGISTER_GPU_KERNEL);
#undef REGISTER_GPU_KERNEL
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define REGISTER_KERNEL(T) \
REGISTER_KERNEL_BUILDER(Name("BatchNormWithGlobalNormalizationGrad") \
.Device(DEVICE_CPU) \
.TypeConstraint<T>("T"), \
BatchNormGradOp<CPUDevice, T>);
TF_CALL_half(REGISTER_KERNEL);
TF_CALL_float(REGISTER_KERNEL);
TF_CALL_double(REGISTER_KERNEL);
#undef REGISTER_KERNEL
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
// Forward declarations of the functor specializations for GPU.
namespace functor {
#define DECLARE_GPU_SPEC(T) \
template <> \
void BatchNormGrad<GPUDevice, T>::operator()( \
const GPUDevice& d, typename TTypes<T, 4>::ConstTensor input, \
typename TTypes<T>::ConstVec mean, typename TTypes<T>::ConstVec var, \
typename TTypes<T>::ConstVec gamma, \
typename TTypes<T, 4>::ConstTensor out_backprop, T variance_epsilon, \
bool scale_after_normalization, typename TTypes<T, 4>::Tensor dx, \
typename TTypes<T>::Vec dm, typename TTypes<T>::Vec dv, \
typename TTypes<T>::Vec db, typename TTypes<T>::Vec dg, \
typename TTypes<T>::Vec scratch1, typename TTypes<T>::Vec scratch2); \
extern template struct BatchNormGrad<GPUDevice, T>;
#define DECLARE_GPU_SPECS(T) DECLARE_GPU_SPEC(T);
TF_CALL_half(DECLARE_GPU_SPECS);
TF_CALL_float(DECLARE_GPU_SPECS);
#undef DECLARE_GPU_SPEC
} // namespace functor
// Registration of the GPU implementations.
#define REGISTER_GPU_KERNEL(T) \
REGISTER_KERNEL_BUILDER(Name("BatchNormWithGlobalNormalizationGrad") \
.Device(DEVICE_GPU) \
.TypeConstraint<T>("T"), \
BatchNormGradOp<GPUDevice, T>);
TF_CALL_half(REGISTER_GPU_KERNEL);
TF_CALL_float(REGISTER_GPU_KERNEL);
#undef REGISTER_GPU_KERNEL
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow
tensorflow/core/kernels/batch_norm_op.h
+143
View File
@@ -0,0 +1,143 @@
/* Copyright 2015 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BATCH_NORM_OP_H_
#define TENSORFLOW_CORE_KERNELS_BATCH_NORM_OP_H_
// Functor definition for BatchNormOp, must be compilable by nvcc.
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/tensor_types.h"
namespace tensorflow {
namespace functor {
// Functor used by BatchNormOp to do the computations.
template <typename Device, typename T>
struct BatchNorm {
void operator()(const Device& d, typename TTypes<T, 4>::ConstTensor input,
typename TTypes<T>::ConstVec mean,
typename TTypes<T>::ConstVec var,
typename TTypes<T>::ConstVec beta,
typename TTypes<T>::ConstVec gamma, T variance_epsilon,
bool scale_after_normalization,
typename TTypes<T, 4>::Tensor output) {
const int depth = mean.dimension(0);
const int rest_size = input.size() / depth;
Eigen::DSizes<int, 2> rest_by_depth(rest_size, depth);
Eigen::IndexList<int, Eigen::type2index<1> > rest_by_one;
rest_by_one.set(0, rest_size);
Eigen::IndexList<Eigen::type2index<1>, int> one_by_depth;
one_by_depth.set(1, depth);
Eigen::IndexList<int, Eigen::type2index<1> > depth_by_one;
depth_by_one.set(0, depth);
if (scale_after_normalization) {
output.reshape(rest_by_depth).device(d) =
(input.reshape(rest_by_depth) -
mean.reshape(one_by_depth).broadcast(rest_by_one)) *
((var + var.constant(variance_epsilon)).rsqrt() * gamma)
.eval()
.reshape(one_by_depth)
.broadcast(rest_by_one) +
beta.reshape(one_by_depth).broadcast(rest_by_one);
} else {
output.reshape(rest_by_depth).device(d) =
(input.reshape(rest_by_depth) -
mean.reshape(one_by_depth).broadcast(rest_by_one)) *
((var + var.constant(variance_epsilon)).rsqrt())
.eval()
.reshape(one_by_depth)
.broadcast(rest_by_one) +
beta.reshape(one_by_depth).broadcast(rest_by_one);
}
}
};
template <typename Device, typename T>
struct BatchNormGrad {
void operator()(const Device& d, typename TTypes<T, 4>::ConstTensor input,
typename TTypes<T>::ConstVec mean,
typename TTypes<T>::ConstVec var,
typename TTypes<T>::ConstVec gamma,
typename TTypes<T, 4>::ConstTensor out_backprop,
T variance_epsilon, bool scale_after_normalization,
typename TTypes<T, 4>::Tensor dx, typename TTypes<T>::Vec dm,
typename TTypes<T>::Vec dv, typename TTypes<T>::Vec db,
typename TTypes<T>::Vec dg, typename TTypes<T>::Vec scratch1,
typename TTypes<T>::Vec scratch2) {
const int depth = mean.dimension(0);
const int rest_size = input.size() / depth;
typedef typename TTypes<T>::ConstVec::Index Index;
Eigen::DSizes<Index, 2> rest_by_depth(rest_size, depth);
Eigen::IndexList<Index, Eigen::type2index<1> > rest_by_one;
rest_by_one.set(0, rest_size);
Eigen::IndexList<Eigen::type2index<1>, Index> one_by_depth;
one_by_depth.set(1, depth);
Eigen::IndexList<Eigen::type2index<0> > reduction_axis;
// db = out_backprop
//
// dg = out_backprop * ((x - m) * rsqrt(v + epsilon))
//
// dv = sum_over_rest(out_backprop * gamma * (x - m)) *
// (-1/2) * (v + epsilon) ^ (-3/2)
//
// dm = sum_over_rest(out_backprop * gamma) * (-1 / rsqrt(v + epsilon))
//
// dx = out_backprop * (gamma * rsqrt(v + epsilon))
db.device(d) = out_backprop.reshape(rest_by_depth).sum(reduction_axis);
// scratch1 = rsqrt(v + epsilon)
scratch1.device(d) = (var + var.constant(variance_epsilon)).rsqrt();
// scratch2 = sum_over_rest(out_backprop * (x - m))
scratch2.device(d) = (out_backprop.reshape(rest_by_depth) *
(input.reshape(rest_by_depth) -
mean.reshape(one_by_depth).broadcast(rest_by_one)))
.sum(reduction_axis);
if (scale_after_normalization) {
dx.reshape(rest_by_depth).device(d) =
out_backprop.reshape(rest_by_depth) * ((scratch1 * gamma)
.eval()
.reshape(one_by_depth)
.broadcast(rest_by_one));
dm.device(d) = -db * (scratch1 * gamma).eval();
dg.device(d) = scratch2 * scratch1;
} else {
dx.reshape(rest_by_depth).device(d) =
out_backprop.reshape(rest_by_depth) *
scratch1.reshape(one_by_depth).broadcast(rest_by_one);
dm.device(d) = -db * scratch1;
dg.device(d) = dg.constant(static_cast<T>(0.0)); // Gamma is not learned.
}
// scratch1 = - 1/2 * (var + epsilon) ^ (-3/2)
scratch1.device(d) = scratch1 * scratch1.constant(static_cast<T>(-0.5f)) /
(var + var.constant(variance_epsilon));
if (scale_after_normalization) {
dv.device(d) = scratch2 * (scratch1 * gamma).eval();
} else {
dv.device(d) = scratch2 * scratch1;
}
}
};
} // namespace functor
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCH_NORM_OP_H_
tensorflow/core/kernels/batch_norm_op_gpu.cu.cc
+34
View File
@@ -0,0 +1,34 @@
/* Copyright 2015 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.
==============================================================================*/
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define EIGEN_USE_GPU
#include "tensorflow/core/kernels/batch_norm_op.h"
#include "tensorflow/core/framework/register_types.h"
namespace tensorflow {
typedef Eigen::GpuDevice GPUDevice;
template struct functor::BatchNorm<GPUDevice, float>;
template struct functor::BatchNorm<GPUDevice, Eigen::half>;
template struct functor::BatchNormGrad<GPUDevice, float>;
template struct functor::BatchNormGrad<GPUDevice, Eigen::half>;
} // namespace tensorflow
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
tensorflow/core/kernels/batch_norm_op_test.cc
+79
View File
@@ -0,0 +1,79 @@
/* Copyright 2015 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 <vector>
#include "tensorflow/core/framework/allocator.h"
#include "tensorflow/core/framework/fake_input.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
template <typename T>
struct BatchNormOpTest : public OpsTestBase {
static constexpr auto TValueType = DataTypeToEnum<T>::value;
void run_me() {
TF_EXPECT_OK(
NodeDefBuilder("batch_norm_op", "BatchNormWithGlobalNormalization")
.Input(FakeInput(TValueType))
.Input(FakeInput(TValueType))
.Input(FakeInput(TValueType))
.Input(FakeInput(TValueType))
.Input(FakeInput(TValueType))
.Attr("scale_after_normalization", false)
.Attr("variance_epsilon", 0.001)
.Finalize(node_def()));
TF_EXPECT_OK(InitOpWithGraphVersion(8));
AddInputFromList<T>(TensorShape({1, 1, 6, 2}),
{1, 4, 2, 5, 3, 6, -1, -4, -2, -5, -3, -6});
AddInputFromList<T>(TensorShape({2}), {10, 20});
AddInputFromList<T>(TensorShape({2}), {0.25, 0.5});
AddInputFromList<T>(TensorShape({2}), {0.1, 0.6});
AddInputFromList<T>(TensorShape({2}), {0.0, 0.0});
TF_ASSERT_OK(RunOpKernel());
double atol = TValueType == DT_FLOAT ? 0.01 : 0.1;
Tensor expected(allocator(), TValueType, TensorShape({1, 1, 6, 2}));
test::FillValues<T>(&expected,
{-17.86f, -22.00f, -15.87f, -20.59f, -13.87f, -19.18f,
-21.86f, -33.31f, -23.85f, -34.72f, -25.85f, -36.13f});
test::ExpectTensorNear<T>(expected, *GetOutput(0), atol);
}
};
TYPED_TEST_SUITE_P(BatchNormOpTest);
TYPED_TEST_P(BatchNormOpTest, Simple) { this->run_me(); }
REGISTER_TYPED_TEST_SUITE_P(BatchNormOpTest, Simple);
// TODO(ezhulenev): Add support for more data types.
using DataTypes = ::testing::Types<float, Eigen::half>;
INSTANTIATE_TYPED_TEST_SUITE_P(Test, BatchNormOpTest, DataTypes);
} // namespace tensorflow
tensorflow/core/kernels/batch_util.h
+23
View File
@@ -0,0 +1,23 @@
/* Copyright 2017 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.
==============================================================================*/
// NOTE(lespeholt): This file is deprecated. Use
// "tensorflow/core/util/batch_util.h" instead.
#ifndef TENSORFLOW_CORE_KERNELS_BATCH_UTIL_H_
#define TENSORFLOW_CORE_KERNELS_BATCH_UTIL_H_
#include "tensorflow/core/util/batch_util.h"
#endif // TENSORFLOW_CORE_KERNELS_BATCH_UTIL_H_
tensorflow/core/kernels/batching_util/BUILD
+574
View File
@@ -0,0 +1,574 @@
# Description: Utilities.
load("//tensorflow:tensorflow.bzl", "if_google", "tf_cc_test")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
cc_library(
name = "periodic_function_dynamic",
srcs = ["periodic_function.cc"],
hdrs = ["periodic_function.h"],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core/lib/core:notification",
"@com_google_absl//absl/functional:any_invocable",
"@com_google_absl//absl/synchronization",
],
)
cc_library(
name = "input_split_metadata",
srcs = ["input_split_metadata.cc"],
hdrs = ["input_split_metadata.h"],
deps = [
"@com_google_absl//absl/container:fixed_array",
"@com_google_absl//absl/strings",
],
)
cc_library(
name = "batch_stats",
hdrs = ["batch_stats.h"],
deps = [
"//tensorflow/core:framework_lite",
"//tensorflow/core:portable_gif_internal",
"@com_google_absl//absl/container:node_hash_map",
"@com_google_absl//absl/time",
],
)
tf_cc_test(
name = "batch_stats_test",
srcs = ["batch_stats_test.cc"],
deps = [
":batch_stats",
"//tensorflow/core:test",
"@com_google_absl//absl/time",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "batch_input_task",
hdrs = ["batch_input_task.h"],
deps = [
":batch_scheduler_hdrs",
":concat_split_util",
":input_split_metadata",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:thread_annotations",
"//tensorflow/core/util:incremental_barrier",
"@com_google_absl//absl/base",
"@com_google_absl//absl/container:fixed_array",
"@com_google_absl//absl/synchronization",
],
)
tf_cc_test(
name = "batch_input_task_test",
srcs = ["batch_input_task_test.cc"],
deps = [
":batch_input_task",
":batch_resource_base",
":input_split_metadata",
":threadsafe_status",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:ops",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/common_runtime:device",
"//tensorflow/core/framework:node_def_proto_cc",
"//tensorflow/core/framework:tensor_testutil",
"//tensorflow/core/framework:types_proto_cc",
"//tensorflow/core/kernels:batch_kernels",
"//tensorflow/core/platform:status_matchers",
"//tensorflow/core/protobuf:error_codes_proto_impl_cc",
"@com_google_absl//absl/status",
],
)
cc_library(
name = "periodic_function",
deps = [
":periodic_function_dynamic",
"//tensorflow/core:lib",
],
)
tf_cc_test(
name = "periodic_function_test",
srcs = ["periodic_function_test.cc"],
deps = [
":fake_clock_env",
":periodic_function_dynamic",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/synchronization",
],
)
cc_library(
name = "batch_scheduler_hdrs",
hdrs = ["batch_scheduler.h"],
deps = [
"//tensorflow/core:framework_lite",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core/lib/core:status",
"//tensorflow/core/platform:thread_annotations",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
"@tsl//tsl/profiler/lib:traceme",
"@xla//xla/tsl/platform:criticality",
],
)
cc_library(
name = "batch_scheduler",
srcs = ["batch_scheduler.cc"],
hdrs = ["batch_scheduler.h"],
deps = [
"//tensorflow/core:lib",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
"@tsl//tsl/profiler/lib:traceme",
"@xla//xla/tsl/platform:criticality",
],
)
cc_library(
name = "batch_scheduler_utils",
srcs = ["batch_scheduler_utils.cc"],
hdrs = ["batch_scheduler_utils.h"],
deps = [
":batch_scheduler_hdrs",
":batch_stats",
"//tensorflow/core:lib",
"//tensorflow/core:portable_gif_internal",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/time",
],
)
cc_library(
name = "threadsafe_status",
srcs = ["threadsafe_status.cc"],
hdrs = ["threadsafe_status.h"],
deps = [
"//tensorflow/core:lib",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/status",
"@com_google_absl//absl/synchronization",
],
)
tf_cc_test(
name = "batch_scheduler_test",
srcs = ["batch_scheduler_test.cc"],
deps = [
":batch_scheduler",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:status_matchers",
"@com_google_absl//absl/status",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
"@com_google_googletest//:gtest_main",
],
)
tf_cc_test(
name = "batch_scheduler_utils_test",
srcs = ["batch_scheduler_utils_test.cc"],
deps = [
":batch_scheduler_hdrs",
":batch_scheduler_utils",
":batch_stats",
"@com_google_absl//absl/time",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "shared_batch_scheduler_hdrs",
hdrs = ["shared_batch_scheduler.h"],
deps = [
":batch_scheduler_hdrs",
":batch_scheduler_utils",
":batch_stats",
":periodic_function_dynamic",
"//tensorflow/core:framework_lite",
"//tensorflow/core:lib",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core/lib/core:errors",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:notification",
"//tensorflow/core/platform:thread_annotations",
"//tensorflow/core/profiler/lib:traceme",
"//tensorflow/core/profiler/lib:traceme_encode",
"@com_google_absl//absl/container:btree",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
"@tsl//tsl/platform:errors",
"@tsl//tsl/profiler/lib:connected_traceme",
"@tsl//tsl/profiler/lib:context_types_hdrs",
"@tsl//tsl/profiler/lib:traceme",
"@xla//xla/tsl/platform:criticality",
],
)
cc_library(
name = "shared_batch_scheduler",
hdrs = ["shared_batch_scheduler.h"],
deps = [
":batch_scheduler",
":batch_scheduler_utils",
":batch_stats",
":periodic_function_dynamic",
"//tensorflow/core:lib",
"//tensorflow/core/profiler/lib:traceme",
"//tensorflow/core/profiler/lib:traceme_encode",
"@com_google_absl//absl/container:btree",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
"@tsl//tsl/profiler/lib:connected_traceme",
"@tsl//tsl/profiler/lib:context_types_hdrs",
"@tsl//tsl/profiler/lib:traceme",
"@xla//xla/tsl/platform:criticality",
],
alwayslink = 1,
)
tf_cc_test(
name = "shared_batch_scheduler_test",
size = "small",
timeout = "moderate",
srcs = ["shared_batch_scheduler_test.cc"],
tags = ["no_windows"],
deps = [
":batch_scheduler",
":batch_scheduler_utils",
":batch_stats",
":fake_clock_env",
":input_split_metadata",
":shared_batch_scheduler",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"@com_google_absl//absl/base",
"@com_google_absl//absl/container:fixed_array",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
"@com_google_benchmark//:benchmark",
"@com_google_googletest//:gtest_main",
"@xla//xla/tsl/lib/monitoring:cell_reader",
],
)
cc_library(
name = "adaptive_shared_batch_scheduler",
hdrs = ["adaptive_shared_batch_scheduler.h"],
deps = [
":batch_scheduler",
":periodic_function_dynamic",
"//tensorflow/core:lib",
"//tensorflow/core/profiler/lib:connected_traceme",
"@com_google_absl//absl/types:optional",
],
)
tf_cc_test(
name = "adaptive_shared_batch_scheduler_test",
srcs = ["adaptive_shared_batch_scheduler_test.cc"],
tags = [
"local",
"manual",
],
deps = [
":adaptive_shared_batch_scheduler",
":fake_clock_env",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/synchronization",
],
)
cc_library(
name = "serial_device_batch_scheduler",
hdrs = ["serial_device_batch_scheduler.h"],
deps = [
":batch_scheduler",
"//tensorflow/core:lib",
],
)
tf_cc_test(
name = "serial_device_batch_scheduler_test",
srcs = ["serial_device_batch_scheduler_test.cc"],
tags = [
"notap", # b/110374108
],
deps = [
":fake_clock_env",
":serial_device_batch_scheduler",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/synchronization",
],
)
cc_library(
name = "basic_batch_scheduler",
hdrs = ["basic_batch_scheduler.h"],
deps = [
":shared_batch_scheduler",
],
)
tf_cc_test(
name = "basic_batch_scheduler_test",
srcs = ["basic_batch_scheduler_test.cc"],
deps = [
":basic_batch_scheduler",
":batch_scheduler",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
],
)
tf_cc_test(
name = "basic_batch_scheduler_benchmark",
srcs = ["basic_batch_scheduler_benchmark_test.cc"],
tags = [
"local",
"manual",
],
deps = [
":basic_batch_scheduler",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:tensorflow",
"//tensorflow/core:test",
],
)
tf_cc_test(
name = "threadsafe_status_test",
srcs = ["threadsafe_status_test.cc"],
deps = [
":threadsafe_status",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
],
)
cc_library(
name = "fake_clock_env",
testonly = 1,
srcs = ["fake_clock_env.cc"],
hdrs = ["fake_clock_env.h"],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:tensorflow",
"@com_google_absl//absl/synchronization",
],
)
cc_library(
name = "bounded_executor",
srcs = ["bounded_executor.cc"],
hdrs = ["bounded_executor.h"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:framework_headers_lib",
"//tensorflow/core:lib",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:status",
"//tensorflow/core/platform:statusor",
"//tensorflow/core/platform:thread_annotations",
"//tensorflow/core/platform:threadpool_interface",
"@com_google_absl//absl/functional:bind_front",
],
)
tf_cc_test(
name = "bounded_executor_test",
size = "small",
srcs = ["bounded_executor_test.cc"],
deps = [
":bounded_executor",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:status_matchers",
"//tensorflow/core/platform:statusor",
"//tensorflow/core/platform:threadpool_interface",
"@com_google_absl//absl/functional:bind_front",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
],
)
cc_library(
name = "concat_split_util",
hdrs = ["concat_split_util.h"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core/kernels:concat_lib",
"//tensorflow/core/kernels:split_lib",
"//tensorflow/core/platform:status",
],
)
cc_library(
name = "batch_resource_base",
srcs = ["batch_resource_base.cc"],
hdrs = ["batch_resource_base.h"],
deps = [
":adaptive_shared_batch_scheduler",
":batch_scheduler",
":batch_scheduler_utils",
":batch_stats",
":concat_split_util",
":input_split_metadata",
":shared_batch_scheduler",
":threadsafe_status",
":warmup",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core/common_runtime:cost_constants",
"//tensorflow/core/common_runtime:cost_measurement",
"//tensorflow/core/common_runtime:cost_measurement_registry",
"//tensorflow/core/common_runtime:cost_util",
"//tensorflow/core/common_runtime:request_cost",
"//tensorflow/core/common_runtime:request_cost_accessor",
"//tensorflow/core/common_runtime:request_cost_accessor_registry",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:status",
"//tensorflow/core/platform:thread_annotations",
"//tensorflow/core/profiler/lib:traceme",
"//tensorflow/core/profiler/lib:traceme_encode",
"//tensorflow/core/protobuf:for_core_protos_cc",
"//tensorflow/core/util:incremental_barrier",
"@com_google_absl//absl/container:fixed_array",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/functional:bind_front",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
"@com_google_absl//absl/types:optional",
"@xla//xla/tsl/platform:criticality",
],
)
cc_library(
name = "warmup",
srcs = ["warmup.cc"],
hdrs = ["warmup.h"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/hash",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@tsl//tsl/platform:logging",
],
)
tf_cc_test(
name = "batch_resource_base_test",
srcs = ["batch_resource_base_test.cc"],
deps = [
":batch_resource_base",
":batch_scheduler_hdrs",
":batch_scheduler_utils",
":batch_stats",
":shared_batch_scheduler",
":threadsafe_status",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:ops",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core/common_runtime:cost_constants",
"//tensorflow/core/common_runtime:cost_measurement",
"//tensorflow/core/common_runtime:cost_measurement_registry",
"//tensorflow/core/common_runtime:no_op_cost_measurement",
"//tensorflow/core/common_runtime:request_cost",
"//tensorflow/core/framework:types_proto_cc",
"//tensorflow/core/kernels:batch_kernels",
"//tensorflow/core/lib/monitoring:cell_reader",
"//tensorflow/core/platform:notification",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
"@com_google_absl//absl/types:span",
"@com_google_googletest//:gtest_main",
"@tsl//tsl/platform:refcount",
"@tsl//tsl/platform:status",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/lib/monitoring:cell_reader",
"@xla//xla/tsl/lib/monitoring:test_utils",
"@xla//xla/tsl/platform:criticality",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "input_split_metadata_test",
srcs = ["input_split_metadata_test.cc"],
deps = [
":input_split_metadata",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/strings",
],
)
tensorflow/core/kernels/batching_util/adaptive_shared_batch_scheduler.h
+871
View File
@@ -0,0 +1,871 @@
/* Copyright 2017 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_ADAPTIVE_SHARED_BATCH_SCHEDULER_H_
#define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_ADAPTIVE_SHARED_BATCH_SCHEDULER_H_
#include <algorithm>
#include <atomic>
#include <functional>
#include <memory>
#include <random>
#include <unordered_map>
#include <vector>
#include "absl/types/optional.h"
#include "tensorflow/core/kernels/batching_util/batch_scheduler.h"
#include "tensorflow/core/kernels/batching_util/periodic_function.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/platform/byte_order.h"
#include "tensorflow/core/platform/cpu_info.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/platform/threadpool_interface.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/profiler/lib/connected_traceme.h"
namespace tensorflow {
namespace serving {
namespace internal {
template <typename TaskType>
class ASBSBatch;
template <typename TaskType>
class ASBSQueue;
} // namespace internal
// Shared batch scheduler designed to minimize latency. The scheduler keeps
// track of a number of queues (one per model or model version) which are
// continuously enqueuing requests. The scheduler groups the requests into
// batches which it periodically sends off for processing (see
// shared_batch_scheduler.h for more details). AdaptiveSharedBatchScheduler
// (ASBS) prioritizes batches primarily by age (i.e. the batch's oldest request)
// along with a configurable preference for scheduling larger batches first.
//
//
// ASBS tries to keep the system busy by maintaining an adjustable number of
// concurrently processed batches. If a new batch is created, and the number of
// in flight batches is below the target, the next (i.e. oldest) batch is
// immediately scheduled. Similarly, when a batch finishes processing, the
// target is rechecked, and another batch may be scheduled. To avoid the need
// to carefully tune the target for workload, model type, platform, etc, it is
// dynamically adjusted in order to provide the lowest average latency.
//
// Some potential use cases:
// Hardware Accelerators (GPUs & TPUs) - If some phase of batch processing
// involves serial processing by a device, from a latency perspective it is
// desirable to keep the device evenly loaded, avoiding the need to wait for
// the device to process prior batches.
// CPU utilization - If the batch processing is cpu dominated, you can reap
// latency gains when underutilized by increasing the processing rate, but
// back the rate off when the load increases to avoid overload.
template <typename TaskType>
class AdaptiveSharedBatchScheduler
: public std::enable_shared_from_this<
AdaptiveSharedBatchScheduler<TaskType>> {
public:
~AdaptiveSharedBatchScheduler() {
// Finish processing batches before destroying other class members.
if (owned_batch_thread_pool_) {
delete batch_thread_pool_;
}
}
struct Options {
// The name to use for the pool of batch threads.
std::string thread_pool_name = {"batch_threads"};
// Number of batch processing threads - the maximum value of
// in_flight_batches_limit_. It is recommended that this value be set by
// running the system under load, observing the learned value for
// in_flight_batches_limit_, and setting this maximum to ~ 2x the value.
// Under low load, in_flight_batches_limit_ has no substantial effect on
// latency and therefore undergoes a random walk. Unreasonably large values
// for num_batch_threads allows for large in_flight_batches_limit_, which
// will harm latency for some time once load increases again.
int64_t num_batch_threads = port::MaxParallelism();
// You can pass a ThreadPool directly rather than the above two
// parameters. If given, the above two parameers are ignored. Ownership of
// the threadpool is not transferred.
thread::ThreadPool* thread_pool = nullptr;
// Lower bound for in_flight_batches_limit_. As discussed above, can be used
// to minimize the damage caused by the random walk under low load.
int64_t min_in_flight_batches_limit = 1;
// Although batch selection is primarily based on age, this parameter
// specifies a preference for larger batches. A full batch will be
// scheduled before an older, nearly empty batch as long as the age gap is
// less than full_batch_scheduling_boost_micros. The optimal value for this
// parameter should be of order the batch processing latency, but must be
// chosen carefully, as too large a value will harm tail latency.
int64_t full_batch_scheduling_boost_micros = 0;
// The environment to use (typically only overridden by test code).
Env* env = Env::Default();
// Initial limit for number of batches being concurrently processed.
// Non-integer values correspond to probabilistic limits - i.e. a value of
// 3.2 results in an actual cap of 3 80% of the time, and 4 20% of the time.
double initial_in_flight_batches_limit = 3;
// Number of batches between adjustments of in_flight_batches_limit. Larger
// numbers will give less noisy latency measurements, but will be less
// responsive to changes in workload.
int64_t batches_to_average_over = 1000;
// If true, schedule batches using FIFO policy.
// Requires that `full_batch_scheduling_boost_micros` is zero.
// NOTE:
// A new parameter is introduced (not re-using
// full_batch_scheduling_boost_micros==zero) for backward compatibility of
// API.
bool fifo_scheduling = false;
};
// Ownership is shared between the caller of Create() and any queues created
// via AddQueue().
static absl::Status Create(
const Options& options,
std::shared_ptr<AdaptiveSharedBatchScheduler<TaskType>>* scheduler);
struct QueueOptions {
// Maximum size of a batch that's formed within
// `ASBSQueue<TaskType>::Schedule`.
int max_batch_size = 1000;
// Maximum size of input task, which is submitted to the queue by
// calling `ASBSQueue<TaskType>::Schedule` and used to form batches.
//
// If specified, it should be larger than or equal to 'max_batch_size'.
absl::optional<int> max_input_task_size = std::nullopt;
// Maximum number of tasks to add to a specific batch.
absl::optional<int> max_tasks_per_batch = std::nullopt;
// Maximum number of enqueued (i.e. non-scheduled) batches.
int max_enqueued_batches = 10;
// Amount of time non-full batches must wait before becoming schedulable.
// A non-zero value can improve performance by limiting the scheduling of
// nearly empty batches.
int64_t batch_timeout_micros = 0;
// If non nullptr, split_input_task_func should split input_task into
// multiple tasks, the first of which has size first_size and the remaining
// not exceeding max_size. This function may acquire ownership of input_task
// and should return a status indicating if the split was successful. Upon
// success, the caller can assume that all output_tasks will be scheduled.
// Including this option allows the scheduler to pack batches better and
// should usually improve overall throughput.
std::function<absl::Status(
std::unique_ptr<TaskType>* input_task, int first_size,
int max_batch_size,
std::vector<std::unique_ptr<TaskType>>* output_tasks)>
split_input_task_func;
// If true, the padding will not be appended.
bool disable_padding = false;
};
using BatchProcessor = std::function<void(std::unique_ptr<Batch<TaskType>>)>;
// Adds queue (and its callback) to be managed by this scheduler.
absl::Status AddQueue(const QueueOptions& options,
BatchProcessor process_batch_callback,
std::unique_ptr<BatchScheduler<TaskType>>* queue);
double in_flight_batches_limit() {
mutex_lock l(mu_);
return in_flight_batches_limit_;
}
private:
// access to AddBatch, MaybeScheduleClosedBatches, RemoveQueue, GetEnv.
friend class internal::ASBSQueue<TaskType>;
explicit AdaptiveSharedBatchScheduler(const Options& options);
// Tracks processing latency and adjusts in_flight_batches_limit to minimize.
void CallbackWrapper(const internal::ASBSBatch<TaskType>* batch,
BatchProcessor callback, bool is_express);
// Schedules batch if in_flight_batches_limit_ is not met.
void MaybeScheduleNextBatch() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Schedules batch using FIFO policy if in_flight_batches_limit_ is not met.
void MaybeScheduleNextBatchFIFO() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Schedules all closed batches in batches_ for which an idle thread is
// available in batch_thread_pool_.
// Batches scheduled this way are called express batches.
// Express batches are not limited by in_flight_batches_limit_, and
// their latencies will not affect in_flight_batches_limit_.
void MaybeScheduleClosedBatches();
void MaybeScheduleClosedBatchesLocked() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
void MaybeScheduleClosedBatchesLockedFIFO() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
void MaybeAdjustInflightLimit() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Notifies scheduler of non-empty batch which is eligible for processing.
void AddBatch(const internal::ASBSBatch<TaskType>* batch);
// Removes queue from scheduler.
void RemoveQueue(const internal::ASBSQueue<TaskType>* queue);
Env* GetEnv() const { return options_.env; }
const Options options_;
// Collection of batches added by AddBatch, ordered by age. Owned by scheduler
// until they are released for processing.
std::vector<const internal::ASBSBatch<TaskType>*> batches_ TF_GUARDED_BY(mu_);
// Collection of batches added by AddBatch, ordered by age. Owned by
// scheduler until they are released for processing.
std::deque<const internal::ASBSBatch<TaskType>*> fifo_batches_
TF_GUARDED_BY(mu_);
// Unowned queues and callbacks added by AddQueue.
std::unordered_map<const internal::ASBSQueue<TaskType>*, BatchProcessor>
queues_and_callbacks_ TF_GUARDED_BY(mu_);
mutex mu_;
// Responsible for running the batch processing callbacks.
thread::ThreadPool* batch_thread_pool_;
bool owned_batch_thread_pool_ = false;
// Limit on number of batches which can be concurrently processed.
// Non-integer values correspond to probabilistic limits - i.e. a value of 3.2
// results in an actual cap of 3 80% of the time, and 4 20% of the time.
double in_flight_batches_limit_ TF_GUARDED_BY(mu_);
// Number of regular batches currently being processed.
int64_t in_flight_batches_ TF_GUARDED_BY(mu_) = 0;
// Number of express batches currently being processed.
int64_t in_flight_express_batches_ TF_GUARDED_BY(mu_) = 0;
// RNG engine and distribution.
std::default_random_engine rand_engine_;
std::uniform_real_distribution<double> rand_double_;
// Fields controlling the dynamic adjustment of in_flight_batches_limit_.
// Number of batches since the last in_flight_batches_limit_ adjustment.
int64_t batch_count_ TF_GUARDED_BY(mu_) = 0;
struct DelayStats {
// Sum of processing latency for batches counted by batch_count_.
int64_t batch_latency_sum = 0;
// Average batch latency for previous value of in_flight_batches_limit_.
double last_avg_latency_ms = 0;
// Did last_avg_latency_ms decrease from the previous last_avg_latency_ms?
bool last_latency_decreased = false;
// Current direction (+-) to adjust in_flight_batches_limit_
int step_direction = 1;
};
// Delay stats between the creation of a batch and the completion of a
// batch.
DelayStats batch_delay_stats_ TF_GUARDED_BY(mu_);
// Max adjustment size (as a fraction of in_flight_batches_limit_).
constexpr static double kMaxStepSizeMultiplier = 0.125; // 1/8;
// Min adjustment size (as a fraction of in_flight_batches_limit_).
constexpr static double kMinStepSizeMultiplier = 0.0078125; // 1/128
// Current adjustment size (as a fraction of in_flight_batches_limit_).
double step_size_multiplier_ TF_GUARDED_BY(mu_) = kMaxStepSizeMultiplier;
AdaptiveSharedBatchScheduler(const AdaptiveSharedBatchScheduler&) = delete;
void operator=(const AdaptiveSharedBatchScheduler&) = delete;
};
//////////////////////////////////////////////////////////
// Implementation details follow. API users need not read.
namespace internal {
// Consolidates tasks into batches, passing them off to the
// AdaptiveSharedBatchScheduler for processing.
template <typename TaskType>
class ASBSQueue : public BatchScheduler<TaskType> {
public:
using QueueOptions =
typename AdaptiveSharedBatchScheduler<TaskType>::QueueOptions;
ASBSQueue(std::shared_ptr<AdaptiveSharedBatchScheduler<TaskType>> scheduler,
const QueueOptions& options);
~ASBSQueue() override;
// Adds task to current batch. Fails if the task size is larger than the batch
// size or if the current batch is full and this queue's number of outstanding
// batches is at its maximum.
absl::Status Schedule(std::unique_ptr<TaskType>* task) override;
// Number of tasks waiting to be scheduled.
size_t NumEnqueuedTasks() const override;
// Number of size 1 tasks which could currently be scheduled without failing.
size_t SchedulingCapacity() const override;
// Notifies queue that a batch is about to be scheduled; the queue should not
// place any more tasks in this batch.
void ReleaseBatch(const ASBSBatch<TaskType>* batch);
size_t max_task_size() const override { return options_.max_batch_size; }
private:
// Number of size 1 tasks which could currently be scheduled without failing.
size_t SchedulingCapacityLocked() const TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Returns uint64 one greater than was returned by the previous call.
// Context id is reused after std::numeric_limits<uint64>::max is exhausted.
static uint64_t NewTraceMeContextIdForBatch();
std::shared_ptr<AdaptiveSharedBatchScheduler<TaskType>> scheduler_;
const QueueOptions options_;
// Owned by scheduler_.
ASBSBatch<TaskType>* current_batch_ TF_GUARDED_BY(mu_) = nullptr;
int64_t num_enqueued_batches_ TF_GUARDED_BY(mu_) = 0;
int64_t num_enqueued_tasks_ TF_GUARDED_BY(mu_) = 0;
mutable mutex mu_;
ASBSQueue(const ASBSQueue&) = delete;
void operator=(const ASBSQueue&) = delete;
};
// Batch which remembers when and by whom it was created.
template <typename TaskType>
class ASBSBatch : public Batch<TaskType> {
public:
ASBSBatch(ASBSQueue<TaskType>* queue, int64_t creation_time_micros,
int64_t batch_timeout_micros, uint64_t traceme_context_id)
: queue_(queue),
creation_time_micros_(creation_time_micros),
schedulable_time_micros_(creation_time_micros + batch_timeout_micros),
traceme_context_id_(traceme_context_id) {}
~ASBSBatch() override {}
ASBSQueue<TaskType>* queue() const { return queue_; }
int64_t creation_time_micros() const { return creation_time_micros_; }
int64_t schedulable_time_micros() const { return schedulable_time_micros_; }
uint64_t traceme_context_id() const { return traceme_context_id_; }
private:
ASBSQueue<TaskType>* queue_;
const int64_t creation_time_micros_;
const int64_t schedulable_time_micros_;
const uint64_t traceme_context_id_;
ASBSBatch(const ASBSBatch&) = delete;
void operator=(const ASBSBatch&) = delete;
};
} // namespace internal
// ---------------- AdaptiveSharedBatchScheduler ----------------
template <typename TaskType>
constexpr double AdaptiveSharedBatchScheduler<TaskType>::kMaxStepSizeMultiplier;
template <typename TaskType>
constexpr double AdaptiveSharedBatchScheduler<TaskType>::kMinStepSizeMultiplier;
template <typename TaskType>
absl::Status AdaptiveSharedBatchScheduler<TaskType>::Create(
const Options& options,
std::shared_ptr<AdaptiveSharedBatchScheduler<TaskType>>* scheduler) {
if (options.num_batch_threads < 1) {
return errors::InvalidArgument("num_batch_threads must be positive; was ",
options.num_batch_threads);
}
if (options.min_in_flight_batches_limit < 1) {
return errors::InvalidArgument(
"min_in_flight_batches_limit must be >= 1; was ",
options.min_in_flight_batches_limit);
}
if (options.min_in_flight_batches_limit > options.num_batch_threads) {
return errors::InvalidArgument(
"min_in_flight_batches_limit (", options.min_in_flight_batches_limit,
") must be <= num_batch_threads (", options.num_batch_threads, ")");
}
if (options.full_batch_scheduling_boost_micros < 0) {
return errors::InvalidArgument(
"full_batch_scheduling_boost_micros can't be negative; was ",
options.full_batch_scheduling_boost_micros);
}
if (options.initial_in_flight_batches_limit > options.num_batch_threads) {
return errors::InvalidArgument(
"initial_in_flight_batches_limit (",
options.initial_in_flight_batches_limit,
") should not be larger than num_batch_threads (",
options.num_batch_threads, ")");
}
if (options.initial_in_flight_batches_limit <
options.min_in_flight_batches_limit) {
return errors::InvalidArgument("initial_in_flight_batches_limit (",
options.initial_in_flight_batches_limit,
"must be >= min_in_flight_batches_limit (",
options.min_in_flight_batches_limit, ")");
}
if (options.batches_to_average_over < 1) {
return errors::InvalidArgument(
"batches_to_average_over should be "
"greater than or equal to 1; was ",
options.batches_to_average_over);
}
scheduler->reset(new AdaptiveSharedBatchScheduler<TaskType>(options));
return absl::OkStatus();
}
template <typename TaskType>
AdaptiveSharedBatchScheduler<TaskType>::AdaptiveSharedBatchScheduler(
const Options& options)
: options_(options),
in_flight_batches_limit_(options.initial_in_flight_batches_limit),
rand_double_(0.0, 1.0) {
std::random_device device;
rand_engine_.seed(device());
if (options.thread_pool == nullptr) {
owned_batch_thread_pool_ = true;
batch_thread_pool_ = new thread::ThreadPool(
GetEnv(), options.thread_pool_name, options.num_batch_threads);
} else {
owned_batch_thread_pool_ = false;
batch_thread_pool_ = options.thread_pool;
}
}
template <typename TaskType>
absl::Status AdaptiveSharedBatchScheduler<TaskType>::AddQueue(
const QueueOptions& options, BatchProcessor process_batch_callback,
std::unique_ptr<BatchScheduler<TaskType>>* queue) {
if (options.max_batch_size <= 0) {
return errors::InvalidArgument("max_batch_size must be positive; was ",
options.max_batch_size);
}
if (options.max_enqueued_batches <= 0) {
return errors::InvalidArgument(
"max_enqueued_batches must be positive; was ",
options.max_enqueued_batches);
}
if (options.max_input_task_size.has_value()) {
if (options.max_input_task_size.value() < options.max_batch_size) {
return errors::InvalidArgument(
"max_input_task_size must be larger than or equal to max_batch_size;"
"got max_input_task_size as ",
options.max_input_task_size.value(), " and max_batch_size as ",
options.max_batch_size);
}
}
internal::ASBSQueue<TaskType>* asbs_queue_raw;
queue->reset(asbs_queue_raw = new internal::ASBSQueue<TaskType>(
this->shared_from_this(), options));
mutex_lock l(mu_);
queues_and_callbacks_[asbs_queue_raw] = process_batch_callback;
return absl::OkStatus();
}
template <typename TaskType>
void AdaptiveSharedBatchScheduler<TaskType>::AddBatch(
const internal::ASBSBatch<TaskType>* batch) {
mutex_lock l(mu_);
if (options_.fifo_scheduling) {
fifo_batches_.push_back(batch);
} else {
batches_.push_back(batch);
}
int64_t delay_micros =
batch->schedulable_time_micros() - GetEnv()->NowMicros();
if (delay_micros <= 0) {
MaybeScheduleNextBatch();
return;
}
// Try to schedule batch once it becomes schedulable. Although scheduler waits
// for all batches to finish processing before allowing itself to be deleted,
// MaybeScheduleNextBatch() is called in other places, and therefore it's
// possible the scheduler could be deleted by the time this closure runs.
// Grab a shared_ptr reference to prevent this from happening.
GetEnv()->SchedClosureAfter(
delay_micros, [this, lifetime_preserver = this->shared_from_this()] {
mutex_lock l(mu_);
MaybeScheduleNextBatch();
});
}
template <typename TaskType>
void AdaptiveSharedBatchScheduler<TaskType>::RemoveQueue(
const internal::ASBSQueue<TaskType>* queue) {
mutex_lock l(mu_);
queues_and_callbacks_.erase(queue);
}
template <typename TaskType>
void AdaptiveSharedBatchScheduler<TaskType>::MaybeScheduleNextBatchFIFO() {
const internal::ASBSBatch<TaskType>* batch = *fifo_batches_.begin();
if (batch->schedulable_time_micros() > GetEnv()->NowMicros()) {
return;
}
fifo_batches_.pop_front();
// Queue may destroy itself after ReleaseBatch is called.
batch->queue()->ReleaseBatch(batch);
batch_thread_pool_->Schedule(std::bind(
&AdaptiveSharedBatchScheduler<TaskType>::CallbackWrapper, this, batch,
queues_and_callbacks_[batch->queue()], false /* is express */));
in_flight_batches_++;
}
template <typename TaskType>
void AdaptiveSharedBatchScheduler<
TaskType>::MaybeScheduleClosedBatchesLockedFIFO() {
// Only schedule closed batches if we have spare capacity.
int available_threads =
static_cast<int>(options_.num_batch_threads - in_flight_batches_ -
in_flight_express_batches_);
for (auto it = fifo_batches_.begin();
it != fifo_batches_.end() && available_threads > 0;
it = fifo_batches_.begin()) {
if ((*it)->IsClosed()) {
const internal::ASBSBatch<TaskType>* batch = *it;
fifo_batches_.pop_front();
batch->queue()->ReleaseBatch(batch);
batch_thread_pool_->Schedule(
std::bind(&AdaptiveSharedBatchScheduler<TaskType>::CallbackWrapper,
this, batch, queues_and_callbacks_[batch->queue()], true));
in_flight_express_batches_++;
available_threads--;
} else {
// Batches are FIFO, so stop iteration after finding the first non-closed
// batches.
break;
}
}
}
template <typename TaskType>
void AdaptiveSharedBatchScheduler<TaskType>::MaybeScheduleNextBatch() {
bool batch_empty =
options_.fifo_scheduling ? fifo_batches_.empty() : batches_.empty();
if (batch_empty || in_flight_batches_ >= in_flight_batches_limit_) return;
// Non-integer limit handled probabilistically.
if (in_flight_batches_limit_ - in_flight_batches_ < 1 &&
rand_double_(rand_engine_) >
in_flight_batches_limit_ - in_flight_batches_) {
return;
}
if (options_.fifo_scheduling) {
MaybeScheduleNextBatchFIFO();
return;
}
auto best_it = batches_.end();
double best_score = (std::numeric_limits<double>::max)();
int64_t now_micros = GetEnv()->NowMicros();
for (auto it = batches_.begin(); it != batches_.end(); it++) {
if ((*it)->schedulable_time_micros() > now_micros) continue;
const double score =
(*it)->creation_time_micros() -
options_.full_batch_scheduling_boost_micros * (*it)->size() /
static_cast<double>((*it)->queue()->max_task_size());
if (best_it == batches_.end() || score < best_score) {
best_score = score;
best_it = it;
}
}
// No schedulable batches.
if (best_it == batches_.end()) return;
const internal::ASBSBatch<TaskType>* batch = *best_it;
batches_.erase(best_it);
// Queue may destroy itself after ReleaseBatch is called.
batch->queue()->ReleaseBatch(batch);
batch_thread_pool_->Schedule(
std::bind(&AdaptiveSharedBatchScheduler<TaskType>::CallbackWrapper, this,
batch, queues_and_callbacks_[batch->queue()], false));
in_flight_batches_++;
}
template <typename TaskType>
void AdaptiveSharedBatchScheduler<TaskType>::MaybeScheduleClosedBatches() {
mutex_lock l(mu_);
MaybeScheduleClosedBatchesLocked();
}
template <typename TaskType>
void AdaptiveSharedBatchScheduler<
TaskType>::MaybeScheduleClosedBatchesLocked() {
if (options_.fifo_scheduling) {
MaybeScheduleClosedBatchesLockedFIFO();
return;
}
// Only schedule closed batches if we have spare capacity.
int available_threads =
static_cast<int>(options_.num_batch_threads - in_flight_batches_ -
in_flight_express_batches_);
for (auto it = batches_.begin();
it != batches_.end() && available_threads > 0;) {
if ((*it)->IsClosed()) {
const internal::ASBSBatch<TaskType>* batch = *it;
it = batches_.erase(it);
batch->queue()->ReleaseBatch(batch);
batch_thread_pool_->Schedule(
std::bind(&AdaptiveSharedBatchScheduler<TaskType>::CallbackWrapper,
this, batch, queues_and_callbacks_[batch->queue()], true));
in_flight_express_batches_++;
available_threads--;
} else {
++it;
}
}
}
template <typename TaskType>
void AdaptiveSharedBatchScheduler<TaskType>::CallbackWrapper(
const internal::ASBSBatch<TaskType>* batch,
AdaptiveSharedBatchScheduler<TaskType>::BatchProcessor callback,
bool is_express) {
tsl::profiler::TraceMeConsumer trace_me(
[&] {
return profiler::TraceMeEncode(
"ProcessBatch", {{"batch_size_before_padding", batch->size()},
{"_r", 2} /*root_event*/});
},
tsl::profiler::ContextType::kAdaptiveSharedBatchScheduler,
batch->traceme_context_id());
const int64_t start_time = batch->creation_time_micros();
callback(std::unique_ptr<Batch<TaskType>>(
const_cast<internal::ASBSBatch<TaskType>*>(batch)));
int64_t end_time = GetEnv()->NowMicros();
mutex_lock l(mu_);
if (is_express) {
in_flight_express_batches_--;
MaybeScheduleClosedBatchesLocked();
return;
}
in_flight_batches_--;
batch_count_++;
batch_delay_stats_.batch_latency_sum += end_time - start_time;
MaybeAdjustInflightLimit();
MaybeScheduleNextBatch();
}
template <typename TaskType>
void AdaptiveSharedBatchScheduler<TaskType>::MaybeAdjustInflightLimit() {
// Occasionally adjust in_flight_batches_limit_ to minimize average latency.
// Although the optimal value may depend on the workload, the latency should
// be a simple convex function of in_flight_batches_limit_, allowing us to
// locate the global minimum relatively quickly.
if (batch_count_ == options_.batches_to_average_over) {
double current_avg_latency_ms =
(batch_delay_stats_.batch_latency_sum / 1000.) / batch_count_;
bool current_latency_decreased =
current_avg_latency_ms < batch_delay_stats_.last_avg_latency_ms;
if (current_latency_decreased) {
// If latency improvement was because we're moving in the correct
// direction, increase step_size so that we can get to the minimum faster.
// If latency improvement was due to backtracking from a previous failure,
// decrease step_size in order to refine our location.
step_size_multiplier_ *=
(batch_delay_stats_.last_latency_decreased ? 2 : 0.5);
step_size_multiplier_ =
std::min(step_size_multiplier_, kMaxStepSizeMultiplier);
step_size_multiplier_ =
std::max(step_size_multiplier_, kMinStepSizeMultiplier);
} else {
// Return (nearly) to previous position and confirm that latency is better
// there before decreasing step size.
batch_delay_stats_.step_direction = -batch_delay_stats_.step_direction;
}
in_flight_batches_limit_ += batch_delay_stats_.step_direction *
in_flight_batches_limit_ *
step_size_multiplier_;
in_flight_batches_limit_ =
std::min(in_flight_batches_limit_,
static_cast<double>(options_.num_batch_threads));
in_flight_batches_limit_ =
std::max(in_flight_batches_limit_,
static_cast<double>(options_.min_in_flight_batches_limit));
batch_delay_stats_.last_avg_latency_ms = current_avg_latency_ms;
batch_delay_stats_.last_latency_decreased = current_latency_decreased;
batch_count_ = 0;
batch_delay_stats_.batch_latency_sum = 0;
}
}
// ---------------- ASBSQueue ----------------
namespace internal {
template <typename TaskType>
ASBSQueue<TaskType>::ASBSQueue(
std::shared_ptr<AdaptiveSharedBatchScheduler<TaskType>> scheduler,
const QueueOptions& options)
: scheduler_(scheduler), options_(options) {}
template <typename TaskType>
ASBSQueue<TaskType>::~ASBSQueue() {
// Wait until last batch has been scheduled.
const int kSleepMicros = 1000;
for (;;) {
{
mutex_lock l(mu_);
if (num_enqueued_batches_ == 0) {
break;
}
}
scheduler_->GetEnv()->SleepForMicroseconds(kSleepMicros);
}
scheduler_->RemoveQueue(this);
}
template <typename TaskType>
absl::Status ASBSQueue<TaskType>::Schedule(std::unique_ptr<TaskType>* task) {
size_t size = (*task)->size();
if (options_.split_input_task_func == nullptr &&
size > options_.max_batch_size) {
return errors::InvalidArgument("Task size ", size,
" is larger than maximum batch size ",
options_.max_batch_size);
}
if (options_.max_input_task_size.has_value() &&
(size > options_.max_input_task_size.value())) {
return errors::InvalidArgument("Task size ", size,
" is larger than max input task size ",
options_.max_input_task_size.value());
}
std::vector<std::unique_ptr<TaskType>> tasks_to_schedule;
std::vector<ASBSBatch<TaskType>*> new_batches;
bool closed_batch = false;
{
mutex_lock l(mu_);
if (size > SchedulingCapacityLocked()) {
return absl::UnavailableError("The batch scheduling queue is full");
}
int remaining_batch_size =
current_batch_ == nullptr
? options_.max_batch_size
: options_.max_batch_size - current_batch_->size();
if (options_.split_input_task_func == nullptr ||
size <= remaining_batch_size) {
// Either we don't allow task splitting or task fits within the current
// batch.
tasks_to_schedule.push_back(std::move(*task));
} else {
// Split task in order to completely fill the current batch.
// Beyond this point Schedule should not fail, as the caller has been
// promised that all of the split tasks will be scheduled.
TF_RETURN_IF_ERROR(options_.split_input_task_func(
task, remaining_batch_size, options_.max_batch_size,
&tasks_to_schedule));
}
for (auto& task : tasks_to_schedule) {
// Can't fit within current batch, close it off and try to create another.
if (current_batch_ &&
current_batch_->size() + task->size() > options_.max_batch_size) {
current_batch_->Close();
closed_batch = true;
current_batch_ = nullptr;
}
if (!current_batch_) {
num_enqueued_batches_++;
// batch.traceme_context_id connects TraceMeProducer and
// TraceMeConsumer.
// When multiple calls to "ASBS::Schedule" accumulate to one batch, they
// are processed in the same batch and should share traceme_context_id.
current_batch_ = new ASBSBatch<TaskType>(
this, scheduler_->GetEnv()->NowMicros(),
options_.batch_timeout_micros, NewTraceMeContextIdForBatch());
new_batches.push_back(current_batch_);
}
// Annotate each task (corresponds to one call of schedule) with a
// TraceMeProducer.
tsl::profiler::TraceMeProducer trace_me(
[task_size = task->size()] {
return profiler::TraceMeEncode(
"ASBSQueue::Schedule",
{{"batching_input_task_size", task_size}});
},
tsl::profiler::ContextType::kAdaptiveSharedBatchScheduler,
this->current_batch_->traceme_context_id());
current_batch_->AddTask(std::move(task));
num_enqueued_tasks_++;
// If current_batch_ is now full, allow it to be processed immediately.
bool reached_max_tasks =
(options_.max_tasks_per_batch.has_value() &&
current_batch_->num_tasks() >= options_.max_tasks_per_batch.value());
if (current_batch_->size() == options_.max_batch_size ||
reached_max_tasks) {
current_batch_->Close();
closed_batch = true;
current_batch_ = nullptr;
}
}
}
// Scheduler functions must be called outside of lock, since they may call
// ReleaseBatch.
for (auto* batch : new_batches) {
scheduler_->AddBatch(batch);
}
if (closed_batch) {
scheduler_->MaybeScheduleClosedBatches();
}
return absl::OkStatus();
}
template <typename TaskType>
void ASBSQueue<TaskType>::ReleaseBatch(const ASBSBatch<TaskType>* batch) {
mutex_lock l(mu_);
num_enqueued_batches_--;
num_enqueued_tasks_ -= batch->num_tasks();
if (batch == current_batch_) {
current_batch_->Close();
current_batch_ = nullptr;
}
}
template <typename TaskType>
size_t ASBSQueue<TaskType>::NumEnqueuedTasks() const {
mutex_lock l(mu_);
return num_enqueued_tasks_;
}
template <typename TaskType>
size_t ASBSQueue<TaskType>::SchedulingCapacity() const {
mutex_lock l(mu_);
return SchedulingCapacityLocked();
}
template <typename TaskType>
size_t ASBSQueue<TaskType>::SchedulingCapacityLocked() const {
const int current_batch_capacity =
current_batch_ ? options_.max_batch_size - current_batch_->size() : 0;
const int spare_batches =
options_.max_enqueued_batches - num_enqueued_batches_;
return spare_batches * options_.max_batch_size + current_batch_capacity;
}
template <typename TaskType>
// static
uint64_t ASBSQueue<TaskType>::NewTraceMeContextIdForBatch() {
static std::atomic<uint64_t> traceme_context_id(0);
return traceme_context_id.fetch_add(1, std::memory_order_relaxed);
}
} // namespace internal
} // namespace serving
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_ADAPTIVE_SHARED_BATCH_SCHEDULER_H_
tensorflow/core/kernels/batching_util/adaptive_shared_batch_scheduler_test.cc
+524
View File
@@ -0,0 +1,524 @@
/* Copyright 2017 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/batching_util/adaptive_shared_batch_scheduler.h"
#include "absl/synchronization/notification.h"
#include "tensorflow/core/kernels/batching_util/fake_clock_env.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/gtl/cleanup.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace serving {
namespace anonymous {
class FakeTask : public BatchTask {
public:
explicit FakeTask(size_t size) : size_(size) {}
~FakeTask() override = default;
size_t size() const override { return size_; }
void set_size(size_t size) { size_ = size; }
private:
size_t size_;
FakeTask(const FakeTask&) = delete;
void operator=(const FakeTask&) = delete;
};
// Creates a FakeTask of size 'task_size', and calls 'scheduler->Schedule()' on
// that task. Returns the resulting status.
absl::Status ScheduleTask(size_t task_size,
BatchScheduler<FakeTask>* scheduler) {
std::unique_ptr<FakeTask> task(new FakeTask(task_size));
absl::Status status = scheduler->Schedule(&task);
// Schedule() should have consumed 'task' iff it returned Status::OK.
CHECK_EQ(status.ok(), task == nullptr);
return status;
}
// Creates a thread that waits on 'start' and then advances the fake clock in
// 'env' in a loop until 'stop' is notified. Useful for allowing objects that
// use the clock to be destroyed.
std::unique_ptr<Thread> CreateFakeClockAdvancerThread(
test_util::FakeClockEnv* env, absl::Notification* start,
absl::Notification* stop) {
return std::unique_ptr<Thread>(Env::Default()->StartThread(
{}, "FakeClockAdvancerThread", [env, start, stop] {
start->WaitForNotification();
while (!stop->HasBeenNotified()) {
env->AdvanceByMicroseconds(10);
Env::Default()->SleepForMicroseconds(10);
}
}));
}
TEST(AdaptiveSharedBatchSchedulerTest, BadOptions) {
using Scheduler = AdaptiveSharedBatchScheduler<FakeTask>;
std::shared_ptr<Scheduler> scheduler;
Scheduler::Options options;
options.num_batch_threads = 0;
EXPECT_FALSE(Scheduler::Create(options, &scheduler).ok());
options = Scheduler::Options();
options.initial_in_flight_batches_limit = 0.5;
EXPECT_FALSE(Scheduler::Create(options, &scheduler).ok());
options = Scheduler::Options();
options.num_batch_threads = 5;
options.initial_in_flight_batches_limit = 8;
EXPECT_FALSE(Scheduler::Create(options, &scheduler).ok());
options = Scheduler::Options();
options.batches_to_average_over = -5;
EXPECT_FALSE(Scheduler::Create(options, &scheduler).ok());
options = Scheduler::Options();
options.min_in_flight_batches_limit = 0;
EXPECT_FALSE(Scheduler::Create(options, &scheduler).ok());
options = Scheduler::Options();
options.min_in_flight_batches_limit = 5;
options.num_batch_threads = 3;
EXPECT_FALSE(Scheduler::Create(options, &scheduler).ok());
options = Scheduler::Options();
options.initial_in_flight_batches_limit = 1;
options.min_in_flight_batches_limit = 2;
options.num_batch_threads = 3;
EXPECT_FALSE(Scheduler::Create(options, &scheduler).ok());
}
TEST(AdaptiveSharedBatchSchedulerTest, InFlightBatchesLimit) {
AdaptiveSharedBatchScheduler<FakeTask>::Options options;
options.initial_in_flight_batches_limit = 2;
options.batches_to_average_over = 1000;
mutex mu;
int processed_batches = 0;
absl::Notification finish_processing;
auto queue_callback = [&mu, &processed_batches, &finish_processing](
std::unique_ptr<Batch<FakeTask>> batch) {
ASSERT_TRUE(batch->IsClosed());
EXPECT_GT(batch->num_tasks(), 0);
mu.lock();
int batch_num = ++processed_batches;
mu.unlock();
if (batch_num == 2) {
// Give third batch a chance to process if it's going to.
Env::Default()->SleepForMicroseconds(1000);
finish_processing.Notify();
}
if (batch_num == 3) {
ASSERT_TRUE(finish_processing.HasBeenNotified());
}
finish_processing.WaitForNotification();
};
std::shared_ptr<AdaptiveSharedBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(
AdaptiveSharedBatchScheduler<FakeTask>::Create(options, &scheduler));
std::unique_ptr<BatchScheduler<FakeTask>> queue;
TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue));
// Enqueue 3 batches.
TF_ASSERT_OK(ScheduleTask(100, queue.get()));
while (queue->NumEnqueuedTasks() > 0) {
}
TF_ASSERT_OK(ScheduleTask(100, queue.get()));
while (queue->NumEnqueuedTasks() > 0) {
}
TF_ASSERT_OK(ScheduleTask(100, queue.get()));
}
TEST(AdaptiveSharedBatchSchedulerTest, InFlightBatchesLimitTuning) {
test_util::FakeClockEnv env(Env::Default());
absl::Notification start_teardown, stop_teardown;
std::unique_ptr<Thread> teardown_thread =
CreateFakeClockAdvancerThread(&env, &start_teardown, &stop_teardown);
{
AdaptiveSharedBatchScheduler<FakeTask>::Options options;
options.env = &env;
options.initial_in_flight_batches_limit = 2;
options.batches_to_average_over = 1;
auto queue_callback = [&env](std::unique_ptr<Batch<FakeTask>> batch) {
ASSERT_TRUE(batch->IsClosed());
switch (batch->size()) {
case 0:
env.AdvanceByMicroseconds(10);
break;
case 1:
env.AdvanceByMicroseconds(15);
break;
case 2:
env.AdvanceByMicroseconds(10);
break;
case 3:
env.AdvanceByMicroseconds(11);
break;
}
};
std::shared_ptr<AdaptiveSharedBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(
AdaptiveSharedBatchScheduler<FakeTask>::Create(options, &scheduler));
std::unique_ptr<BatchScheduler<FakeTask>> queue;
TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue));
TF_ASSERT_OK(ScheduleTask(0, queue.get()));
double in_flight_batches_limit = 2;
while (scheduler->in_flight_batches_limit() == in_flight_batches_limit) {
}
// Initial direction will be negative.
EXPECT_LT(scheduler->in_flight_batches_limit(), in_flight_batches_limit);
in_flight_batches_limit = scheduler->in_flight_batches_limit();
TF_ASSERT_OK(ScheduleTask(1, queue.get()));
while (scheduler->in_flight_batches_limit() == in_flight_batches_limit) {
}
// Latency increased -> change direction.
EXPECT_GT(scheduler->in_flight_batches_limit(), in_flight_batches_limit);
in_flight_batches_limit = scheduler->in_flight_batches_limit();
TF_ASSERT_OK(ScheduleTask(2, queue.get()));
while (scheduler->in_flight_batches_limit() == in_flight_batches_limit) {
}
// Latency decreased -> keep going in same direction.
EXPECT_GT(scheduler->in_flight_batches_limit(), in_flight_batches_limit);
in_flight_batches_limit = scheduler->in_flight_batches_limit();
TF_ASSERT_OK(ScheduleTask(3, queue.get()));
while (scheduler->in_flight_batches_limit() == in_flight_batches_limit) {
}
// Latency increased -> change direction.
EXPECT_LT(scheduler->in_flight_batches_limit(), in_flight_batches_limit);
start_teardown.Notify();
}
stop_teardown.Notify();
}
TEST(AdaptiveSharedBatchSchedulerTest, FullBatchSchedulingBoostMicros) {
test_util::FakeClockEnv env(Env::Default());
absl::Notification start_teardown, stop_teardown;
std::unique_ptr<Thread> teardown_thread =
CreateFakeClockAdvancerThread(&env, &start_teardown, &stop_teardown);
{
AdaptiveSharedBatchScheduler<FakeTask>::Options options;
options.env = &env;
options.initial_in_flight_batches_limit = 1;
options.num_batch_threads = 1;
options.batches_to_average_over = 1000;
options.full_batch_scheduling_boost_micros = 100;
mutex mu;
int processed_batches = 0;
absl::Notification finish_processing;
auto queue_callback = [&mu, &processed_batches, &finish_processing](
std::unique_ptr<Batch<FakeTask>> batch) {
ASSERT_TRUE(batch->IsClosed());
finish_processing.WaitForNotification();
mutex_lock l(mu);
processed_batches++;
switch (processed_batches) {
case 1:
EXPECT_EQ(100, batch->size());
break;
case 2:
EXPECT_EQ(50, batch->size());
break;
case 3:
EXPECT_EQ(900, batch->size());
break;
case 4:
EXPECT_EQ(200, batch->size());
break;
default:
EXPECT_TRUE(false) << "Should only have 4 batches";
}
};
std::shared_ptr<AdaptiveSharedBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(
AdaptiveSharedBatchScheduler<FakeTask>::Create(options, &scheduler));
AdaptiveSharedBatchScheduler<FakeTask>::QueueOptions queue_options;
std::unique_ptr<BatchScheduler<FakeTask>> queue1;
std::unique_ptr<BatchScheduler<FakeTask>> queue2;
queue_options.max_batch_size = 1000;
TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue1));
queue_options.max_batch_size = 100;
TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue2));
// First batch immediately processed.
TF_ASSERT_OK(ScheduleTask(100, queue1.get()));
while (queue1->NumEnqueuedTasks() > 0) {
}
TF_ASSERT_OK(ScheduleTask(100, queue1.get()));
env.AdvanceByMicroseconds(10);
TF_ASSERT_OK(ScheduleTask(100, queue1.get()));
env.AdvanceByMicroseconds(10);
TF_ASSERT_OK(ScheduleTask(50, queue2.get()));
env.AdvanceByMicroseconds(45);
TF_ASSERT_OK(ScheduleTask(900, queue1.get()));
// Second batch - creation time: 0, fullness: 0.2, sched score: -20
// Third batch - creation time: 20, fullness: 0.5, sched score: -30
// Fourth batch - creation time: 65, fullness: 0.9, sched score: -25
finish_processing.Notify();
start_teardown.Notify();
}
stop_teardown.Notify();
}
TEST(AdaptiveSharedBatchSchedulerTest, FIFO) {
test_util::FakeClockEnv env(Env::Default());
absl::Notification start_teardown, stop_teardown;
std::unique_ptr<Thread> teardown_thread =
CreateFakeClockAdvancerThread(&env, &start_teardown, &stop_teardown);
{
AdaptiveSharedBatchScheduler<FakeTask>::Options options;
options.env = &env;
options.initial_in_flight_batches_limit = 1;
options.num_batch_threads = 1;
options.batches_to_average_over = 1000;
options.full_batch_scheduling_boost_micros = 0;
options.fifo_scheduling = true;
mutex mu;
int processed_batches = 0;
absl::Notification finish_processing;
auto queue_callback = [&mu, &processed_batches, &finish_processing](
std::unique_ptr<Batch<FakeTask>> batch) {
ASSERT_TRUE(batch->IsClosed());
finish_processing.WaitForNotification();
mutex_lock l(mu);
processed_batches++;
switch (processed_batches) {
case 1:
EXPECT_EQ(100, batch->size());
break;
case 2:
EXPECT_EQ(200, batch->size());
break;
case 3:
EXPECT_EQ(50, batch->size());
break;
case 4:
EXPECT_EQ(900, batch->size());
break;
default:
EXPECT_TRUE(false) << "Should only have 4 batches";
}
};
std::shared_ptr<AdaptiveSharedBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(
AdaptiveSharedBatchScheduler<FakeTask>::Create(options, &scheduler));
AdaptiveSharedBatchScheduler<FakeTask>::QueueOptions queue_options;
std::unique_ptr<BatchScheduler<FakeTask>> queue1;
std::unique_ptr<BatchScheduler<FakeTask>> queue2;
queue_options.max_batch_size = 1000;
TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue1));
queue_options.max_batch_size = 100;
TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue2));
// First batch immediately processed.
TF_ASSERT_OK(ScheduleTask(100, queue1.get()));
env.AdvanceByMicroseconds(30);
TF_ASSERT_OK(ScheduleTask(100, queue1.get()));
env.AdvanceByMicroseconds(10);
TF_ASSERT_OK(ScheduleTask(100, queue1.get()));
env.AdvanceByMicroseconds(10);
TF_ASSERT_OK(ScheduleTask(50, queue2.get()));
env.AdvanceByMicroseconds(45);
TF_ASSERT_OK(ScheduleTask(900, queue1.get()));
finish_processing.Notify();
start_teardown.Notify();
}
stop_teardown.Notify();
}
TEST(AdaptiveSharedBatchSchedulerTest, DeleteQueue) {
AdaptiveSharedBatchScheduler<FakeTask>::Options options;
options.initial_in_flight_batches_limit = 1;
options.num_batch_threads = 1;
options.batches_to_average_over = 1000;
mutex mu;
int processed_batches = 0;
absl::Notification finish_processing;
auto queue_callback = [&mu, &processed_batches, &finish_processing](
std::unique_ptr<Batch<FakeTask>> batch) {
ASSERT_TRUE(batch->IsClosed());
EXPECT_GT(batch->num_tasks(), 0);
finish_processing.WaitForNotification();
mu.lock();
processed_batches++;
mu.unlock();
};
auto processed_checker = gtl::MakeCleanup([&mu, &processed_batches] {
mutex_lock l(mu);
EXPECT_EQ(processed_batches, 2);
});
std::shared_ptr<AdaptiveSharedBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(
AdaptiveSharedBatchScheduler<FakeTask>::Create(options, &scheduler));
std::unique_ptr<BatchScheduler<FakeTask>> queue;
TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue));
// Enqueue 2 tasks, should result in 2 batches.
TF_ASSERT_OK(ScheduleTask(100, queue.get()));
while (queue->NumEnqueuedTasks() > 0) {
}
TF_ASSERT_OK(ScheduleTask(100, queue.get()));
// Queue destructor should block until second batch has been scheduled.
Env::Default()->SchedClosureAfter(
1000, [&finish_processing] { finish_processing.Notify(); });
}
TEST(AdaptiveSharedBatchSchedulerTest, QueueCapacityInfo) {
AdaptiveSharedBatchScheduler<FakeTask>::Options options;
options.initial_in_flight_batches_limit = 1;
options.batches_to_average_over = 1000;
mutex mu;
int processed_batches = 0;
absl::Notification finish_processing;
auto queue_callback = [&mu, &processed_batches, &finish_processing](
std::unique_ptr<Batch<FakeTask>> batch) {
ASSERT_TRUE(batch->IsClosed());
EXPECT_GT(batch->num_tasks(), 0);
mu.lock();
int batch_num = ++processed_batches;
mu.unlock();
if (batch_num == 1) {
finish_processing.WaitForNotification();
}
};
std::shared_ptr<AdaptiveSharedBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(
AdaptiveSharedBatchScheduler<FakeTask>::Create(options, &scheduler));
std::unique_ptr<BatchScheduler<FakeTask>> queue;
TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue));
// Enqueue 2 tasks, should result in 2 batches.
TF_ASSERT_OK(ScheduleTask(100, queue.get()));
while (queue->NumEnqueuedTasks() > 0) {
}
TF_ASSERT_OK(ScheduleTask(100, queue.get()));
// First batch was immediately processed, no longer counts as enqueued.
EXPECT_EQ(queue->NumEnqueuedTasks(), 1);
EXPECT_EQ(queue->SchedulingCapacity(), 9 * 1000 + 900);
// Enqueue 2 more tasks, should fall in same batch.
TF_ASSERT_OK(ScheduleTask(100, queue.get()));
TF_ASSERT_OK(ScheduleTask(200, queue.get()));
EXPECT_EQ(queue->NumEnqueuedTasks(), 3);
EXPECT_EQ(queue->SchedulingCapacity(), 9 * 1000 + 600);
// Enqueue 1 more task, should create new batch and start processing the
// previous batch.
TF_ASSERT_OK(ScheduleTask(700, queue.get()));
EXPECT_EQ(queue->NumEnqueuedTasks(), 1);
EXPECT_EQ(queue->SchedulingCapacity(), 9 * 1000 + 300);
finish_processing.Notify();
}
TEST(AdaptiveSharedBatchSchedulerTest, FullBatches) {
std::shared_ptr<AdaptiveSharedBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(AdaptiveSharedBatchScheduler<FakeTask>::Create({}, &scheduler));
auto queue_callback = [](std::unique_ptr<Batch<FakeTask>> batch) {
ASSERT_TRUE(batch->IsClosed());
};
AdaptiveSharedBatchScheduler<FakeTask>::QueueOptions queue_options;
queue_options.max_batch_size = 100;
queue_options.batch_timeout_micros = 1000000000000;
std::unique_ptr<BatchScheduler<FakeTask>> queue;
TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue));
TF_ASSERT_OK(ScheduleTask(100, queue.get()));
// Full batches should not have to wait batch_timeout_micros.
}
TEST(AdaptiveSharedBatchSchedulerTest, TruncateBatches) {
mutex mu;
int processed_batches = 0;
auto queue_callback =
[&mu, &processed_batches](std::unique_ptr<Batch<FakeTask>> batch) {
ASSERT_TRUE(batch->IsClosed());
mutex_lock l(mu);
++processed_batches;
};
std::shared_ptr<AdaptiveSharedBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(AdaptiveSharedBatchScheduler<FakeTask>::Create({}, &scheduler));
std::unique_ptr<BatchScheduler<FakeTask>> queue;
AdaptiveSharedBatchScheduler<FakeTask>::QueueOptions queue_options;
queue_options.max_batch_size = 100;
queue_options.batch_timeout_micros = 1000000;
queue_options.split_input_task_func =
[](std::unique_ptr<FakeTask>* input_task, int first_size, int max_size,
std::vector<std::unique_ptr<FakeTask>>* output_tasks) {
EXPECT_EQ(first_size, 70);
output_tasks->push_back(std::move(*input_task));
int remaining_size = output_tasks->back()->size() - first_size;
output_tasks->back()->set_size(first_size);
while (remaining_size > 0) {
int task_size = std::min(remaining_size, max_size);
output_tasks->emplace_back(new FakeTask(task_size));
remaining_size -= task_size;
}
return absl::OkStatus();
};
TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue));
TF_ASSERT_OK(ScheduleTask(30, queue.get()));
TF_ASSERT_OK(ScheduleTask(350, queue.get()));
// Second task should be split into a task of size 70, 2 tasks of size 100,
// and one task of size 80.
while (true) {
mutex_lock l(mu);
if (processed_batches == 4) break;
}
}
TEST(AdaptiveSharedBatchSchedulerTest, MaxTasksPerBatch) {
mutex mu;
int processed_batches = 0;
auto queue_callback =
[&mu, &processed_batches](std::unique_ptr<Batch<FakeTask>> batch) {
ASSERT_TRUE(batch->IsClosed());
mutex_lock l(mu);
++processed_batches;
};
std::shared_ptr<AdaptiveSharedBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(AdaptiveSharedBatchScheduler<FakeTask>::Create({}, &scheduler));
std::unique_ptr<BatchScheduler<FakeTask>> queue;
AdaptiveSharedBatchScheduler<FakeTask>::QueueOptions queue_options;
queue_options.max_batch_size = 100;
queue_options.batch_timeout_micros = 1000000;
queue_options.max_tasks_per_batch = 2;
TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue));
TF_ASSERT_OK(ScheduleTask(10, queue.get()));
// Only one task in the batch, so batch not processed yet.
EXPECT_EQ(queue->NumEnqueuedTasks(), 1);
TF_ASSERT_OK(ScheduleTask(10, queue.get()));
// Two tasks were added to the batch, so batch was processed.
EXPECT_EQ(queue->NumEnqueuedTasks(), 0);
TF_ASSERT_OK(ScheduleTask(10, queue.get()));
TF_ASSERT_OK(ScheduleTask(10, queue.get()));
TF_ASSERT_OK(ScheduleTask(10, queue.get()));
TF_ASSERT_OK(ScheduleTask(10, queue.get()));
// We get 3 batches since only two tasks are allowed per batch.
while (true) {
mutex_lock l(mu);
if (processed_batches == 3) break;
}
}
} // namespace anonymous
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/basic_batch_scheduler.h
+366
View File
@@ -0,0 +1,366 @@
/* Copyright 2016 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BASIC_BATCH_SCHEDULER_H_
#define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BASIC_BATCH_SCHEDULER_H_
#include <stddef.h>
#include <cstddef>
#include <functional>
#include <memory>
#include <string>
#include "tensorflow/core/kernels/batching_util/shared_batch_scheduler.h"
namespace tensorflow {
namespace serving {
// A BatchScheduler implementation geared toward handling a single request type
// running on a specific set of hardware resources. A typical scenario is one in
// which all requests invoke the same machine-learned model on one GPU.
//
// If there are, say, two GPUs and two models each bound to one of the GPUs, one
// could use two BasicBatchScheduler instances to schedule the two model/GPU
// combinations independently. If multiple models must share a given GPU or
// other hardware resource, consider using SharedBatchScheduler instead.
//
//
// PARAMETERS AND BEHAVIOR:
//
// BasicBatchScheduler runs a fixed pool of threads, which it uses to process
// batches of tasks. It enforces a maximum batch size, and enqueues a bounded
// number of tasks. If the queue is nearly empty, such that a full batch cannot
// be formed, when a thread becomes free, it anyway schedules a batch
// immediately if a task has been in the queue for longer than a given timeout
// parameter. If the timeout parameter is set to 0, then the batch threads will
// always be kept busy (unless there are zero tasks waiting to be processed).
//
// For online serving, it is recommended to set the maximum number of enqueued
// batches worth of tasks equal to the number of batch threads, which allows
// enqueuing of enough tasks s.t. if every thread becomes available it can be
// kept busy, but no more. For bulk processing jobs and throughput-oriented
// benchmarks, you may want to set it much higher.
//
// When Schedule() is called, if the queue is full the call will fail with an
// UNAVAILABLE error (after which the client may retry again later). If the call
// succeeds, the maximum time the task will spend in the queue before being
// placed in a batch and assigned to a thread for processing, is the greater of:
// - the maximum time to process ceil(max_enqueued_batches/num_batch_threads)
// (1 in the recommended configuration) batches of previously-submitted tasks
// - the configured timeout parameter (which can be 0, as mentioned above)
//
// Unlike StreamingBatchScheduler, when BasicBatchScheduler assigns a batch to a
// thread, it closes the batch. The process-batch callback may assume that every
// batch it receives is closed at the outset.
//
//
// RECOMMENDED USE-CASES:
//
// BasicBatchScheduler is suitable for use-cases that feature a single kind of
// request (e.g. a server performing inference with a single machine-learned
// model, possibly evolving over time), with loose versioning semantics.
// Concretely, the following conditions should hold:
//
// A. All requests batched onto a given resource (e.g. a hardware accelerator,
// or a pool accelerators) are of the same type. For example, they all
// invoke the same machine-learned model.
//
// These variations are permitted:
// - The model may reside in a single servable, or it may be spread across
// multiple servables that are used in unison (e.g. a vocabulary lookup
// table servable and a tensorflow session servable).
// - The model's servable(s) may be static, or they may evolve over time
// (successive servable versions).
// - Zero or more of the servables are used in the request thread; the rest
// are used in the batch thread. In our running example, the vocabulary
// lookups and tensorflow runs may both be performed in the batch thread,
// or alternatively the vocabulary lookup may occur in the request thread
// with only the tensorflow run performed in the batch thread.
//
// In contrast, BasicBatchScheduler is not a good fit if the server
// hosts multiple distinct models running on a pool accelerators, with each
// request specifying which model it wants to use. BasicBatchScheduler
// has no facility to time-multiplex the batch threads across multiple
// models in a principled way. More basically, it cannot ensure that a given
// batch doesn't contain a mixture of requests for different models.
//
// B. Requests do not specify a particular version of the servable(s) that must
// be used. Instead, each request is content to use the "latest" version.
//
// BasicBatchScheduler does not constrain which requests get grouped
// together into a batch, so using this scheduler there is no way to achieve
// cohesion of versioned requests to version-specific batches.
//
// C. No servable version coordination needs to be performed between the
// request threads and the batch threads. Often, servables are only used in
// the batch threads, in which case this condition trivially holds. If
// servables are used in both threads, then the use-case must tolerate
// version skew across the servables used in the two kinds of threads.
//
//
// EXAMPLE USE-CASE FLOW:
//
// For such use-cases, request processing via BasicBatchScheduler generally
// follows this flow (given for illustration; variations are possible):
// 1. Optionally perform some pre-processing on each request in the request
// threads.
// 2. Route the requests to the batch scheduler, as batching::Task objects.
// (Since all requests are of the same type and are not versioned, the
// scheduler is free to group them into batches arbitrarily.)
// 3. Merge the requests into a single batched representation B.
// 4. Obtain handles to the servable(s) needed to process B. The simplest
// approach is to obtain the latest version of each servable. Alternatively,
// if cross-servable consistency is required (e.g. the vocabulary lookup
// table's version number must match that of the tensorflow session),
// identify an appropriate version number and obtain the servable handles
// accordingly.
// 5. Process B using the obtained servable handles, and split the result into
// individual per-request units.
// 6. Perform any post-processing in the batch thread and/or request thread.
//
//
// PERFORMANCE TUNING: See README.md.
//
template <typename TaskType>
class BasicBatchScheduler : public BatchScheduler<TaskType> {
public:
// TODO(b/25089730): Tune defaults based on best practices as they develop.
// (Keep them mirrored to the ones in SharedBatchScheduler::QueueOptions and
// SharedBatchScheduler::Options.)
struct Options {
// Options related with (underlying) shared batch scheduler.
// 'thread_pool_name' and 'num_batch_threads' are used to initialize
// a shared batch scheduler underlyingly iff 'shared_batch_scheduler' is
// nullptr.
//
// There are two ways to specify threading:
// 1) Have each session create its own pool.
// 2) Have multiple sessions share the same pool.
//
// In general, the number of threads should be tied to roughly the number of
// compute resources (CPU cores or accelerator cores) backing the threads.
// Sharing a thread pool helps alleviate potential over allocation of
// threads to limited compute resources.
// To have each session create its own thread pool (1) set
// thread_pool_name/num_batch_threads.
// To share a thread pool (2) create a scheduler and pass it in.
// The name to use for the pool of batch threads.
std::string thread_pool_name = {"batch_threads"};
// The number of threads to use to process batches.
// Must be >= 1, and should be tuned carefully.
int num_batch_threads = port::MaxParallelism();
// If specified, this scheduler will be used underlyingly to schedule
// batches. Note setting this means `thread_pool_name` and
// `num_batch_threads` are ignored.
std::shared_ptr<SharedBatchScheduler<TaskType>> shared_batch_scheduler =
nullptr;
// Options for queue.
// The maximum size of each batch.
//
// The scheduler may form batches of any size between 1 and this number
// (inclusive). If there is a need to quantize the batch sizes, i.e. only
// submit batches whose size is in a small set of allowed sizes, that can be
// done by adding padding in the process-batch callback.
int max_batch_size = 1000;
// If a task has been enqueued for this amount of time (in microseconds),
// and a thread is available, the scheduler will immediately form a batch
// from enqueued tasks and assign the batch to the thread for processing,
// even if the batch's size is below 'max_batch_size'.
//
// This parameter offers a way to bound queue latency, so that a task isn't
// stuck in the queue indefinitely waiting for enough tasks to arrive to
// make a full batch. (The latency bound is given in the class documentation
// above.)
//
// The goal is to smooth out batch sizes under low request rates, and thus
// avoid latency spikes.
int64_t batch_timeout_micros = 0;
// The maximum allowable number of enqueued (accepted by Schedule() but
// not yet being processed on a batch thread) tasks in terms of batches.
// If this limit is reached, Schedule() will return an UNAVAILABLE error.
// See the class documentation above for guidelines on how to tune this
// parameter.
int max_enqueued_batches = 10;
// If true, an input task (i.e., input of `BasicBatchScheduler::Schedule`)
// with a large size (i.e., larger than the largest value of
// `allowed_batch_sizes`) will be split into multiple smaller batch tasks
// and possibly put into different batches for processing. If false, each
// input task is put into one batch as a whole for processing.
//
// API note:
// The value of this option doesn't affect processing output given the same
// input; it affects implementation details as stated below:
// 1. Improve batching efficiency by eliminating unnecessary padding in the
// following scenario: when an open batch has M slots while an input of size
// N is scheduled (M < N), the input can be split to fill remaining slots
// of an open batch as opposed to padding.
// 2.`max_batch_size` specifies the limit of input and
// `max_execution_batch_size` specifies the limit of a task to be processed.
// API user can give an input of size 128 when 'max_execution_batch_size'
// is 32 -> implementation can split input of 128 into 4 x 32, schedule
// concurrent processing, and then return concatenated results corresponding
// to 128.
bool enable_large_batch_splitting = false;
// `split_input_task_func` specifies how to split `input_task` into
// `output_tasks`.
//
// `input_task`: a unit of task to be split.
// `first_output_task_size`: task size of first output.
// `max_batch_size`: Maximum size of each batch.
// `output_tasks`: A list of output tasks after split.
//
// REQUIRED:
// 1) All `output_tasks` should be non-empty tasks.
// 2) Sizes of `output_tasks` add up to size of `input_task`.
//
// NOTE:
// Instantiations of `TaskType` may vary, so it's up to caller to define
// how (e.g., which members to access) to split input tasks.
std::function<absl::Status(
std::unique_ptr<TaskType>* input_task, int first_output_task_size,
int input_batch_size_limit,
std::vector<std::unique_ptr<TaskType>>* output_tasks)>
split_input_task_func;
// The maximum size of each enqueued batch (i.e., in `batches_`).
//
// The scheduler may form batches of any size between 1 and this number
// (inclusive). If there is a need to quantize the batch sizes, i.e. only
// submit batches whose size is in a small set of allowed sizes, that can be
// done by adding padding in the process-batch callback.
//
// REQUIRES:
// - If enable_large_batch_splitting is true, `max_execution_batch_size` is
// less than or equal to `max_batch_size`.
// - If enable_large_batch_splitting is false, `max_execution_batch_size` is
// equal to `max_batch_size`.
int max_execution_batch_size = 10;
// The following options are typically only overridden by test code.
// The environment to use.
Env* env = Env::Default();
};
static absl::Status Create(
const Options& options,
std::function<void(std::unique_ptr<Batch<TaskType>>)>
process_batch_callback,
std::unique_ptr<BasicBatchScheduler>* scheduler);
~BasicBatchScheduler() override = default;
absl::Status Schedule(std::unique_ptr<TaskType>* task) override;
size_t NumEnqueuedTasks() const override;
size_t SchedulingCapacity() const override;
size_t max_task_size() const override {
return shared_scheduler_queue_->max_task_size();
}
private:
explicit BasicBatchScheduler(
std::unique_ptr<BatchScheduler<TaskType>> shared_scheduler_queue);
// This class is merely a thin wrapper around a SharedBatchScheduler with a
// single queue.
std::unique_ptr<BatchScheduler<TaskType>> shared_scheduler_queue_;
BasicBatchScheduler(const BasicBatchScheduler&) = delete;
void operator=(const BasicBatchScheduler&) = delete;
};
//////////
// Implementation details follow. API users need not read.
template <typename TaskType>
absl::Status BasicBatchScheduler<TaskType>::Create(
const Options& options,
std::function<void(std::unique_ptr<Batch<TaskType>>)>
process_batch_callback,
std::unique_ptr<BasicBatchScheduler>* scheduler) {
std::shared_ptr<SharedBatchScheduler<TaskType>> shared_scheduler;
if (options.shared_batch_scheduler == nullptr) {
typename SharedBatchScheduler<TaskType>::Options shared_scheduler_options;
shared_scheduler_options.thread_pool_name = options.thread_pool_name;
shared_scheduler_options.num_batch_threads = options.num_batch_threads;
shared_scheduler_options.env = options.env;
TF_RETURN_IF_ERROR(SharedBatchScheduler<TaskType>::Create(
shared_scheduler_options, &shared_scheduler));
} else {
shared_scheduler = options.shared_batch_scheduler;
}
typename SharedBatchScheduler<TaskType>::QueueOptions
shared_scheduler_queue_options;
shared_scheduler_queue_options.input_batch_size_limit =
options.max_batch_size;
shared_scheduler_queue_options.batch_timeout_micros =
options.batch_timeout_micros;
shared_scheduler_queue_options.max_enqueued_batches =
options.max_enqueued_batches;
shared_scheduler_queue_options.enable_large_batch_splitting =
options.enable_large_batch_splitting;
shared_scheduler_queue_options.split_input_task_func =
options.split_input_task_func;
shared_scheduler_queue_options.max_execution_batch_size =
options.max_execution_batch_size;
std::unique_ptr<BatchScheduler<TaskType>> shared_scheduler_queue;
TF_RETURN_IF_ERROR(shared_scheduler->AddQueue(shared_scheduler_queue_options,
process_batch_callback,
&shared_scheduler_queue));
scheduler->reset(
new BasicBatchScheduler<TaskType>(std::move(shared_scheduler_queue)));
return absl::OkStatus();
}
template <typename TaskType>
absl::Status BasicBatchScheduler<TaskType>::Schedule(
std::unique_ptr<TaskType>* task) {
return shared_scheduler_queue_->Schedule(task);
}
template <typename TaskType>
size_t BasicBatchScheduler<TaskType>::NumEnqueuedTasks() const {
return shared_scheduler_queue_->NumEnqueuedTasks();
}
template <typename TaskType>
size_t BasicBatchScheduler<TaskType>::SchedulingCapacity() const {
return shared_scheduler_queue_->SchedulingCapacity();
}
template <typename TaskType>
BasicBatchScheduler<TaskType>::BasicBatchScheduler(
std::unique_ptr<BatchScheduler<TaskType>> shared_scheduler_queue)
: shared_scheduler_queue_(std::move(shared_scheduler_queue)) {}
} // namespace serving
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BASIC_BATCH_SCHEDULER_H_
tensorflow/core/kernels/batching_util/basic_batch_scheduler_benchmark_test.cc
+403
View File
@@ -0,0 +1,403 @@
/* Copyright 2016 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.
==============================================================================*/
// Benchmarks for performance (throughput and latency) of BasicBatchScheduler
// under various rates of task injection.
#include <memory>
#include "tensorflow/core/kernels/batching_util/basic_batch_scheduler.h"
#include "tensorflow/core/lib/histogram/histogram.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/util/command_line_flags.h"
namespace tensorflow {
namespace serving {
namespace {
using ::tensorflow::histogram::Histogram;
// Fixed duration to run latency benchmark.
static int latency_benchmark_duration_secs = 100;
// An abstract class for injecting load into a system at a specific rate.
class LoadInjector {
public:
virtual ~LoadInjector() = default;
// Run 'injector' 'num_injection' times, with average inter-injection spacing
// as 'average_injection_interval_micros' (in microseconds).
virtual void InjectLoad(std::function<void()> injector, int num_injections,
int64_t average_injection_interval_micros) const = 0;
};
// A load injector that uses uniform inter-injection spacing, i.e. each pair of
// injections is separated in time by 'average_injection_interval_micros' (as
// best as possible).
class UniformLoadInjector : public LoadInjector {
public:
UniformLoadInjector() = default;
~UniformLoadInjector() override = default;
void InjectLoad(std::function<void()> injector, int num_injections,
int64_t average_injection_interval_micros) const override;
private:
UniformLoadInjector(const UniformLoadInjector&) = delete;
void operator=(const UniformLoadInjector&) = delete;
};
void UniformLoadInjector::InjectLoad(
std::function<void()> injector, const int num_injections,
const int64_t average_injection_interval_micros) const {
int num_injections_performed = 0;
const int64_t start_time_micros = Env::Default()->NowMicros();
while (num_injections_performed < num_injections) {
// Inject.
injector();
++num_injections_performed;
// Wait until it's time for the next injection.
const int64_t next_injection_time_micros =
start_time_micros +
(num_injections_performed * average_injection_interval_micros);
int64_t now_micros = Env::Default()->NowMicros();
while (now_micros < next_injection_time_micros) {
const int64_t kSleepThresholdMicros = 1000;
if (next_injection_time_micros - now_micros >= kSleepThresholdMicros) {
Env::Default()->SleepForMicroseconds(1 /* minimum time */);
}
now_micros = Env::Default()->NowMicros();
}
}
}
class BenchmarkBatchTask : public BatchTask {
public:
BenchmarkBatchTask();
BenchmarkBatchTask(const BenchmarkBatchTask&) = delete;
BenchmarkBatchTask& operator=(const BenchmarkBatchTask&) = delete;
~BenchmarkBatchTask() override = default;
size_t size() const override { return 1; }
uint64_t start_time_micros() const { return start_time_micros_; }
private:
// The time at which the task was created, in microseconds.
const uint64_t start_time_micros_;
};
BenchmarkBatchTask::BenchmarkBatchTask()
: start_time_micros_(Env::Default()->NowMicros()) {}
// The state associated with a throughput benchmark
class ThroughputBenchmark {
public:
explicit ThroughputBenchmark(
const BasicBatchScheduler<BenchmarkBatchTask>::Options& scheduler_options)
: scheduler_options_(scheduler_options) {
auto process_batch_callback =
[this](std::unique_ptr<Batch<BenchmarkBatchTask>> batch) {
ProcessBatch(std::move(batch));
};
TF_CHECK_OK(BasicBatchScheduler<BenchmarkBatchTask>::Create(
scheduler_options_, process_batch_callback, &scheduler_));
}
ThroughputBenchmark(const ThroughputBenchmark&) = delete;
ThroughputBenchmark& operator=(const ThroughputBenchmark&) = delete;
BasicBatchScheduler<BenchmarkBatchTask>* GetScheduler() const {
return scheduler_.get();
}
// Reset scheduler. This has a side-effect of waiting for all work to be
// completed prior to reset.
void ResetScheduler() { return scheduler_.reset(); }
private:
// Processes a batch of tasks. (Invoked by 'scheduler_' on one of its batch
// threads.)
void ProcessBatch(std::unique_ptr<Batch<BenchmarkBatchTask>> batch) {
// No-op
}
// Parameters for the BasicBatchScheduler being benchmarked.
const BasicBatchScheduler<BenchmarkBatchTask>::Options scheduler_options_;
// The BasicBatchScheduler being benchmarked.
std::unique_ptr<BasicBatchScheduler<BenchmarkBatchTask>> scheduler_;
};
// The state associated with a latency benchmark, which injects tasks into a
// batch scheduler at a controlled rate and measures the distribution of task
// completion latencies.
class LatencyBenchmark {
public:
LatencyBenchmark(
const BasicBatchScheduler<BenchmarkBatchTask>::Options& scheduler_options,
int64_t task_injection_interval_micros, int batch_cpu_cost);
LatencyBenchmark(const LatencyBenchmark&) = delete;
LatencyBenchmark& operator=(const LatencyBenchmark&) = delete;
// Inject tasks at specified rate for `latency_benchmark_duration_secs`.
void InjectLoad();
// Return latency and batch size stat.
std::string ReportLatencyBatchSz();
// Reset scheduler. This has a side-effect of waiting for all work to be
// completed prior to reset.
void ResetScheduler() { return scheduler_.reset(); }
private:
// Processes a batch of tasks. (Invoked by 'scheduler_' on one of its batch
// threads.)
void ProcessBatch(std::unique_ptr<Batch<BenchmarkBatchTask>> batch);
// Performs one batch's dummy CPU work.
void PerformBatchCpuWork() const;
// Parameters for the BasicBatchScheduler being benchmarked.
const BasicBatchScheduler<BenchmarkBatchTask>::Options scheduler_options_;
// The time interval between successively injected tasks, in microseconds.
// A large interval corresponds to a slow rate of task injection, and vice-
// versa.
const int64_t task_injection_interval_micros_;
// The amount of work to do while processing one batch of tasks. (The cost is
// independent of the number of tasks in the batch.)
const int batch_cpu_cost_;
// The BasicBatchScheduler being benchmarked.
std::unique_ptr<BasicBatchScheduler<BenchmarkBatchTask>> scheduler_;
mutable mutex mu_;
// A histogram of the task latencies, i.e. queue time plus processing time, in
// milliseconds.
Histogram task_latency_millis_histogram_ TF_GUARDED_BY(mu_);
// A histogram of the batch sizes.
Histogram batch_size_histogram_ TF_GUARDED_BY(mu_);
};
LatencyBenchmark::LatencyBenchmark(
const BasicBatchScheduler<BenchmarkBatchTask>::Options& scheduler_options,
int64_t task_injection_interval_micros, int batch_cpu_cost)
: scheduler_options_(scheduler_options),
task_injection_interval_micros_(task_injection_interval_micros),
batch_cpu_cost_(batch_cpu_cost) {
auto process_batch_callback =
[this](std::unique_ptr<Batch<BenchmarkBatchTask>> batch) {
ProcessBatch(std::move(batch));
};
TF_CHECK_OK(BasicBatchScheduler<BenchmarkBatchTask>::Create(
scheduler_options_, process_batch_callback, &scheduler_));
}
void LatencyBenchmark::InjectLoad() {
// Arrange to inject tasks at the specified rate, for a total duration of
// of kTimeDurationMicros.
const int kTimeDurationMicros = latency_benchmark_duration_secs * 1000 * 1000;
const int kNumTasks = kTimeDurationMicros / task_injection_interval_micros_;
const int64_t start_time_micros = Env::Default()->NowMicros();
// Inject the tasks.
UniformLoadInjector injector;
injector.InjectLoad(
[this] {
auto task = std::make_unique<BenchmarkBatchTask>();
TF_CHECK_OK(scheduler_->Schedule(&task));
},
kNumTasks, task_injection_interval_micros_);
// Be sure we were able to more-or-less match our target injection rate.
const int64_t target_injection_time_micros =
kNumTasks * task_injection_interval_micros_;
const int64_t actual_injection_time_micros =
Env::Default()->NowMicros() - start_time_micros;
if (actual_injection_time_micros > 1.1 * target_injection_time_micros) {
LOG(FATAL) << "Unable to inject tasks at the requested rate";
}
// Be sure the scheduler was able to process the tasks at close to the
// injection rate. If not, our latency measurements will be dominated by queue
// waiting time
const int64_t actual_processing_time_micros =
Env::Default()->NowMicros() - start_time_micros;
if (actual_processing_time_micros > 1.01 * actual_injection_time_micros) {
LOG(FATAL) << "Unable to keep up with task injection rate";
}
}
void LatencyBenchmark::ProcessBatch(
std::unique_ptr<Batch<BenchmarkBatchTask>> batch) {
PerformBatchCpuWork();
const uint64_t batch_completion_time = Env::Default()->NowMicros();
{
mutex_lock l(mu_);
batch_size_histogram_.Add(batch->num_tasks());
}
for (int i = 0; i < batch->num_tasks(); ++i) {
const uint64_t task_latency_micros =
batch_completion_time - batch->task(i).start_time_micros();
{
mutex_lock l(mu_);
task_latency_millis_histogram_.Add(task_latency_micros / 1000.0);
}
}
}
void LatencyBenchmark::PerformBatchCpuWork() const {
int dummy = 1;
for (int i = 0; i < batch_cpu_cost_; ++i) {
dummy += dummy * 2;
}
CHECK_NE(dummy, 0);
}
std::string LatencyBenchmark::ReportLatencyBatchSz() {
mutex_lock l(mu_);
return absl::StrCat(
"lat_p99.9=", task_latency_millis_histogram_.Percentile(99.9),
"ms,batchsz_p99=", batch_size_histogram_.Percentile(99));
}
// Injects a large number of tasks into a batch scheduler and measures
// the total time to process all the tasks.
//
// Multi-threaded (thread > 1) version simulates N concurrent request streams.
void ThroughputBM(::testing::benchmark::State& state) {
static std::unique_ptr<ThroughputBenchmark> bm;
if (state.thread_index() == 0) {
BasicBatchScheduler<BenchmarkBatchTask>::Options scheduler_options;
const int kMaxBatchSize = 100;
scheduler_options.max_batch_size = kMaxBatchSize;
scheduler_options.batch_timeout_micros = state.range(0) * 1000;
scheduler_options.num_batch_threads = state.range(1);
scheduler_options.max_enqueued_batches = INT_MAX; // Unbounded queue.
bm = std::make_unique<ThroughputBenchmark>(scheduler_options);
}
// Have each iteration issue a reasonably large number of tasks, to ensure our
// measurements reflect steady-state behavior.
const int kNumTasksPerIteration = 100 * 1000;
// Schedule 'num_iterations_*kNumTasksPerIteration' tasks.
for (auto s : state) {
for (int j = 0; j < kNumTasksPerIteration; ++j) {
auto task = std::make_unique<BenchmarkBatchTask>();
TF_CHECK_OK(bm->GetScheduler()->Schedule(&task));
}
}
if (state.thread_index() == 0) {
state.ResumeTiming();
// Wait for the scheduler to process all tasks.
bm->ResetScheduler();
state.PauseTiming();
bm.reset();
}
state.SetItemsProcessed(state.iterations() * kNumTasksPerIteration);
}
BENCHMARK(ThroughputBM)
->UseRealTime()
->Threads(1)
->Threads(8)
->Threads(16)
->ArgNames({"timeout", "batch_threads"})
->ArgsProduct({{0, 2, 10}, {1, 4, 8, 16}});
// Latency benchmark is a long running fixed interval (by time) benchmark and is
// run once (see ->Iterations(1) below). We measure and report latency over this
// fixed interval.
//
// Multi-threaded (thread > 1) version simulates N concurrent request streams.
void LatencyBM(::testing::benchmark::State& state) {
static std::unique_ptr<LatencyBenchmark> bm;
if (state.thread_index() == 0) {
BasicBatchScheduler<BenchmarkBatchTask>::Options scheduler_options;
const int kMaxBatchSize = 100;
scheduler_options.max_batch_size = kMaxBatchSize;
scheduler_options.batch_timeout_micros = state.range(0);
scheduler_options.num_batch_threads = state.range(1);
scheduler_options.max_enqueued_batches = INT_MAX; // Unbounded queue.
const int kBatchCpuCost = 10 * 1000 * 1000;
const int64_t kQps = state.range(2);
const int64_t kInjectionIntervalMicros = 1000000 / (kQps / state.threads());
const int64_t kNumTasks = latency_benchmark_duration_secs * kQps;
if (kNumTasks <= 10000) {
LOG(WARNING) << "Not enough tasks (" << kNumTasks << ")"
<< " to report meaningful 99.9% latency!"
<< " duration: " << latency_benchmark_duration_secs
<< " interval: " << kInjectionIntervalMicros;
}
bm = std::make_unique<LatencyBenchmark>(
scheduler_options, kInjectionIntervalMicros, kBatchCpuCost);
}
for (auto s : state) {
bm->InjectLoad();
}
if (state.thread_index() == 0) {
state.ResumeTiming();
// Wait for the scheduler to process all tasks.
bm->ResetScheduler();
state.PauseTiming();
state.SetLabel(bm->ReportLatencyBatchSz());
bm.reset();
}
}
BENCHMARK(LatencyBM)
->UseRealTime()
->Iterations(1)
->Threads(1)
->Threads(8)
->Threads(16)
->ArgNames({"timeout", "batch_threads", "qps"})
->ArgsProduct({{0, 2, 10}, {1, 4, 8, 16}, {50000, 20000, 1000}});
} // namespace
} // namespace serving
} // namespace tensorflow
int main(int argc, char** argv) {
const std::vector<tensorflow::Flag> flag_list = {tensorflow::Flag(
"scheduler_latency_bm_fixed_duration_secs",
&tensorflow::serving::latency_benchmark_duration_secs,
"Fixed duration that the latency benchmark must be run.")};
if (!tensorflow::Flags::Parse(&argc, argv, flag_list)) {
std::cout << tensorflow::Flags::Usage(argv[0], flag_list);
return -1;
}
::benchmark::Initialize(&argc, argv);
tensorflow::port::InitMain(argv[0], &argc, &argv);
::benchmark::RunSpecifiedBenchmarks();
return 0;
}
tensorflow/core/kernels/batching_util/basic_batch_scheduler_test.cc
+93
View File
@@ -0,0 +1,93 @@
/* Copyright 2016 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/batching_util/basic_batch_scheduler.h"
#include <utility>
#include "tensorflow/core/kernels/batching_util/batch_scheduler.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace serving {
namespace {
class FakeTask : public BatchTask {
public:
explicit FakeTask(size_t size) : size_(size) {}
~FakeTask() override = default;
size_t size() const override { return size_; }
private:
const size_t size_;
FakeTask(const FakeTask&) = delete;
void operator=(const FakeTask&) = delete;
};
// Creates a FakeTask of size 'task_size', and calls 'scheduler->Schedule()'
// on that task. Returns the resulting status.
absl::Status ScheduleTask(size_t task_size,
BatchScheduler<FakeTask>* scheduler) {
std::unique_ptr<FakeTask> task(new FakeTask(task_size));
absl::Status status = scheduler->Schedule(&task);
// Schedule() should have consumed 'task' iff it returned Status::OK.
CHECK_EQ(status.ok(), task == nullptr);
return status;
}
// Since BasicBatchScheduler is implemented as a thin wrapper around
// SharedBatchScheduler, we only do some basic testing. More comprehensive
// testing is done in shared_batch_scheduler_test.cc.
TEST(BasicBatchSchedulerTest, Basic) {
bool callback_called = false;
auto callback = [&callback_called](std::unique_ptr<Batch<FakeTask>> batch) {
callback_called = true;
ASSERT_TRUE(batch->IsClosed());
ASSERT_EQ(2, batch->num_tasks());
EXPECT_EQ(3, batch->task(0).size());
EXPECT_EQ(5, batch->task(1).size());
};
{
BasicBatchScheduler<FakeTask>::Options options;
options.max_batch_size = 10;
options.batch_timeout_micros = 100 * 1000; // 100 milliseconds
options.num_batch_threads = 1;
options.max_enqueued_batches = 3;
std::unique_ptr<BasicBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(
BasicBatchScheduler<FakeTask>::Create(options, callback, &scheduler));
EXPECT_EQ(10, scheduler->max_task_size());
EXPECT_EQ(0, scheduler->NumEnqueuedTasks());
EXPECT_EQ(3 * 10, scheduler->SchedulingCapacity());
TF_ASSERT_OK(ScheduleTask(3, scheduler.get()));
EXPECT_EQ(1, scheduler->NumEnqueuedTasks());
EXPECT_EQ((3 * 10) - 3, scheduler->SchedulingCapacity());
TF_ASSERT_OK(ScheduleTask(5, scheduler.get()));
EXPECT_EQ(2, scheduler->NumEnqueuedTasks());
EXPECT_EQ((3 * 10) - (3 + 5), scheduler->SchedulingCapacity());
}
EXPECT_TRUE(callback_called);
}
} // namespace
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/batch_input_task.h
+267
View File
@@ -0,0 +1,267 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_INPUT_TASK_H_
#define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_INPUT_TASK_H_
#include <algorithm>
#include <atomic>
#include <functional>
#include <memory>
#include <utility>
#include "absl/base/call_once.h"
#include "absl/container/fixed_array.h"
#include "absl/synchronization/mutex.h"
#include "tensorflow/core/kernels/batching_util/batch_scheduler.h"
#include "tensorflow/core/kernels/batching_util/concat_split_util.h"
#include "tensorflow/core/kernels/batching_util/input_split_metadata.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/util/incremental_barrier.h"
namespace tensorflow {
namespace serving {
namespace internal {
template <typename TaskType>
class BatchInputTaskHandleTestAccess;
template <typename TaskType>
class BatchInputTaskTestAccess;
template <typename TaskType>
class BatchInputTask;
// A RAII-style object that holds a ref-counted batch-input-task, and
// represents a slice of batch-input-task.
// To be handed out to callers of `BatchInputTask::ToTaskHandles` quickly
// (i.e. not necessarily waiting for input split)
//
// `BatchInputTaskHandle::GetSplitTask` evaluates to the slice of task.
template <typename TaskType>
class BatchInputTaskHandle : public BatchTask {
public:
BatchInputTaskHandle(
std::shared_ptr<BatchInputTask<TaskType>> batch_input_task, int split_id,
size_t task_size);
// Should be called once. Returns nullptr on subsequent calls.
std::unique_ptr<TaskType> GetSplitTask();
// Returns the size of this task.
size_t size() const override { return task_size_; }
private:
template <typename T>
friend class internal::BatchInputTaskHandleTestAccess;
int split_id() const { return split_id_; }
std::shared_ptr<BatchInputTask<TaskType>> batch_input_task_;
// The handle evaluates to the N-th slice of original task, and
// N is `split_id_`.
const int split_id_;
const size_t task_size_;
std::atomic<bool> once_{false};
};
// BatchInputTask encapsulates a input (`input_task`) to be batched and the
// information to get task splits after it's enqueued, so as to support lazy
// split of a task.
//
// Input split could reduce excessive padding for efficiency; lazy split
// moves task-split out of the critical path of enqueue and dequeue and reduces
// contention.
//
// BatchInputTask is thread safe.
//
// Usage
//
// ... a deque with frequent enqueue and dequeue operations ...
// ... Note, a deque of Batch of BatchInputTaskHandle is used to form batches
// at enqueue time (split is lazy at deque time);
// ... For use cases to form batches at dequeue time, we can use a deque of
// BatchInputTaskHandle directly, and "peek" metadata to form a batch by
// then.
// std::deque<std::unique_ptr<Batch<BatchInputTaskHandle<TaskType>>>> deque_
// TF_GUARDED_BY(mu_);
//
// std::unique_ptr<TaskType> input_task;
//
// ... Enqueue path ...
//
// {
// mutex_lock l(mu_);
// std::shared_ptr<BatchInputTask<TaskType>> batch_input_task =
// ConstructLazyBatchWithoutSplit(input_task);
//
// std::vector<std::unique_ptr<BatchInputTaskHandle<TaskType>>> task_handles;
// input_batch->ToTaskHandles(&task_handles);
// for (int i = 0; i < task_handles.size(); ++i) {
// EnqueueTaskHandleIntoDeque(deque_);
// }
//
// ... Dequeue path ...
// std::unique_ptr<Batch<BatchInputTaskHandle<TaskType>>> handles_to_schedule;
// {
// mutex_lock l(mu_);
// ... HasBatchToSchedule could be customized or specialized
// ... (e.g., readiness depending on enqueue time)
// if (HasBatchToSchedule(deque_)) {
// handles_to_schedule = std::move(deque_.front());
// deque_.pop_front();
// }
// }
// ...... `mu_` is released ......
//
// std::vector<std::unique_ptr<BatchInputTaskHandle<TaskType>>> tasks_in_batch =
// RemoveAllTasksFromBatch(handles_to_schedule);
//
// std::unique_ptr<Batch<TaskType>> batch_to_schedule;
// for (int i = 0; i < tasks_in_batch.size(); i++) {
// batch_to_schedule->AddTask(std::move(tasks_in_batch[i]->GetSplitTask()));
// }
// batch_to_schedule->Close();
//
// `batch_to_schedule` is ready for schedule.
template <typename TaskType>
class BatchInputTask
: public std::enable_shared_from_this<BatchInputTask<TaskType>> {
public:
using SplitInputFunc = std::function<absl::Status(
std::unique_ptr<TaskType>* input_task, int first_output_task_size,
int input_batch_size_limit,
std::vector<std::unique_ptr<TaskType>>* output_tasks)>;
BatchInputTask(std::unique_ptr<TaskType> input_task,
int open_batch_remaining_slot, int batch_size_limit,
SplitInputFunc split_input_func);
// Outputs the task handles for the input task.
// Each task handle represents a slice of task after input task is split, and
// could evaluate to that slice.
//
// NOTE:
// Each task handle in `output_task_handles` takes ownership of a reference of
// this BatchInputTask.
void ToTaskHandles(
std::vector<std::unique_ptr<BatchInputTaskHandle<TaskType>>>*
output_task_handles);
private:
friend class BatchInputTaskHandle<TaskType>;
template <typename T>
friend class internal::BatchInputTaskTestAccess;
std::unique_ptr<TaskType> GetSplitTask(int split_id);
absl::Status SplitBatches(
std::vector<std::unique_ptr<TaskType>>* output_tasks);
std::unique_ptr<TaskType> input_task_;
const int input_task_size_ = 0;
const int open_batch_remaining_slot_;
const int batch_size_limit_;
const SplitInputFunc split_func_;
const InputSplitMetadata input_split_metadata_;
mutable absl::once_flag once_;
std::vector<std::unique_ptr<TaskType>> task_splits_;
absl::Status split_status_;
};
//
// Implementation details. API readers may skip.
//
template <typename TaskType>
BatchInputTaskHandle<TaskType>::BatchInputTaskHandle(
std::shared_ptr<BatchInputTask<TaskType>> batch_input_task, int split_id,
size_t task_size)
: batch_input_task_(batch_input_task),
split_id_(split_id),
task_size_(task_size) {}
template <typename TaskType>
std::unique_ptr<TaskType> BatchInputTaskHandle<TaskType>::GetSplitTask() {
if (once_.load(std::memory_order_acquire)) {
return nullptr;
}
once_.store(true, std::memory_order_release);
return batch_input_task_->GetSplitTask(split_id_);
}
template <typename TaskType>
BatchInputTask<TaskType>::BatchInputTask(std::unique_ptr<TaskType> input_task,
int open_batch_remaining_slot,
int batch_size_limit,
SplitInputFunc split_input_func)
: input_task_(std::move(input_task)),
input_task_size_(input_task_->size()),
open_batch_remaining_slot_(open_batch_remaining_slot),
batch_size_limit_(batch_size_limit),
split_func_(split_input_func),
input_split_metadata_(input_task_size_, open_batch_remaining_slot,
batch_size_limit) {}
template <typename TaskType>
void BatchInputTask<TaskType>::ToTaskHandles(
std::vector<std::unique_ptr<BatchInputTaskHandle<TaskType>>>*
task_handles) {
const absl::FixedArray<int>& task_sizes = input_split_metadata_.task_sizes();
task_handles->resize(task_sizes.size());
for (int i = 0; i < task_handles->size(); i++) {
(*task_handles)[i] = std::make_unique<BatchInputTaskHandle<TaskType>>(
this->shared_from_this(), i, task_sizes[i]);
}
}
template <typename TaskType>
std::unique_ptr<TaskType> BatchInputTask<TaskType>::GetSplitTask(int split_id) {
absl::call_once(once_,
[this]() { split_status_ = SplitBatches(&task_splits_); });
if (!split_status_.ok()) {
LOG_EVERY_N_SEC(WARNING, 60 /* seconds */)
<< "Split task with error: " << split_status_ << " split metadata is "
<< input_split_metadata_.DebugString();
return nullptr;
}
if (split_id >= 0 && split_id < task_splits_.size()) {
return std::move(task_splits_[split_id]);
}
return nullptr;
}
template <typename TaskType>
absl::Status BatchInputTask<TaskType>::SplitBatches(
std::vector<std::unique_ptr<TaskType>>* output_tasks) {
return split_func_(&input_task_, open_batch_remaining_slot_,
batch_size_limit_, output_tasks);
}
} // namespace internal
} // namespace serving
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_INPUT_TASK_H_
tensorflow/core/kernels/batching_util/batch_input_task_test.cc
+209
View File
@@ -0,0 +1,209 @@
/* 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/kernels/batching_util/batch_input_task.h"
#include <functional>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include "absl/status/status.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/framework/allocator.h"
#include "tensorflow/core/framework/device_factory.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/node_properties.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/kernels/batching_util/batch_resource_base.h"
#include "tensorflow/core/kernels/batching_util/input_split_metadata.h"
#include "tensorflow/core/kernels/batching_util/threadsafe_status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
#include "tensorflow/core/public/session_options.h"
#include "tensorflow/core/public/version.h"
namespace tensorflow {
namespace serving {
namespace internal {
template <typename TaskType>
class BatchInputTaskHandleTestAccess {
public:
explicit BatchInputTaskHandleTestAccess(
BatchInputTaskHandle<TaskType>* handle)
: handle_(handle) {}
int split_id() const { return handle_->split_id(); }
private:
BatchInputTaskHandle<TaskType>* const handle_;
};
namespace {
using TensorMatrix = std::vector<std::vector<Tensor>>;
using SplitFunc = std::function<absl::Status(
std::unique_ptr<BatchResourceBase::BatchTask>* input_task,
int first_output_task_size, int input_batch_size_limit,
std::vector<std::unique_ptr<BatchResourceBase::BatchTask>>* output_tasks)>;
// Creates a tensor with the specified dtype, shape, and value.
template <typename T>
static Tensor CreateTensor(const TensorShape& input_shape,
gtl::ArraySlice<T> input_data) {
Tensor tensor(DataTypeToEnum<T>::value, input_shape);
test::FillValues<T>(&tensor, input_data);
return tensor;
}
NodeDef CreateBatchKernelNodeDef() {
NodeDef batch_kernel_node_def;
NodeDefBuilder batch_function_builder("BatchTPUInput", "BatchFunction");
batch_function_builder.Attr("max_batch_size", 128);
batch_function_builder.Attr("num_batch_threads", 8);
batch_function_builder.Attr("allowed_batch_sizes", {2, 4, 8});
batch_function_builder.Attr("batch_timeout_micros", 1000);
batch_function_builder.Attr("max_enqueued_batches", 100);
batch_function_builder.Attr("enable_large_batch_splitting", true);
std::vector<DataType> input_dtypes({DataType::DT_INT64, DataType::DT_INT64});
std::vector<NodeDefBuilder::NodeOut> inputs;
inputs.resize(2);
inputs[0] = NodeDefBuilder::NodeOut({"n1", 0, DataType::DT_INT64});
inputs[1] = NodeDefBuilder::NodeOut({"n2", 1, DataType::DT_INT64});
batch_function_builder.Attr("Tin", input_dtypes);
batch_function_builder.Input(inputs);
batch_function_builder.Attr("Tcaptured",
std::vector<DataType>{DataType::DT_INT64});
batch_function_builder.Input(std::vector<NodeDefBuilder::NodeOut>{
NodeDefBuilder::NodeOut({"n3", 1, DataType::DT_INT64})});
batch_function_builder.Attr("Tout",
std::vector<DataType>(4, DataType::DT_INT64));
NameAttrList f;
f.set_name("func_to_batch");
batch_function_builder.Attr("f", f);
TF_CHECK_OK(batch_function_builder.Finalize(&batch_kernel_node_def));
return batch_kernel_node_def;
}
class BatchInputTaskTest : public ::testing::Test {
protected:
BatchInputTaskTest() {
device_ = DeviceFactory::NewDevice("CPU", SessionOptions{},
"/job:a/replica:0/task:0");
absl::Status op_kernel_creation_status;
batch_kernel_ = CreateOpKernel(
DEVICE_CPU, device_.get(), device_->GetAllocator(AllocatorAttributes{}),
CreateBatchKernelNodeDef(), TF_GRAPH_DEF_VERSION,
&op_kernel_creation_status);
TF_CHECK_OK(op_kernel_creation_status);
EXPECT_NE(batch_kernel_, nullptr);
op_kernel_context_params_.device = device_.get();
op_kernel_context_params_.op_kernel = batch_kernel_.get();
op_kernel_context_ = std::make_unique<OpKernelContext>(
&op_kernel_context_params_, 4 /* num outputs */);
}
OpKernelContext* op_kernel_context() const {
return op_kernel_context_.get();
}
private:
std::unique_ptr<Device> device_;
std::unique_ptr<OpKernel> batch_kernel_;
OpKernelContext::Params op_kernel_context_params_;
std::unique_ptr<OpKernelContext> op_kernel_context_;
};
TEST_F(BatchInputTaskTest, BatchInputToSplitTasks) {
bool batch_task_done_callback_executed = false;
auto status = std::make_shared<ThreadSafeStatus>();
auto done_callback = [&batch_task_done_callback_executed]() {
batch_task_done_callback_executed = true;
};
auto batch_task = std::make_unique<BatchResourceBase::BatchTask>(
std::move(done_callback), std::move(status));
batch_task->inputs.push_back(CreateTensor<int64_t>(
TensorShape({5, 2, 1}), {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}));
batch_task->inputs.push_back(CreateTensor<int64_t>(
TensorShape({5, 1, 2}), {11, 12, 13, 14, 15, 16, 17, 18, 19, 20}));
batch_task->captured_inputs.push_back(
CreateTensor<int64_t>(TensorShape{1}, {0}));
batch_task->context = op_kernel_context();
batch_task->output = std::make_shared<TensorMatrix>();
auto batch_input_task =
std::make_shared<BatchInputTask<BatchResourceBase::BatchTask>>(
std::move(batch_task), /*open_batch_remaining_slot=*/1,
/*batch_size_limit=*/3, BatchResourceBase::SplitInputTask);
std::vector<
std::unique_ptr<BatchInputTaskHandle<BatchResourceBase::BatchTask>>>
output_tasks;
batch_input_task->ToTaskHandles(&output_tasks);
// Output tasks haven't invoked `done_callback`, so
// `batch_task->done_callback` hasn't run yet.
ASSERT_FALSE(batch_task_done_callback_executed);
const std::vector<int> expected_task_sizes{1, 3, 1};
// Call `done_callback` for each output task, so `batch_task->done_callback`
// can be executed and batch_task_done_callback_executed will be updated.
for (int i = 0; i < output_tasks.size(); i++) {
// When emitting output tasks, split_id starts from zero and increases by
// one.
EXPECT_EQ(
internal::BatchInputTaskHandleTestAccess<BatchResourceBase::BatchTask>(
output_tasks[i].get())
.split_id(),
i);
auto batch_task = output_tasks[i]->GetSplitTask();
ASSERT_NE(batch_task, nullptr);
EXPECT_EQ(batch_task->size(), expected_task_sizes[i]);
batch_task->FinishTask(absl::OkStatus());
// `GetSplitTask` returns nullptr from the 2nd call and on.
EXPECT_EQ(output_tasks[i]->GetSplitTask(), nullptr);
}
// Each output task completed, so `batch_task->done_callback` ran.
ASSERT_TRUE(batch_task_done_callback_executed);
}
} // namespace
} // namespace internal
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/batch_resource_base.cc
+1609
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/kernels/batching_util/batch_resource_base.h
+456
View File
@@ -0,0 +1,456 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_RESOURCE_BASE_H_
#define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_RESOURCE_BASE_H_
#include <cstdint>
#include <functional>
#include <map>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include "absl/strings/str_join.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/blocking_counter.h"
#include "absl/time/time.h"
#include "xla/tsl/platform/criticality.h"
#include "tensorflow/core/common_runtime/cost_measurement_registry.h"
#include "tensorflow/core/common_runtime/request_cost.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/kernels/batching_util/adaptive_shared_batch_scheduler.h"
#include "tensorflow/core/kernels/batching_util/batch_scheduler.h"
#include "tensorflow/core/kernels/batching_util/batch_scheduler_utils.h"
#include "tensorflow/core/kernels/batching_util/shared_batch_scheduler.h"
#include "tensorflow/core/kernels/batching_util/threadsafe_status.h"
#include "tensorflow/core/platform/context.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/protobuf/config.pb.h"
namespace tensorflow {
namespace serving {
// Options used to create a batch resource.
struct BatchResourceOptions {
int32_t num_batch_threads;
int32_t max_batch_size;
int32_t batch_timeout_micros;
int32_t max_enqueued_batches;
std::vector<int32_t> allowed_batch_sizes;
std::string batch_padding_policy{kPadUpPolicy};
int32_t low_priority_max_batch_size;
int32_t low_priority_batch_timeout_micros;
int32_t low_priority_max_enqueued_batches;
std::vector<int32_t> low_priority_allowed_batch_sizes;
MixedPriorityBatchingPolicy mixed_priority_batching_policy;
bool enable_priority_aware_batch_scheduler;
bool enable_priority_aware_batch_scheduler_resplit;
bool enable_batching_task_lazy_cancellation;
int32_t num_warmup_batch_threads;
};
// Base class for resource that encapsulating the state and logic for batching
// tensors.
class BatchResourceBase : public ResourceBase {
public:
// Given a BatchTask (from one op invocation) with 'num_outputs'== M and
// split into N sub tasks, TensorMatrix is a N X M matrix.
// Namely, TensorMatrix[i][j] indicates the i-th split tensor of j-th output;
// concatenating tensors along the 2nd dimension gives a output tensor.
typedef std::vector<std::vector<Tensor>> TensorMatrix;
// One task to be batched, corresponds to a `slice` of input from one batch-op
// invocation.
//
// Given input from one batch-op invocation or a subtask that is split from
// the input from an invocation, a `slice` is:
// 1) Split each Tensor in `BatchTask::inputs` along the 0th dimension.
// 2) 'split_index' is calculated along the 0-th dimension.
// A subtask may itself be re-split, allowing arbitrary nesting.
//
// Note input from one batch-op invocation is valid and considered a
// specialized `slice`.
struct BatchTask : public tensorflow::serving::BatchTask {
BatchTask() : criticality_val(tsl::criticality::GetCriticality()) {};
BatchTask(AsyncOpKernel::DoneCallback done_callback,
std::shared_ptr<ThreadSafeStatus> status)
: status(status),
done_callback(std::move(done_callback)),
criticality_val(tsl::criticality::GetCriticality()) {}
// A unique ID to identify this invocation of Batch.
int64_t guid;
Context propagated_context;
std::vector<Tensor> inputs;
std::vector<Tensor> captured_inputs;
OpKernelContext* context;
// The index of this split, along the 0-th dimension of input from op
// invocation.
int split_index = 0;
// Two-dimensional tensor matrix where output of all splits of task is
// collected. The splits can be further split; BatchTask supports
// arbitrary levels of splits.
// Ownership shared by:
// 1) each split of task (to fill one row in this matrix) and
// 2) callback that runs to merge output of individual splits for an op
// invocation, after all splits complete.
std::shared_ptr<TensorMatrix> output;
bool is_partial = false;
uint64 start_time;
// Absolute RPC deadline. When set, the task is considered expired if
// absl::Now() > rpc_deadline. Defaults to nullopt (no enforcement).
std::optional<absl::Time> rpc_deadline;
// Callback: returns true if client actively cancelled the RPC.
std::function<bool()> is_rpc_cancelled;
size_t size() const override { return inputs[0].shape().dim_size(0); }
bool is_subtask() const override { return is_partial; }
bool IsDeadlineExceeded(absl::Time now) const override;
bool IsCancelled() const override;
// Create a split task from this one. The caller needs to setup the inputs
// of the new task
std::unique_ptr<BatchTask> CreateSplitTask(
int split_index, AsyncOpKernel::DoneCallback done_callback);
// RequestCost is for collecting the cost and must outlive the batching
// processing.
//
// For example, to collect cost in rpc processing, `request_cost` is owned
// by rpc handler and points to the RequestCost of an rpc which provides
// the inputs to this BatchTask.
//
// After the batch processing, the request cost will be incremented with
// this task's processing costs.
RequestCost* request_cost = nullptr;
// Returns the criticality associated with the task.
tsl::criticality::Criticality criticality() const override {
return criticality_val;
};
// If nonzero, make a batch of this size entirely out of padding. This
// batch is processed, but is not propagated to the kernel outputs.
int forced_warmup_batch_size = 0;
// If true, the task is a warmup task.
bool is_warmup_task = false;
bool is_warmup() const override { return is_warmup_task; }
// 'status' records error (could be from any split) if at least one split
// returns error, OK otherwise.
// Ownership is shared by individual splits and callback.
std::shared_ptr<ThreadSafeStatus> status;
// Note the done callback MUST NOT (although unlikely in practice) attempt
// to acquire the queue's lock — e.g., by calling methods like Schedule() on
// the same queue — as this may lead to deadlock.
void set_done_callback(AsyncOpKernel::DoneCallback callback) {
done_callback = std::move(callback);
}
protected:
void FinishTaskImpl(const absl::Status& status) override {
WithContext wc(this->propagated_context);
if (!status.ok()) {
this->status->Update(status);
}
if (this->done_callback) {
this->done_callback();
// Clear the callback to avoid double deletion.
this->done_callback = nullptr;
}
}
virtual std::unique_ptr<BatchTask> CreateDerivedTask() {
#if defined(PLATFORM_GOOGLE)
// ScopedCriticality is needed to ensure that the criticality is set
// correctly for the derived task.
tsl::criticality::ScopedCriticality scoped_criticality(
this->criticality());
#endif
return std::make_unique<BatchTask>();
}
private:
AsyncOpKernel::DoneCallback done_callback;
// Criticality associated with the task.
::tsl::criticality::Criticality criticality_val;
};
// Appending a T suffix to make the type alias different to those in
// tensorflow::serving namespace, because some versions of compiler complain
// about changing meaning of the symbols.
using BatcherT = SharedBatchScheduler<BatchResourceBase::BatchTask>;
using AdaptiveBatcherT =
AdaptiveSharedBatchScheduler<BatchResourceBase::BatchTask>;
using BatcherQueueT = BatchScheduler<BatchResourceBase::BatchTask>;
using BatchT = Batch<BatchResourceBase::BatchTask>;
BatchResourceBase(bool has_process_batch_function,
std::shared_ptr<BatcherT> batcher,
const BatcherT::QueueOptions& batcher_queue_options,
std::vector<int32> allowed_batch_sizes)
: has_process_batch_function_(has_process_batch_function),
batcher_(std::move(batcher)),
batcher_queue_options_(batcher_queue_options),
allowed_batch_sizes_(std::move(allowed_batch_sizes)),
allowed_batch_sizes_str_(absl::StrJoin(allowed_batch_sizes_, ",")) {}
BatchResourceBase(bool has_process_batch_function,
std::shared_ptr<AdaptiveBatcherT> batcher,
const AdaptiveBatcherT::QueueOptions& batcher_queue_options,
std::vector<int32> allowed_batch_sizes)
: has_process_batch_function_(has_process_batch_function),
adaptive_batcher_(std::move(batcher)),
adaptive_batcher_queue_options_(batcher_queue_options),
allowed_batch_sizes_(std::move(allowed_batch_sizes)),
allowed_batch_sizes_str_(absl::StrJoin(allowed_batch_sizes_, ",")) {}
void set_session_metadata(tensorflow::SessionMetadata session_metadata) {
session_metadata_ = std::move(session_metadata);
}
const SessionMetadata& session_metadata() const { return session_metadata_; }
using CreateBatchTaskFn =
std::function<StatusOr<std::unique_ptr<BatchTask>>()>;
// Like `RegisterInput`, but extra "dummy" batches are processed for each
// batch size. Only the real request's outputs are propagated to the caller.
Status RegisterWarmupInputs(int64_t guid, OpKernelContext* context,
const string& batcher_queue_name,
const CreateBatchTaskFn& create_batch_task_fn,
AsyncOpKernel::DoneCallback done);
// Ingests data from one invocation of the batch op. The data is enqueued to
// be combined with others into a batch, asynchronously.
// `CreateBatchTaskFn` should be used to instantiate fields added to a
// child class of `BatchTask` by the caller.
Status RegisterInput(int64_t guid, OpKernelContext* context,
const string& batcher_queue_name,
const CreateBatchTaskFn& create_batch_task_fn,
AsyncOpKernel::DoneCallback done_callback,
int forced_warmup_batch_size = 0);
static BatcherT::QueueOptions GetBatcherQueueOptions(
int32_t num_batch_threads, int32_t max_batch_size,
int32_t batch_timeout_micros, int32_t max_enqueued_batches,
const std::vector<int32>& allowed_batch_sizes,
bool enable_large_batch_splitting, bool disable_padding);
static BatcherT::QueueOptions GetBatcherQueueOptions(
int32_t num_batch_threads, int32_t max_batch_size,
int32_t batch_timeout_micros, int32_t max_enqueued_batches,
const std::vector<int32>& allowed_batch_sizes,
bool enable_large_batch_splitting, bool disable_padding,
absl::string_view batch_padding_policy,
int32_t low_priority_max_batch_size,
int32_t low_priority_batch_timeout_micros,
int32_t low_priority_max_enqueued_batches,
const std::vector<int32>& low_priority_allowed_batch_sizes,
MixedPriorityBatchingPolicy mixed_priority_batching_policy,
bool enable_priority_aware_batch_scheduler,
bool enable_priority_aware_batch_scheduler_resplit,
bool enable_batching_task_lazy_cancellation);
static AdaptiveBatcherT::QueueOptions GetAdaptiveBatcherQueueOptions(
int32_t max_batch_size, int32_t batch_timeout_micros,
int32_t max_enqueued_batches, bool enable_large_batch_splitting,
const std::vector<int32>& allowed_batch_sizes, bool disable_padding);
// Split 'input' of 'input_task_ptr' along 0th dimension, into a list of
// 'output_tasks'.
// Task sizes are determined by
// 1) open_batch_remaining_slot
// 2) max_batch_size
// 3) size-of-input-task
// in a way that
// 1) Task sizes add up to `size-of-input-task`.
// 2) Task sizes from left to right are like
// [open_batch_remaining_slot, max_batch_size, max_batch_size, ...,
// `size-of-input-task` - `sum-of-previous-elements`].
//
// REQUIRES:
// Caller should make sure size-of-input-task is greater than
// open_batch_remaining_slot.
static Status SplitInputTask(
std::unique_ptr<BatchTask>* input_task_ptr, int open_batch_remaining_slot,
int max_batch_size,
std::vector<std::unique_ptr<BatchTask>>* output_tasks);
// Splits the batch costs to each task.
//
// Inputs:
// 1) batch_cost_measurements, which provides the total cost of each type;
// 2) processed_size, it's the batch size plus the padding amount;
// 3) batch, provides the batch size and input sizes.
//
// Outputs:
// The request_cost in each batch task will be updated.
// - This function will use two approaches to split the batch cost (if it's
// non-zero), thus two costs will be output.
// 1) smeared cost: batch cost is split proportionally to each task's size,
// and paddings do not share any cost;
// 2) non-smeared cost: batch cost is split proportionally to each task or
// padding's size. Here padding's cost is not assigned to any tasks.
// - This function will also record the metrics of this batch in each task,
// including:
// 1) the batch size;
// 2) the input size from this task;
// 3) the padding amount.
static void SplitBatchCostsAndRecordMetrics(
const std::string& model_name, const std::string& op_name,
const std::vector<std::unique_ptr<CostMeasurement>>&
batch_cost_measurements,
int64_t processed_size, BatchT& batch);
// Records information about the delay between a task being registered and
// that task being scheduled into a batch.
static void RecordBatchDelayMetrics(
const BatchResourceBase::BatchT& batch, const std::string& model_name,
const std::string& op_name, int64_t processed_size,
absl::Time batch_schedule_time,
std::optional<absl::Duration> batch_timeout);
private:
// Implementation of calling the process batch function.
virtual void ProcessFuncBatchImpl(
const BatchResourceBase::BatchTask& last_task,
absl::Span<const Tensor> inputs, std::vector<Tensor>* combined_outputs,
std::function<void(const Status&)> done) const = 0;
// Validates that it's legal to combine the tasks in 'batch' into a batch.
// Assumes the batch is non-empty.
static Status ValidateBatch(const BatchT& batch);
// Returns a boolean indicating whether a batch is formed from low priority
// tasks only or not.
bool IsLowPriorityBatch(const BatchT& batch) const;
// Returns the smallest entry in 'allowed_batch_sizes_' that is greater than
// or equal to 'batch_size'. If 'allowed_batch_sizes_' is empty, simply
// returns 'batch_size'.
int RoundToLowestAllowedBatchSize(int batch_size,
bool is_low_priority_batch = false) const;
// Helper function to propagate the status to the task's context and call the
// done callback on the task.
void CleanUpFunctionHelper(BatchTask& task, const Status& status) const;
// Concatenates the input tensors of the tasks from the batch and the
// unbatched task vector. When padding is enabled in the batcher queue, they
// are padded with garbage value up to the nearest allowed batch size.
Status ConcatInputTensors(
const BatchT& batch,
const std::vector<std::unique_ptr<BatchTask>>& unbatched_tasks,
OpKernelContext* context,
std::vector<Tensor>* concatenated_tensors) const;
Status SplitOutputTensors(
const std::vector<Tensor>& combined_outputs, BatchT* batch,
std::vector<std::unique_ptr<BatchTask>>& unbatched_tasks) const;
void ProcessFuncBatch(
std::unique_ptr<BatchT> batch,
std::vector<std::unique_ptr<BatchTask>> unbatched_tasks = {}) const;
// Processes a batch of one or more BatchTask entries.
void ProcessBatch(std::unique_ptr<BatchT> batch) const;
// Callback function that wraps the Process*Batch functions above. The caller
// of the callback must guarantee that the unique pointers passed as argument
// are not null.
void ProcessBatchCallBack(
std::unique_ptr<Batch<BatchTask>> batch,
std::vector<std::unique_ptr<BatchTask>> unbatched_tasks);
// Emits an index tensor, which the Unbatch op will use to un-concatenate
// the tensor and attribute the pieces to the right batch keys. The index
// tensor contains, for each input: [batch_key, start_offset, end_offset]
// where start_offset and end_offset represent the range of entries in the
// concatenated tensors that belong to that input.
//
// Emits the result to the output at 'output_index' using 'context'.
static Status EmitIndexTensor(OpKernelContext* context, const BatchT& batch,
int output_index);
// Looks up the batcher queue for 'queue_name'. If it didn't previously exist,
// creates it.
//
// The model_name and op_name are the names of the current model and
// operation, respectively.
Status LookupOrCreateBatcherQueue(const string& queue_name,
const string& model_name,
const string& op_name,
BatcherQueueT** queue);
// Returns the batch timeout for the configured scheduler, or nullopt if the
// scheduler does not have such a parameter.
std::optional<absl::Duration> GetBatchTimeout() const;
SessionMetadata session_metadata_;
absl::Mutex outstanding_batch_mu_;
int num_outstanding_batched_items_ TF_GUARDED_BY(outstanding_batch_mu_) = 0;
// True if user specified a batch processing function for this resource.
const bool has_process_batch_function_;
// A batch scheduler, and options for creating queues.
std::shared_ptr<BatcherT> batcher_;
BatcherT::QueueOptions batcher_queue_options_;
// A batch scheduler, and options for creating queues.
std::shared_ptr<AdaptiveBatcherT> adaptive_batcher_;
AdaptiveBatcherT::QueueOptions adaptive_batcher_queue_options_;
// A collection of batcher queues, keyed on queue name.
// TODO(olston): Garbage-collect unused queues (perhaps simply remove empty
// ones (with a time delay?); it's okay if they get recreated later).
mutable mutex batcher_queues_mu_;
std::map<string, std::unique_ptr<BatcherQueueT>> batcher_queues_
TF_GUARDED_BY(batcher_queues_mu_);
std::vector<int32> allowed_batch_sizes_;
// A concatenated string of <allowed_batch_sizes_>, separated by ",". This is
// used to record batching parameter.
string allowed_batch_sizes_str_;
};
} // namespace serving
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_RESOURCE_BASE_H_
tensorflow/core/kernels/batching_util/batch_resource_base_test.cc
+1888
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/kernels/batching_util/batch_scheduler.cc
+62
View File
@@ -0,0 +1,62 @@
/* Copyright 2024 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/batching_util/batch_scheduler.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
namespace tensorflow {
namespace serving {
absl::StatusOr<MixedPriorityBatchingPolicy> GetMixedPriorityBatchingPolicy(
absl::string_view attr_value) {
if (attr_value == kLowPriorityPaddingWithMaxBatchSizeAttrValue) {
return MixedPriorityBatchingPolicy::kLowPriorityPaddingWithMaxBatchSize;
} else if (attr_value ==
kLowPriorityPaddingWithNextAllowedBatchSizeAttrValue) {
return MixedPriorityBatchingPolicy::
kLowPriorityPaddingWithNextAllowedBatchSize;
} else if (attr_value == kPriorityIsolationAttrValue) {
return MixedPriorityBatchingPolicy::kPriorityIsolation;
} else if (attr_value == kPriorityMergeAttrValue) {
return MixedPriorityBatchingPolicy::kPriorityMerge;
}
return absl::InvalidArgumentError(absl::StrFormat(
"Unknown mixed priority batching policy: %s", attr_value));
}
absl::StatusOr<absl::string_view> GetMixedPriorityBatchingPolicyString(
MixedPriorityBatchingPolicy policy) {
switch (policy) {
case MixedPriorityBatchingPolicy::kLowPriorityPaddingWithMaxBatchSize:
return kLowPriorityPaddingWithMaxBatchSizeAttrValue;
case MixedPriorityBatchingPolicy::
kLowPriorityPaddingWithNextAllowedBatchSize:
return kLowPriorityPaddingWithNextAllowedBatchSizeAttrValue;
case MixedPriorityBatchingPolicy::kPriorityIsolation:
return kPriorityIsolationAttrValue;
case MixedPriorityBatchingPolicy::kPriorityMerge:
return kPriorityMergeAttrValue;
default:
return absl::InvalidArgumentError(absl::StrFormat(
"Unknown mixed priority batching policy: %d", policy));
}
}
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/batch_scheduler.h
+632
View File
@@ -0,0 +1,632 @@
/* Copyright 2016 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.
==============================================================================*/
// Abstractions for processing small tasks in a batched fashion, to reduce
// processing times and costs that can be amortized across multiple tasks.
//
// The core class is BatchScheduler, which groups tasks into batches.
//
// BatchScheduler encapsulates logic for aggregating multiple tasks into a
// batch, and kicking off processing of a batch on a thread pool it manages.
//
// This file defines an abstract BatchScheduler class.
#ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_SCHEDULER_H_
#define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_SCHEDULER_H_
#include <stddef.h>
#include <sys/types.h>
#include <algorithm>
#include <atomic>
#include <cstddef>
#include <deque>
#include <iterator>
#include <memory>
#include <optional>
#include <utility>
#include <vector>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/notification.h"
#include "absl/time/time.h"
#include "xla/tsl/platform/criticality.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/platform/types.h"
#include "tsl/profiler/lib/traceme.h"
namespace tensorflow {
namespace serving {
const absl::string_view kLowPriorityPaddingWithMaxBatchSizeAttrValue =
"low_priority_padding_with_max_batch_size";
const absl::string_view kLowPriorityPaddingWithNextAllowedBatchSizeAttrValue =
"low_priority_padding_with_next_allowed_batch_size";
const absl::string_view kPriorityIsolationAttrValue = "priority_isolation";
const absl::string_view kPriorityMergeAttrValue = "priority_merge";
enum class MixedPriorityBatchingPolicy {
kLowPriorityPaddingWithMaxBatchSize,
kLowPriorityPaddingWithNextAllowedBatchSize,
kPriorityIsolation,
kPriorityMerge,
};
absl::StatusOr<MixedPriorityBatchingPolicy> GetMixedPriorityBatchingPolicy(
absl::string_view attr_value);
absl::StatusOr<absl::string_view> GetMixedPriorityBatchingPolicyString(
MixedPriorityBatchingPolicy policy);
// The abstract superclass for a unit of work to be done as part of a batch.
//
// An implementing subclass typically contains (or points to):
// (a) input data;
// (b) a thread-safe completion signal (e.g. a Notification);
// (c) a place to store the outcome (success, or some error), upon completion;
// (d) a place to store the output data, upon success.
//
// Items (b), (c) and (d) are typically non-owned pointers to data homed
// elsewhere, because a task's ownership gets transferred to a BatchScheduler
// (see below) and it may be deleted as soon as it is done executing.
class BatchTask {
public:
virtual ~BatchTask() = default;
// Returns the size of the task, in terms of how much it contributes to the
// size of a batch. (A batch's size is the sum of its task sizes.)
virtual size_t size() const = 0;
// Returns the criticality of associated with the task. It defaults to
// kCritical.
virtual tsl::criticality::Criticality criticality() const {
return tsl::criticality::Criticality::kCritical;
}
// If true, the task is a warmup task.
virtual bool is_warmup() const { return false; }
virtual bool is_subtask() const { return false; }
// Returns true if the task's deadline has expired.
// `now` is passed by the caller to amortize absl::Now() across iterations.
virtual bool IsDeadlineExceeded(absl::Time now) const { return false; }
// Returns true if the RPC has been cancelled by the client.
virtual bool IsCancelled() const { return false; }
// Called when the task is finished, either successfully or with an error.
//
// BatchScheduler guarantees that this method is invoked exactly once for each
// BatchTask. This frees the task implementation from worrying about thread
// safety or idempotency of this method.
void FinishTask(const absl::Status& status) { FinishTaskImpl(status); }
protected:
// TODO: Make this method pure-virtual once all subclasses implement it.
virtual void FinishTaskImpl(const absl::Status& status) {
// Default implementation does nothing. Subclasses should override.
}
};
// A thread-safe collection of BatchTasks. Tasks can be either added or removed
// from the TaskQueue. It is mainly used to hold the registered tasks without
// forming batches, so that the batches can be formed more flexibly right before
// they get scheduled for execution.
//
// Type parameter TaskType must be a subclass of BatchTask.
template <typename TaskType>
class TaskQueue {
public:
TaskQueue() = default;
struct TaskWrapper {
std::unique_ptr<TaskType> task;
uint64_t start_time_micros;
TaskWrapper(std::unique_ptr<TaskType> task, uint64_t start_time_micros)
: task(std::move(task)), start_time_micros(start_time_micros) {}
};
// Appends a task to the end of the queue with the given start time.
void AddTask(std::unique_ptr<TaskType> task, uint64_t start_time_micros);
// Adds a task to the front of the queue with the given start time.
void PrependTask(std::unique_ptr<TaskType> task, uint64_t start_time_micros);
// Removes a task from the front of the queue, i.e., the oldest task in the
// queue.
std::unique_ptr<TaskType> RemoveTask();
// Removes tasks from the front of the queue as many as possible as long as
// the sum of sizes of the removed tasks don't exceed the 'size' given as the
// argument.
std::vector<std::unique_ptr<TaskType>> RemoveTask(int size);
// Returns the start time of the earliest task in the queue. If the queue is
// empty, return the null value.
std::optional<uint64_t> EarliestTaskStartTime() const;
// Returns true iff the queue contains 0 tasks.
bool empty() const;
// Returns the number of tasks in the queue.
int num_tasks() const;
// Returns the sum of the task sizes.
int size() const;
private:
mutable mutex mu_;
// Tasks in the queue.
std::deque<TaskWrapper> tasks_ TF_GUARDED_BY(mu_);
// The sum of the sizes of the tasks in 'tasks_'.
int size_ TF_GUARDED_BY(mu_) = 0;
// Whether the queue is empty.
std::atomic<bool> empty_ TF_GUARDED_BY(mu_){true};
// The copy constructor and the assign op are deleted.
TaskQueue(const TaskQueue&) = delete;
void operator=(const TaskQueue&) = delete;
};
template <typename TaskType>
void TaskQueue<TaskType>::AddTask(std::unique_ptr<TaskType> task,
uint64_t start_time_micros) {
{
mutex_lock l(mu_);
size_ += task->size();
tasks_.emplace_back(std::move(task), start_time_micros);
empty_.store(false);
}
}
template <typename TaskType>
void TaskQueue<TaskType>::PrependTask(std::unique_ptr<TaskType> task,
uint64_t start_time_micros) {
{
mutex_lock l(mu_);
size_ += task->size();
tasks_.emplace_front(std::move(task), start_time_micros);
empty_.store(false);
}
}
template <typename TaskType>
std::unique_ptr<TaskType> TaskQueue<TaskType>::RemoveTask() {
{
mutex_lock l(mu_);
if (tasks_.empty()) {
return nullptr;
}
std::unique_ptr<TaskType> task = std::move(tasks_.front().task);
size_ -= task->size();
tasks_.pop_front();
if (tasks_.empty()) {
empty_.store(true);
}
return task;
}
}
template <typename TaskType>
std::vector<std::unique_ptr<TaskType>> TaskQueue<TaskType>::RemoveTask(
int size) {
{
mutex_lock l(mu_);
if (tasks_.empty()) {
return {};
}
int size_lower_bound = size_ - size;
std::vector<std::unique_ptr<TaskType>> remove_tasks;
while (!tasks_.empty() &&
size_ - static_cast<int>(tasks_.front().task->size()) >=
size_lower_bound) {
size_ -= static_cast<int>(tasks_.front().task->size());
remove_tasks.push_back(std::move(tasks_.front().task));
tasks_.pop_front();
if (tasks_.empty()) {
empty_.store(true);
}
}
return remove_tasks;
}
}
template <typename TaskType>
bool TaskQueue<TaskType>::empty() const {
{
mutex_lock l(mu_);
return empty_.load();
}
}
template <typename TaskType>
std::optional<uint64_t> TaskQueue<TaskType>::EarliestTaskStartTime() const {
{
mutex_lock l(mu_);
if (tasks_.empty()) {
return std::nullopt;
}
return tasks_.front().start_time_micros;
}
}
template <typename TaskType>
int TaskQueue<TaskType>::num_tasks() const {
{
mutex_lock l(mu_);
return tasks_.size();
}
}
template <typename TaskType>
int TaskQueue<TaskType>::size() const {
{
mutex_lock l(mu_);
return size_;
}
}
// A thread-safe collection of BatchTasks, to be executed together in some
// fashion.
//
// At a given time, a batch is either "open" or "closed": an open batch can
// accept new tasks; a closed one cannot. A batch is monotonic: initially it is
// open and tasks can be added to it; then it is closed and its set of tasks
// remains fixed for the remainder of its life. A closed batch cannot be re-
// opened.
//
// Type parameter TaskType must be a subclass of BatchTask.
template <typename TaskType>
class Batch {
public:
Batch();
explicit Batch(uint64_t traceme_context_id);
virtual ~Batch(); // Blocks until the batch is closed.
// Appends 'task' to the batch. After calling AddTask(), the newly-added task
// can be accessed via task(num_tasks()-1) or mutable_task(num_tasks()-1).
// Dies if the batch is closed.
void AddTask(std::unique_ptr<TaskType> task, uint64_t start_time_micros = 0);
// Removes the most recently added task. Returns nullptr if the batch is
// empty.
std::unique_ptr<TaskType> RemoveTask();
// Caller takes ownership of returned tasks.
// Must be called after a batch is closed.
std::vector<std::unique_ptr<TaskType>> RemoveAllTasks();
// Returns the number of tasks in the batch.
int num_tasks() const;
// Returns true iff the batch contains 0 tasks.
bool empty() const;
// Returns a reference to the ith task (in terms of insertion order).
const TaskType& task(int i) const;
// Returns a pointer to the ith task (in terms of insertion order).
//
// Caller doesn't take ownership.
TaskType* mutable_task(int i);
// Returns the sum of the task sizes.
size_t size() const;
// Returns true iff the batch is currently closed.
bool IsClosed() const;
// Blocks until the batch is closed.
void WaitUntilClosed() const;
// Marks the batch as closed. Dies if called more than once.
void Close();
// Returns the TraceMe context id of this batch.
uint64_t traceme_context_id() const;
// Attempts to trim this batch to a new, smaller size (not to be confused with
// the number of tasks in the batch). On success, the trimmed tasks go into
// 'out_trimmed_tasks' in the same order the tasks were in this batch.
//
// The method might not succeed if it needs to split a large task to hit the
// correct size.
void TryTrimToNewSize(
int new_size, std::vector<std::unique_ptr<TaskType>>& out_trimmed_tasks);
// Returns the start time of the earliest task in the queue. If the queue is
// empty, return the null value.
std::optional<uint64_t> EarliestTaskStartTime() const;
private:
mutable mutex mu_;
// The tasks in the batch.
std::vector<std::unique_ptr<TaskType>> tasks_ TF_GUARDED_BY(mu_);
// The sum of the sizes of the tasks in 'tasks_'.
size_t size_ TF_GUARDED_BY(mu_) = 0;
std::atomic<bool> empty_ TF_GUARDED_BY(mu_){true};
// Whether the batch has been closed.
absl::Notification closed_;
// The TracMe context id.
const uint64_t traceme_context_id_;
// The minimum start time of all tasks in the batch.
// If the batch is empty, the value is undefined.
uint64_t earliest_task_start_time_micros_ TF_GUARDED_BY(mu_);
Batch(const Batch&) = delete;
void operator=(const Batch&) = delete;
};
// An abstract batch scheduler class. Collects individual tasks into batches,
// and processes each batch on a pool of "batch threads" that it manages. The
// actual logic for processing a batch is accomplished via a callback.
//
// Type parameter TaskType must be a subclass of BatchTask.
template <typename TaskType>
class BatchScheduler {
public:
virtual ~BatchScheduler() = default;
// Submits a task to be processed as part of a batch.
//
// Ownership of '*task' is transferred to the callee iff the method returns
// Status::OK. In that case, '*task' is left as nullptr. Otherwise, '*task' is
// left as-is.
//
// If no batch processing capacity is available to process this task at the
// present time, and any task queue maintained by the implementing subclass is
// full, this method returns an UNAVAILABLE error code. The client may retry
// later.
//
// Other problems, such as the task size being larger than the maximum batch
// size, yield other, permanent error types.
//
// In all cases, this method returns "quickly" without blocking for any
// substantial amount of time. If the method returns Status::OK, the task is
// processed asynchronously, and any errors that occur during the processing
// of the batch that includes the task can be reported to 'task'.
virtual absl::Status Schedule(std::unique_ptr<TaskType>* task) = 0;
// Returns the number of tasks that have been scheduled (i.e. accepted by
// Schedule()), but have yet to be handed to a thread for execution as part of
// a batch. Note that this returns the number of tasks, not the aggregate task
// size (so if there is one task of size 3 and one task of size 5, this method
// returns 2 rather than 8).
virtual size_t NumEnqueuedTasks() const = 0;
// Returns a guaranteed number of size 1 tasks that can be Schedule()d without
// getting an UNAVAILABLE error. In a typical implementation, returns the
// available space on a queue.
//
// There are two important caveats:
// 1. The guarantee does not extend to varying-size tasks due to possible
// internal fragmentation of batches.
// 2. The guarantee only holds in a single-thread environment or critical
// section, i.e. if an intervening thread cannot call Schedule().
//
// This method is useful for monitoring, or for guaranteeing a future slot in
// the schedule (but being mindful about the caveats listed above).
virtual size_t SchedulingCapacity() const = 0;
// Returns the maximum allowed size of tasks submitted to the scheduler. (This
// is typically equal to a configured maximum batch size.)
virtual size_t max_task_size() const = 0;
};
//////////
// Implementation details follow. API users need not read.
template <typename TaskType>
Batch<TaskType>::Batch() : Batch(0) {}
template <typename TaskType>
Batch<TaskType>::Batch(uint64_t traceme_context_id)
: traceme_context_id_(traceme_context_id) {}
template <typename TaskType>
Batch<TaskType>::~Batch() {
WaitUntilClosed();
}
template <typename TaskType>
void Batch<TaskType>::AddTask(std::unique_ptr<TaskType> task,
uint64_t start_time_micros) {
DCHECK(!IsClosed());
{
mutex_lock l(mu_);
size_ += task->size();
tasks_.push_back(std::move(task));
empty_.store(false);
if (tasks_.size() == 1) {
earliest_task_start_time_micros_ = start_time_micros;
} else {
earliest_task_start_time_micros_ =
std::min(earliest_task_start_time_micros_, start_time_micros);
}
}
}
template <typename TaskType>
std::optional<uint64_t> Batch<TaskType>::EarliestTaskStartTime() const {
{
mutex_lock l(mu_);
if (tasks_.empty()) {
return std::nullopt;
}
return earliest_task_start_time_micros_;
}
}
template <typename TaskType>
std::vector<std::unique_ptr<TaskType>> Batch<TaskType>::RemoveAllTasks() {
DCHECK(IsClosed());
{
mutex_lock l(mu_);
size_ = 0;
empty_.store(true);
std::vector<std::unique_ptr<TaskType>> tasks_to_return;
// Swapping vector takes constant time.
tasks_to_return.swap(tasks_);
return std::move(tasks_to_return);
}
}
template <typename TaskType>
std::unique_ptr<TaskType> Batch<TaskType>::RemoveTask() {
{
mutex_lock l(mu_);
if (tasks_.empty()) {
return nullptr;
}
std::unique_ptr<TaskType> task = std::move(tasks_.back());
size_ -= task->size();
tasks_.pop_back();
if (tasks_.empty()) {
empty_.store(true);
}
return task;
}
}
template <typename TaskType>
int Batch<TaskType>::num_tasks() const {
{
mutex_lock l(mu_);
return tasks_.size();
}
}
template <typename TaskType>
bool Batch<TaskType>::empty() const TF_NO_THREAD_SAFETY_ANALYSIS {
// tracer is added to zoom in about this method.
// TODO(b/160249203): Remove tracer after evaluating a change to reduce
// lock contention and cpu usage (which is observed in profiler and
// very data-driven).
tsl::profiler::TraceMe tracer("BatchTask::empty");
return empty_.load();
}
template <typename TaskType>
const TaskType& Batch<TaskType>::task(int i) const {
DCHECK_GE(i, 0);
{
mutex_lock l(mu_);
DCHECK_LT(i, tasks_.size());
return *tasks_[i].get();
}
}
template <typename TaskType>
TaskType* Batch<TaskType>::mutable_task(int i) {
DCHECK_GE(i, 0);
{
mutex_lock l(mu_);
DCHECK_LT(i, tasks_.size());
return tasks_[i].get();
}
}
template <typename TaskType>
size_t Batch<TaskType>::size() const {
{
mutex_lock l(mu_);
return size_;
}
}
template <typename TaskType>
bool Batch<TaskType>::IsClosed() const {
return const_cast<absl::Notification*>(&closed_)->HasBeenNotified();
}
template <typename TaskType>
void Batch<TaskType>::WaitUntilClosed() const {
const_cast<absl::Notification*>(&closed_)->WaitForNotification();
}
template <typename TaskType>
void Batch<TaskType>::Close() {
closed_.Notify();
}
template <typename TaskType>
uint64_t Batch<TaskType>::traceme_context_id() const {
return traceme_context_id_;
}
template <typename TaskType>
void Batch<TaskType>::TryTrimToNewSize(
int new_size, std::vector<std::unique_ptr<TaskType>>& out_trimmed_tasks) {
mutex_lock l(mu_);
DCHECK_GT(new_size, 0);
DCHECK_LT(new_size, size_);
DCHECK(out_trimmed_tasks.empty());
// Index of the first task to trim away. It is possible that it is the index
// of a task of size larger than 1 that will have to be split in order to get
// to the target new_size.
int32_t first_task_to_move = 0;
// The sum of sizes of tasks i, where i < first_task_to_move.
int32_t size_of_previous_tasks = 0;
while (size_of_previous_tasks + tasks_[first_task_to_move]->size() <=
new_size) {
size_of_previous_tasks += tasks_[first_task_to_move]->size();
first_task_to_move++;
// The loop must always stop before this check is tripped because new_size
// must never be larger than the size of the batch.
DCHECK_LT(first_task_to_move, tasks_.size());
}
// Check whether task 'first_task_to_move' will have to be split.
if (size_of_previous_tasks < new_size) {
// TODO: b/325954758 - Consider supporting splitting large tasks and then
// drop 'Try' from the method name.
return;
}
DCHECK_EQ(size_of_previous_tasks, new_size);
// Actually trim.
out_trimmed_tasks.reserve(tasks_.size() - first_task_to_move);
std::move(tasks_.begin() + first_task_to_move, tasks_.end(),
std::back_inserter(out_trimmed_tasks));
tasks_.resize(first_task_to_move);
size_ = new_size;
}
} // namespace serving
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_SCHEDULER_H_
tensorflow/core/kernels/batching_util/batch_scheduler_test.cc
+517
View File
@@ -0,0 +1,517 @@
/* Copyright 2016 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/batching_util/batch_scheduler.h"
#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <tuple>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/synchronization/notification.h"
#include "absl/time/clock.h"
#include "xla/tsl/platform/criticality.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/notification.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace serving {
namespace {
using ::testing::ElementsAre;
using ::testing::Eq;
using ::testing::Pointer;
using ::testing::Property;
TEST(MixedPriorityBatchingPolicyTest, InvalidAttrValueError) {
EXPECT_THAT(
GetMixedPriorityBatchingPolicy("invalid_attr_value"),
absl_testing::StatusIs(
absl::StatusCode::kInvalidArgument,
::testing::HasSubstr(
"Unknown mixed priority batching policy: invalid_attr_value")));
EXPECT_THAT(
GetMixedPriorityBatchingPolicyString(
static_cast<MixedPriorityBatchingPolicy>(4)),
absl_testing::StatusIs(
absl::StatusCode::kInvalidArgument,
::testing::HasSubstr("Unknown mixed priority batching policy: 4")));
}
using MixedPriorityBatchingPolicyParameterizedTest = ::testing::TestWithParam<
std::tuple<std::string, MixedPriorityBatchingPolicy>>;
TEST_P(MixedPriorityBatchingPolicyParameterizedTest,
GetMixedPriorityBatchingPolicySuccess) {
auto [attr_name, policy] = GetParam();
EXPECT_THAT(GetMixedPriorityBatchingPolicy(attr_name),
absl_testing::IsOkAndHolds(Eq(policy)));
EXPECT_THAT(GetMixedPriorityBatchingPolicyString(policy),
absl_testing::IsOkAndHolds(Eq(attr_name)));
}
INSTANTIATE_TEST_SUITE_P(
Parameter, MixedPriorityBatchingPolicyParameterizedTest,
::testing::Values(
std::make_tuple(
/*attr_name=*/kLowPriorityPaddingWithMaxBatchSizeAttrValue,
/*policy=*/MixedPriorityBatchingPolicy::
kLowPriorityPaddingWithMaxBatchSize),
std::make_tuple(
/*attr_name=*/kLowPriorityPaddingWithNextAllowedBatchSizeAttrValue,
/*policy=*/MixedPriorityBatchingPolicy::
kLowPriorityPaddingWithNextAllowedBatchSize),
std::make_tuple(
/*attr_name=*/kPriorityIsolationAttrValue,
/*policy=*/MixedPriorityBatchingPolicy::kPriorityIsolation),
std::make_tuple(
/*attr_name=*/kPriorityMergeAttrValue,
/*policy=*/MixedPriorityBatchingPolicy::kPriorityMerge)));
class FakeTask : public BatchTask {
public:
explicit FakeTask(size_t size) : size_(size) {}
~FakeTask() override = default;
size_t size() const override { return size_; }
private:
const size_t size_;
FakeTask(const FakeTask&) = delete;
void operator=(const FakeTask&) = delete;
};
TEST(TaskCriticalityTest, CriticalityDefaultsToCritical) {
FakeTask fake_task(0);
EXPECT_EQ(fake_task.criticality(), tsl::criticality::Criticality::kCritical);
}
TEST(BatchTaskTest, IsCancelledDefaultsToFalse) {
FakeTask fake_task(0);
EXPECT_FALSE(fake_task.IsCancelled());
}
TEST(BatchTaskTest, IsDeadlineExceededDefaultsToFalse) {
FakeTask fake_task(0);
EXPECT_FALSE(fake_task.IsDeadlineExceeded(absl::Now()));
}
TEST(TaskQueueTest, EmptyTaskQueue) {
TaskQueue<FakeTask> task_queue;
EXPECT_TRUE(task_queue.empty());
EXPECT_EQ(0, task_queue.num_tasks());
EXPECT_EQ(0, task_queue.size());
}
TEST(TaskQueueTest, AddTaskToTaskQueue) {
TaskQueue<FakeTask> task_queue;
task_queue.AddTask(std::make_unique<FakeTask>(1), 1);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(1, task_queue.num_tasks());
EXPECT_EQ(1, task_queue.size());
}
TEST(TaskQueueTest, AddTasksToTaskQueue) {
TaskQueue<FakeTask> task_queue;
task_queue.AddTask(std::make_unique<FakeTask>(1), 1);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(1, task_queue.num_tasks());
EXPECT_EQ(1, task_queue.size());
task_queue.AddTask(std::make_unique<FakeTask>(2), 2);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(2, task_queue.num_tasks());
EXPECT_EQ(3, task_queue.size());
task_queue.AddTask(std::make_unique<FakeTask>(3), 3);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(3, task_queue.num_tasks());
EXPECT_EQ(6, task_queue.size());
}
TEST(TaskQueueTest, RemoveTaskFromTaskQueueWithSingleTask) {
TaskQueue<FakeTask> task_queue;
task_queue.AddTask(std::make_unique<FakeTask>(1), 1);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(1, task_queue.num_tasks());
EXPECT_EQ(1, task_queue.size());
EXPECT_THAT(task_queue.RemoveTask(),
Pointee(Property(&FakeTask::size, Eq(1))));
EXPECT_TRUE(task_queue.empty());
EXPECT_EQ(0, task_queue.num_tasks());
EXPECT_EQ(0, task_queue.size());
}
TEST(TaskQueueTest, RemoveTaskFromTaskQueueWithMultipleTasks) {
TaskQueue<FakeTask> task_queue;
task_queue.AddTask(std::make_unique<FakeTask>(2), 1);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(1, task_queue.num_tasks());
EXPECT_EQ(2, task_queue.size());
task_queue.AddTask(std::make_unique<FakeTask>(1), 2);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(2, task_queue.num_tasks());
EXPECT_EQ(3, task_queue.size());
EXPECT_THAT(task_queue.RemoveTask(),
Pointee(Property(&FakeTask::size, Eq(2))));
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(1, task_queue.num_tasks());
EXPECT_EQ(1, task_queue.size());
}
TEST(TaskQueueTest, RemoveTasksFromTaskQueue) {
TaskQueue<FakeTask> task_queue;
task_queue.AddTask(std::make_unique<FakeTask>(1), 1);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(1, task_queue.num_tasks());
EXPECT_EQ(1, task_queue.size());
task_queue.AddTask(std::make_unique<FakeTask>(2), 2);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(2, task_queue.num_tasks());
EXPECT_EQ(3, task_queue.size());
task_queue.AddTask(std::make_unique<FakeTask>(3), 3);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(3, task_queue.num_tasks());
EXPECT_EQ(6, task_queue.size());
// The first two tasks are removed because they sum up to the size 3 as
// specified.
EXPECT_THAT(task_queue.RemoveTask(3),
ElementsAre(Pointee(Property(&FakeTask::size, Eq(1))),
Pointee(Property(&FakeTask::size, Eq(2)))));
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(1, task_queue.num_tasks());
EXPECT_EQ(3, task_queue.size());
}
TEST(TaskQueueTest, RemoveTasksFewerThanArgFromTaskQueue) {
TaskQueue<FakeTask> task_queue;
task_queue.AddTask(std::make_unique<FakeTask>(1), 1);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(1, task_queue.num_tasks());
EXPECT_EQ(1, task_queue.size());
task_queue.AddTask(std::make_unique<FakeTask>(2), 2);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(2, task_queue.num_tasks());
EXPECT_EQ(3, task_queue.size());
task_queue.AddTask(std::make_unique<FakeTask>(3), 3);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(3, task_queue.num_tasks());
EXPECT_EQ(6, task_queue.size());
// The last task is not removed since that will end up removing tasks that
// sum up to the size larger than 5.
EXPECT_THAT(task_queue.RemoveTask(5),
ElementsAre(Pointee(Property(&FakeTask::size, Eq(1))),
Pointee(Property(&FakeTask::size, Eq(2)))));
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(1, task_queue.num_tasks());
EXPECT_EQ(3, task_queue.size());
}
TEST(TaskQueueTest, RemoveAllTasksWhenArgGreaterThanTaskSize) {
TaskQueue<FakeTask> task_queue;
task_queue.AddTask(std::make_unique<FakeTask>(1), 1);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(1, task_queue.num_tasks());
EXPECT_EQ(1, task_queue.size());
task_queue.AddTask(std::make_unique<FakeTask>(2), 2);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(2, task_queue.num_tasks());
EXPECT_EQ(3, task_queue.size());
task_queue.AddTask(std::make_unique<FakeTask>(3), 3);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(3, task_queue.num_tasks());
EXPECT_EQ(6, task_queue.size());
// All tasks upto the size 6 should be remove when the size 8 is specified.
EXPECT_THAT(task_queue.RemoveTask(8),
ElementsAre(Pointee(Property(&FakeTask::size, Eq(1))),
Pointee(Property(&FakeTask::size, Eq(2))),
Pointee(Property(&FakeTask::size, Eq(3)))));
EXPECT_TRUE(task_queue.empty());
EXPECT_EQ(0, task_queue.num_tasks());
EXPECT_EQ(0, task_queue.size());
}
TEST(TaskQueueTest, EarliestStartTimeWithEmptyQueue) {
TaskQueue<FakeTask> task_queue;
EXPECT_FALSE(task_queue.EarliestTaskStartTime().has_value());
}
TEST(TaskQueueTest, EarliestStartTimeWithMultipleTasksInQueue) {
TaskQueue<FakeTask> task_queue;
task_queue.AddTask(std::make_unique<FakeTask>(1), 1);
task_queue.AddTask(std::make_unique<FakeTask>(2), 2);
std::optional<uint64_t> result = task_queue.EarliestTaskStartTime();
EXPECT_TRUE(result.has_value());
EXPECT_EQ(*result, 1);
}
TEST(TaskQueueTest, EarliestStartTimeAfterTaskRemoval) {
TaskQueue<FakeTask> task_queue;
task_queue.AddTask(std::make_unique<FakeTask>(1), 1);
task_queue.AddTask(std::make_unique<FakeTask>(2), 2);
task_queue.AddTask(std::make_unique<FakeTask>(3), 3);
std::optional<uint64_t> result = task_queue.EarliestTaskStartTime();
EXPECT_TRUE(result.has_value());
EXPECT_EQ(*result, 1);
EXPECT_THAT(task_queue.RemoveTask(3),
ElementsAre(Pointee(Property(&FakeTask::size, Eq(1))),
Pointee(Property(&FakeTask::size, Eq(2)))));
result = task_queue.EarliestTaskStartTime();
EXPECT_TRUE(result.has_value());
EXPECT_EQ(*result, 3);
}
TEST(TaskQueueTest, PrependSingleTask) {
TaskQueue<FakeTask> task_queue;
task_queue.AddTask(std::make_unique<FakeTask>(1), 1);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(1, task_queue.num_tasks());
EXPECT_EQ(1, task_queue.size());
task_queue.AddTask(std::make_unique<FakeTask>(2), 2);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(2, task_queue.num_tasks());
EXPECT_EQ(3, task_queue.size());
task_queue.PrependTask(std::make_unique<FakeTask>(3), 3);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(3, task_queue.num_tasks());
EXPECT_EQ(6, task_queue.size());
EXPECT_THAT(task_queue.RemoveTask(),
Pointee(Property(&FakeTask::size, Eq(3))));
EXPECT_THAT(task_queue.RemoveTask(),
Pointee(Property(&FakeTask::size, Eq(1))));
}
TEST(TaskQueueTest, PrependMultipleTasks) {
TaskQueue<FakeTask> task_queue;
task_queue.AddTask(std::make_unique<FakeTask>(1), 1);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(1, task_queue.num_tasks());
EXPECT_EQ(1, task_queue.size());
task_queue.AddTask(std::make_unique<FakeTask>(2), 2);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(2, task_queue.num_tasks());
EXPECT_EQ(3, task_queue.size());
task_queue.PrependTask(std::make_unique<FakeTask>(3), 3);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(3, task_queue.num_tasks());
EXPECT_EQ(6, task_queue.size());
task_queue.PrependTask(std::make_unique<FakeTask>(4), 4);
EXPECT_FALSE(task_queue.empty());
EXPECT_EQ(4, task_queue.num_tasks());
EXPECT_EQ(10, task_queue.size());
EXPECT_THAT(task_queue.RemoveTask(8),
ElementsAre(Pointee(Property(&FakeTask::size, Eq(4))),
Pointee(Property(&FakeTask::size, Eq(3))),
Pointee(Property(&FakeTask::size, Eq(1)))));
}
TEST(BatchTest, Basic) {
Batch<FakeTask> batch;
EXPECT_EQ(0, batch.num_tasks());
EXPECT_TRUE(batch.empty());
EXPECT_EQ(0, batch.size());
EXPECT_FALSE(batch.IsClosed());
auto task0 = new FakeTask(3);
batch.AddTask(std::unique_ptr<FakeTask>(task0));
EXPECT_EQ(1, batch.num_tasks());
EXPECT_FALSE(batch.empty());
EXPECT_EQ(task0->size(), batch.size());
EXPECT_EQ(task0->size(), batch.task(0).size());
EXPECT_FALSE(batch.IsClosed());
auto task1 = new FakeTask(7);
batch.AddTask(std::unique_ptr<FakeTask>(task1));
EXPECT_EQ(2, batch.num_tasks());
EXPECT_FALSE(batch.empty());
EXPECT_EQ(task0->size() + task1->size(), batch.size());
EXPECT_EQ(task1->size(), batch.task(1).size());
EXPECT_EQ(task1->size(), batch.mutable_task(1)->size());
EXPECT_FALSE(batch.IsClosed());
batch.Close();
EXPECT_TRUE(batch.IsClosed());
EXPECT_EQ(2, batch.num_tasks());
EXPECT_FALSE(batch.empty());
EXPECT_EQ(task0->size() + task1->size(), batch.size());
EXPECT_EQ(task0->size(), batch.task(0).size());
EXPECT_EQ(task1->size(), batch.task(1).size());
EXPECT_EQ(7, batch.RemoveTask()->size());
EXPECT_EQ(3, batch.size());
EXPECT_EQ(3, batch.RemoveTask()->size());
EXPECT_EQ(0, batch.size());
EXPECT_TRUE(batch.empty());
}
TEST(BatchTest, WaitUntilClosed) {
Batch<FakeTask> batch;
batch.AddTask(std::make_unique<FakeTask>(3));
EXPECT_FALSE(batch.IsClosed());
std::unique_ptr<Thread> close_thread(
Env::Default()->StartThread(ThreadOptions(), "test", [&batch]() {
Env::Default()->SleepForMicroseconds(100);
batch.Close();
}));
batch.WaitUntilClosed();
EXPECT_TRUE(batch.IsClosed());
}
TEST(BatchTest, DeletionBlocksUntilClosed) {
Batch<FakeTask>* batch = new Batch<FakeTask>;
batch->AddTask(std::make_unique<FakeTask>(3));
EXPECT_FALSE(batch->IsClosed());
absl::Notification do_delete, deleted;
std::unique_ptr<Thread> delete_thread(Env::Default()->StartThread(
ThreadOptions(), "test", [&batch, &do_delete, &deleted]() {
do_delete.WaitForNotification();
delete batch;
deleted.Notify();
}));
do_delete.Notify();
Env::Default()->SleepForMicroseconds(10 * 1000 /* 10 milliseconds */);
EXPECT_FALSE(deleted.HasBeenNotified());
batch->Close();
deleted.WaitForNotification();
}
TEST(BatchTest, RemoveAllTasks) {
Batch<FakeTask> batch;
auto task0 = new FakeTask(3);
batch.AddTask(std::unique_ptr<FakeTask>(task0));
auto task1 = new FakeTask(7);
batch.AddTask(std::unique_ptr<FakeTask>(task1));
batch.Close();
EXPECT_TRUE(batch.IsClosed());
std::vector<std::unique_ptr<FakeTask>> tasks_in_batch =
batch.RemoveAllTasks();
EXPECT_EQ(2, tasks_in_batch.size());
EXPECT_TRUE(batch.empty());
EXPECT_EQ(task0, tasks_in_batch[0].get());
EXPECT_EQ(task1, tasks_in_batch[1].get());
// RemoveAllTasks returns empty vector from the second call and on, since
// batch is closed.
EXPECT_THAT(batch.RemoveAllTasks(), ::testing::IsEmpty()); // second call
EXPECT_THAT(batch.RemoveAllTasks(), ::testing::IsEmpty()); // third call
}
TEST(BatchTest, TryTrimToNewSizeTrimsAndReturnsTrimmedElementsInOrder) {
Batch<FakeTask> batch;
auto task0 = new FakeTask(3);
batch.AddTask(std::unique_ptr<FakeTask>(task0));
auto task1 = new FakeTask(5);
batch.AddTask(std::unique_ptr<FakeTask>(task1));
auto task2 = new FakeTask(7);
batch.AddTask(std::unique_ptr<FakeTask>(task2));
auto task3 = new FakeTask(9);
batch.AddTask(std::unique_ptr<FakeTask>(task3));
std::vector<std::unique_ptr<FakeTask>> trimmed_tasks;
batch.TryTrimToNewSize(/* new_size= */ 8,
/* out_trimmed_tasks= */ trimmed_tasks);
EXPECT_EQ(batch.size(), 8);
EXPECT_EQ(batch.num_tasks(), 2);
EXPECT_THAT(trimmed_tasks, ElementsAre(Pointer(task2), Pointer(task3)));
batch.Close(); // Batch::~Batch blocks until the batch is closed.
}
TEST(BatchTest, TryTrimToNewSizeDoesNotTrimWhenItWouldNeedToSplitATask) {
Batch<FakeTask> batch;
auto task0 = new FakeTask(3);
batch.AddTask(std::unique_ptr<FakeTask>(task0));
auto task1 = new FakeTask(5);
batch.AddTask(std::unique_ptr<FakeTask>(task1));
std::vector<std::unique_ptr<FakeTask>> trimmed_tasks;
batch.TryTrimToNewSize(/* new_size= */ 4,
/* out_trimmed_tasks= */ trimmed_tasks);
EXPECT_EQ(batch.size(), 8);
EXPECT_EQ(batch.num_tasks(), 2);
EXPECT_TRUE(trimmed_tasks.empty());
batch.Close(); // Batch::~Batch blocks until the batch is closed.
}
} // namespace
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/batch_scheduler_utils.cc
+179
View File
@@ -0,0 +1,179 @@
/* Copyright 2024 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/batching_util/batch_scheduler_utils.h"
#include <algorithm>
#include <cstdint>
#include <optional>
#include <string>
#include <vector>
#include "absl/algorithm/container.h"
#include "absl/log/log.h"
#include "absl/strings/string_view.h"
#include "absl/time/time.h"
#include "tensorflow/core/kernels/batching_util/batch_stats.h"
#include "tensorflow/core/lib/monitoring/counter.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace serving {
int GetNextAllowedBatchSize(int batch_size,
const std::vector<int32_t>& allowed_batch_sizes,
bool disable_padding) {
if (disable_padding || allowed_batch_sizes.empty()) {
return batch_size;
}
DCHECK(absl::c_is_sorted(allowed_batch_sizes));
DCHECK_GT(batch_size, 0);
for (int allowed_size : allowed_batch_sizes) {
if (allowed_size >= batch_size) {
return allowed_size;
}
}
LOG(ERROR) << "Batch size " << batch_size
<< " is greater than largest allowed size; ignoring allowed sizes "
"constraint.";
return batch_size;
}
int32_t GetPrevAllowedBatchSize(int batch_size,
const std::vector<int32_t>& allowed_batch_sizes,
bool disable_padding) {
if (disable_padding || allowed_batch_sizes.empty()) {
return batch_size;
}
DCHECK(absl::c_is_sorted(allowed_batch_sizes));
DCHECK_GT(batch_size, 0);
// First from the end allowed_batch_size not larger than batch_size.
auto result = std::find_if(
allowed_batch_sizes.rbegin(), allowed_batch_sizes.rend(),
[&](int allowed_size) { return allowed_size <= batch_size; });
if (result == allowed_batch_sizes.rend()) {
// No such element exists.
return batch_size;
}
return *result;
}
int ApplyBatchPaddingPolicy(int candidate_size,
const std::vector<int32_t>& allowed_batch_sizes,
bool disable_padding,
absl::string_view batch_padding_policy,
ModelBatchStats* model_batch_stats) {
if (candidate_size == 0) {
return candidate_size;
}
if (batch_padding_policy == kPadUpPolicy) {
return candidate_size;
}
bool minimize_tpu_cost_per_request;
if (batch_padding_policy == kBatchDownPolicy) {
minimize_tpu_cost_per_request = false;
} else if (batch_padding_policy == kMinimizeTpuCostPerRequestPolicy) {
if (model_batch_stats == nullptr) {
LOG_FIRST_N(ERROR, 1)
<< kMinimizeTpuCostPerRequestPolicy
<< " batch padding policy has been chosen "
"but no ModelBatchStats passed to the batch scheduler; will "
"fall back on the "
<< kPadUpPolicy << " policy.";
return candidate_size;
}
minimize_tpu_cost_per_request = true;
} else {
LOG_FIRST_N(ERROR, 1) << "Unsupported batch_padding_policy: "
<< batch_padding_policy << ", falling back on the "
<< kPadUpPolicy << " policy.";
return candidate_size;
}
int32_t pad_up_size = GetNextAllowedBatchSize(
candidate_size, allowed_batch_sizes, disable_padding);
if (pad_up_size == candidate_size) {
return candidate_size; // Good, no padding is necessary.
}
int32_t batch_down_size = GetPrevAllowedBatchSize(
candidate_size, allowed_batch_sizes, disable_padding);
if (batch_down_size == candidate_size) {
return candidate_size; // Can't batch down (e.g. no smaller batch size
// available).
}
if (minimize_tpu_cost_per_request) {
// TODO: b/325954758 - Consider logging a warning here or elsewhere if
// a larger batch doesn't cost meaningfully cheaper than a smaller batch.
// TODO: b/325954758 - Consider logging a warning here or elsewhere if a
// smaller batch costs unreasonably cheaper than a larger one (assuming
// a batch cost model = constant_cost + batch_size * per_element_cost).
// TODO: b/325954758 - Consider occasionally picking either batch size so
// that we learn fresh costs of each batch size. For this code, it is not a
// large priority though because if we are in between two allowed batch
// sizes (say, 16 and 32), chances are that will occasionally organically
// get batches of exact sizes 16 and 32 (and then we pick those
// unconditionally). But if we explicitly occasionally explored other batch
// sizes, we wouldn't have to rely on this "chances are". For other
// applications of batch costs, we might also want to occasionally explore
// all allowed batch sizes and not just 16 and 32 from this example.
std::optional<absl::Duration> down_batch_cost =
model_batch_stats->batch_size(batch_down_size).tpu_cost().mean();
std::optional<absl::Duration> up_batch_cost =
model_batch_stats->batch_size(pad_up_size).tpu_cost().mean();
if (!down_batch_cost.has_value() || !up_batch_cost.has_value()) {
// We have no data about batch costs, let's just do nothing.
return candidate_size;
}
auto batch_down_cost_per_request = *down_batch_cost / batch_down_size;
auto pad_up_cost_per_request = *up_batch_cost / candidate_size;
if (pad_up_cost_per_request < batch_down_cost_per_request) {
// Abort batching down because it's cheaper to pad up.
return candidate_size;
}
}
return batch_down_size;
}
namespace internal {
void RecordLazyCancelledTaskMetrics(int64_t size, absl::string_view reason) {
static auto* count_cell = tensorflow::monitoring::Counter<1>::New(
"/tensorflow/serving/batching/lazy_cancelled_task_count",
"Tracks the number of tasks cancelled due to deadline exceeded or RPC "
"cancellation before batch formation.",
"reason");
count_cell->GetCell(std::string(reason))->IncrementBy(1);
static auto* size_cell = tensorflow::monitoring::Counter<1>::New(
"/tensorflow/serving/batching/lazy_cancelled_task_size",
"Tracks the sum of task sizes dropped due to deadline exceeded or RPC "
"cancellation before batch formation.",
"reason");
size_cell->GetCell(std::string(reason))->IncrementBy(size);
}
} // namespace internal
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/batch_scheduler_utils.h
+111
View File
@@ -0,0 +1,111 @@
/* Copyright 2024 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_SCHEDULER_UTILS_H_
#define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_SCHEDULER_UTILS_H_
#include <cstdint>
#include <memory>
#include <optional>
#include <vector>
#include "absl/strings/string_view.h"
#include "absl/time/time.h"
#include "tensorflow/core/kernels/batching_util/batch_scheduler.h"
#include "tensorflow/core/kernels/batching_util/batch_stats.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace serving {
// Returns the next allowed batch size, which is the smallest allowed batch size
// greater than or equal to the given batch size. If allowed_batch_sizes,
// returns batch_size as is.
int GetNextAllowedBatchSize(int batch_size,
const std::vector<int32_t>& allowed_batch_sizes,
bool disable_padding);
// Returns the largest allowed batch size that is smaller than or equal to
// batch_size. Returns batch_size if no such size exists.
int GetPrevAllowedBatchSize(int batch_size,
const std::vector<int32_t>& allowed_batch_sizes,
bool disable_padding);
// Applies the batch padding policy to the candidate size and returns the target
// size depending on the policy.
int ApplyBatchPaddingPolicy(int candidate_size,
const std::vector<int32_t>& allowed_batch_sizes,
bool disable_padding,
absl::string_view batch_padding_policy,
ModelBatchStats* model_batch_stats);
// Constants containing possible values for the batch_padding_policy argument
// of MaybeBatchDown. This argument specifies the policy that a batch scheduler
// is using when deciding what to do when, say, 18 requests need to be batched,
// but only 16 and 32 batch sizes are allowed. The following options are
// available.
//
// - PAD_UP: pad to size 32.
// - BATCH_DOWN: schedule a batch of size 16 and leave 2 requests in the
// batch buffer.
// - MINIMIZE_TPU_COST_PER_REQUEST: a smarter greedy policy that chooses
// to either PAD_UP or BATCH_DOWN so as to minimize the TPU costs per
// real request. In this case, it would compare (batch_16_cost / 16) and
// (batch_32_cost / 18).
//
inline constexpr absl::string_view kBatchDownPolicy = "BATCH_DOWN";
inline constexpr absl::string_view kPadUpPolicy = "PAD_UP";
inline constexpr absl::string_view kMinimizeTpuCostPerRequestPolicy =
"MINIMIZE_TPU_COST_PER_REQUEST";
// Trims the batch to the next allowed batch size when possible and when
// configured by batch_padding_policy.
//
// When trimming, this function puts the trimmed tasks go into the
// out_trimmed_tasks vector in the same order as they were in the batch.
template <typename TaskType>
void MaybeBatchDown(Batch<TaskType>& batch,
const std::vector<int32_t>& allowed_batch_sizes,
bool disable_padding,
absl::string_view batch_padding_policy,
ModelBatchStats* model_batch_stats,
std::vector<std::unique_ptr<TaskType>>& out_trimmed_tasks) {
if (batch.empty()) {
return;
}
int32_t batch_size = batch.size();
int32_t target_size =
ApplyBatchPaddingPolicy(batch_size, allowed_batch_sizes, disable_padding,
batch_padding_policy, model_batch_stats);
if (target_size < batch_size) {
batch.TryTrimToNewSize(target_size, out_trimmed_tasks);
}
}
namespace internal {
inline constexpr absl::string_view kLazyCancellationReasonDeadlineExceeded =
"deadline_exceeded";
inline constexpr absl::string_view kLazyCancellationReasonRpcCancelled =
"rpc_cancelled";
void RecordLazyCancelledTaskMetrics(int64_t size, absl::string_view reason);
} // namespace internal
} // namespace serving
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_SCHEDULER_UTILS_H_
tensorflow/core/kernels/batching_util/batch_scheduler_utils_test.cc
+256
View File
@@ -0,0 +1,256 @@
/* Copyright 2024 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/batching_util/batch_scheduler_utils.h"
#include <cstddef>
#include <memory>
#include <vector>
#include <gtest/gtest.h>
#include "absl/time/time.h"
#include "tensorflow/core/kernels/batching_util/batch_scheduler.h"
#include "tensorflow/core/kernels/batching_util/batch_stats.h"
namespace tensorflow {
namespace serving {
namespace {
TEST(GetNextAllowedBatchSizeTest, PaddingDisallowed) {
EXPECT_EQ(GetNextAllowedBatchSize(3, {2, 4, 8}, true), 3);
}
TEST(GetNextAllowedBatchSizeTest, EmptyAllowedBatchSizes) {
EXPECT_EQ(GetNextAllowedBatchSize(3, {}, false), 3);
}
TEST(GetNextAllowedBatchSizeTest, NextAllowedBatchSizeFound) {
EXPECT_EQ(GetNextAllowedBatchSize(3, {2, 4, 8}, false), 4);
}
TEST(GetNextAllowedBatchSizeTest, AlreadyAllowedBatchSize) {
EXPECT_EQ(GetNextAllowedBatchSize(2, {2, 4, 8}, false), 2);
}
TEST(GetNextAllowedBatchSizeTest, GreaterThanAllowedBatchSize) {
EXPECT_EQ(GetNextAllowedBatchSize(10, {2, 4, 8}, false), 10);
}
TEST(GetPrevAllowedBatchSizeTest, PaddingDisallowed) {
EXPECT_EQ(GetPrevAllowedBatchSize(3, {2, 4, 8}, true), 3);
}
TEST(GetPrevAllowedBatchSizeTest, EmptyAllowedBatchSizes) {
EXPECT_EQ(GetPrevAllowedBatchSize(3, {}, false), 3);
}
TEST(GetPrevAllowedBatchSizeTest, PrevAllowedBatchSizeFound) {
EXPECT_EQ(GetPrevAllowedBatchSize(3, {1, 2, 4, 8}, false), 2);
}
TEST(GetPrevAllowedBatchSizeTest, NoSmallerAllowedBatchSizeFound) {
EXPECT_EQ(GetPrevAllowedBatchSize(3, {4, 8}, false), 3);
}
TEST(GetPrevAllowedBatchSizeTest, AlreadyAllowedBatchSize) {
EXPECT_EQ(GetPrevAllowedBatchSize(2, {1, 2, 4, 8}, false), 2);
}
TEST(GetPrevAllowedBatchSizeTest, GreaterThanMaxAllowedBatchSize) {
EXPECT_EQ(GetPrevAllowedBatchSize(10, {2, 4, 8}, false), 8);
}
class FakeTask : public BatchTask {
public:
explicit FakeTask(size_t size) : size_(size) {}
size_t size() const override { return size_; }
private:
const size_t size_;
};
TEST(MaybeBatchDownTest, EmptyBatch) {
Batch<FakeTask> batch;
batch.Close();
std::vector<std::unique_ptr<FakeTask>> out_trimmed_tasks;
MaybeBatchDown(
/* batch= */ batch, /* allowed_batch_sizes= */ {1, 2, 4, 8},
/* disable_padding= */ false,
/* batch_padding_policy= */ kBatchDownPolicy,
/* model_batch_stats= */ nullptr,
/* out_trimmed_tasks= */ out_trimmed_tasks);
EXPECT_TRUE(batch.empty());
EXPECT_TRUE(out_trimmed_tasks.empty());
}
TEST(MaybeBatchDownTest, PadUp) {
Batch<FakeTask> batch;
batch.AddTask(std::make_unique<FakeTask>(1));
batch.AddTask(std::make_unique<FakeTask>(1));
batch.AddTask(std::make_unique<FakeTask>(1));
batch.Close();
std::vector<std::unique_ptr<FakeTask>> out_trimmed_tasks;
MaybeBatchDown(
/* batch= */ batch, /* allowed_batch_sizes= */ {1, 2, 4, 8},
/* disable_padding= */ false,
/* batch_padding_policy= */ kPadUpPolicy,
/* model_batch_stats= */ nullptr,
/* out_trimmed_tasks= */ out_trimmed_tasks);
// The batch must stay unchanged (for the batch resource to then pad it to the
// next allowed batch size, thus ending up in a pad-up behavior.)
EXPECT_EQ(batch.size(), 3);
}
TEST(MaybeBatchDownTest, BatchDown) {
Batch<FakeTask> batch;
batch.AddTask(std::make_unique<FakeTask>(1));
batch.AddTask(std::make_unique<FakeTask>(1));
batch.AddTask(std::make_unique<FakeTask>(1));
batch.Close();
std::vector<std::unique_ptr<FakeTask>> out_trimmed_tasks;
MaybeBatchDown(
/* batch= */ batch, /* allowed_batch_sizes= */ {1, 2, 4, 8},
/* disable_padding= */ false,
/* batch_padding_policy= */ kBatchDownPolicy,
/* model_batch_stats= */ nullptr,
/* out_trimmed_tasks= */ out_trimmed_tasks);
// The scheduler should trim the batch to a smaller allowed size that requires
// no padding.
EXPECT_EQ(batch.size(), 2);
// The trimmed part.
EXPECT_EQ(out_trimmed_tasks.size(), 1);
}
TEST(MaybeBatchDownTest, BatchDownDoesNotSplitTasks) {
// Add tasks for size 3, but the second task is large and will have to be
// split if doing batch-down.
Batch<FakeTask> batch;
batch.AddTask(std::make_unique<FakeTask>(1));
batch.AddTask(std::make_unique<FakeTask>(2));
batch.Close();
std::vector<std::unique_ptr<FakeTask>> out_trimmed_tasks;
MaybeBatchDown(
/* batch= */ batch, /* allowed_batch_sizes= */ {1, 2, 4, 8},
/* disable_padding= */ false,
/* batch_padding_policy= */ kBatchDownPolicy,
/* model_batch_stats= */ nullptr,
/* out_trimmed_tasks= */ out_trimmed_tasks);
// The batch must stay unchanged due the fact that the current implementation
// doesn's support splitting large tasks.
EXPECT_EQ(batch.size(), 3);
}
TEST(ApplyBatchPaddingPolicyTest, ZeroCandidateSize) {
EXPECT_EQ(
ApplyBatchPaddingPolicy(0, {2, 4}, false, kBatchDownPolicy, nullptr), 0);
}
TEST(ApplyBatchPaddingPolicyTest, PadUp) {
EXPECT_EQ(ApplyBatchPaddingPolicy(3, {2, 4}, false, kPadUpPolicy, nullptr),
3);
}
TEST(ApplyBatchPaddingPolicyTest, BatchDown) {
EXPECT_EQ(
ApplyBatchPaddingPolicy(3, {2, 4}, false, kBatchDownPolicy, nullptr), 2);
}
TEST(ApplyBatchPaddingPolicyTest,
BatchDownDoesNothingWhenTheBatchSizeIsAlreadyAllowed) {
EXPECT_EQ(
ApplyBatchPaddingPolicy(2, {2, 4}, false, kBatchDownPolicy, nullptr), 2);
}
TEST(ApplyBatchPaddingPolicyTest,
BatchDownDoesNothingWhenNoSmallerAllowedSize) {
EXPECT_EQ(
ApplyBatchPaddingPolicy(1, {2, 4}, false, kBatchDownPolicy, nullptr), 1);
}
TEST(ApplyBatchPaddingPolicyTest,
BatchDownDoesNothingWhenCandidateExceedsAllAllowedSizes) {
// When candidate_size > max(allowed_batch_sizes), pad_up_size falls back to
// candidate_size and the early-return guard fires.
// This should not happen in practice since the candidate size is capped to
// the max execution batch size, which is equal to the max of the allowed
// batch sizes.
EXPECT_EQ(
ApplyBatchPaddingPolicy(10, {2, 4, 8}, false, kBatchDownPolicy, nullptr),
10);
}
TEST(ApplyBatchPaddingPolicyTest, MinimizeTpuCostPerRequestPicksBatchDown) {
ModelBatchStats model_batch_stats;
model_batch_stats.batch_size(2).tpu_cost().Register(absl::Seconds(2));
model_batch_stats.batch_size(4).tpu_cost().Register(absl::Seconds(3.1));
EXPECT_EQ(ApplyBatchPaddingPolicy(3, {2, 4}, false,
kMinimizeTpuCostPerRequestPolicy,
&model_batch_stats),
2);
}
TEST(ApplyBatchPaddingPolicyTest, MinimizeTpuCostPerRequestPicksPadUp) {
ModelBatchStats model_batch_stats;
model_batch_stats.batch_size(2).tpu_cost().Register(absl::Seconds(2));
model_batch_stats.batch_size(4).tpu_cost().Register(absl::Seconds(2.9));
EXPECT_EQ(ApplyBatchPaddingPolicy(3, {2, 4}, false,
kMinimizeTpuCostPerRequestPolicy,
&model_batch_stats),
3);
}
TEST(ApplyBatchPaddingPolicyTest,
MinimizeTpuCostPerRequestMissingCostsReturnsCandidateSize) {
ModelBatchStats model_batch_stats;
model_batch_stats.batch_size(2).tpu_cost().Register(absl::Seconds(2));
EXPECT_EQ(ApplyBatchPaddingPolicy(3, {2, 4}, false,
kMinimizeTpuCostPerRequestPolicy,
&model_batch_stats),
3);
}
TEST(ApplyBatchPaddingPolicyTest,
MinimizeTpuCostPerRequestNoModelStatsReturnsCandidateSize) {
EXPECT_EQ(ApplyBatchPaddingPolicy(3, {2, 4}, false,
kMinimizeTpuCostPerRequestPolicy, nullptr),
3);
}
TEST(ApplyBatchPaddingPolicyTest, UnsupportedPolicy) {
EXPECT_EQ(ApplyBatchPaddingPolicy(3, {2, 4}, false, "UNSUPPORTED", nullptr),
3);
}
} // namespace
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/batch_stats.h
+274
View File
@@ -0,0 +1,274 @@
/* 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.
==============================================================================*/
// The API for reporting and querying batch statistics such as the average batch
// costs for in-process use.
//
// All these statistics can also be retrieved from metrics reported by various
// modules (e.g., batch_resource_base), but it would be slow. This API, on the
// other hand, was designed to be queried on every request.
//
// The classes defined here are not supposed to be instantiated by the user.
// Instead, this file provides a single entry point:
//
// BatchStatsRegistry& GlobalBatchStatsRegistry();
//
// For example, to register batch cost, do:
//
// GlobalBatchStatsRegistry()
// .model(/* model_name= */ "m", /* op_name= */ "o")
// .batch_size(4)
// .tpu_cost
// .Register(cost);
//
// To get the mean cost later, do:
//
// std::optional<absl::Duration> cost =
// .GlobalBatchStatsRegistry()
// .model(/* model_name= */ "m", /* op_name= */ "o")
// .batch_size(4)
// .tpu_cost
// .mean();
//
// It is allowed and safe to store references to intermediate objects here
// because all intermediate objects are guaranteed to never be destroyed.
//
// All operations supported by this API are thread-safe.
#ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_STATS_H_
#define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_STATS_H_
#include <atomic>
#include <cstdint>
#include <optional>
#include <string>
#include <tuple>
#include <vector>
#include "absl/container/node_hash_map.h"
#include "absl/time/time.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/types.h"
#include "tsl/platform/thread_annotations.h"
namespace tensorflow::serving {
// Default values for when there is no recorded statistic in ModelBatchStats.
constexpr int64_t kNumBatchThreadsUnknown = -1;
constexpr int64_t kBatchTimeoutMicrosUnknown = -1;
// Tracks the average cost of registered samples.
//
// Thread-safe.
class CostTracker {
public:
// Registers a cost sample.
void Register(absl::Duration cost) {
DCHECK_GT(cost, absl::ZeroDuration());
mutex_lock l(mu_);
sample_count_++;
sample_sum_ += cost;
};
// Returns the average cost of all registered samples, giving each sample
// the same weight.
//
// Returns std::nullopt if no samples have been registered.
//
// TODO: b/325954758 - Switch this to an exponentially-decaying average. It's
// likely enough to set the half-life to the last 100-1000 samples.
std::optional<absl::Duration> mean() const {
int64_t count;
absl::Duration sum;
{
// We only hold the lock to read the values and release it before later
// performing a relatively slow division operation.
mutex_lock l(mu_);
count = sample_count_;
sum = sample_sum_;
}
if (count == 0) return std::nullopt;
return sum / count;
};
private:
mutable mutex mu_;
int64_t sample_count_ TF_GUARDED_BY(mu_) = 0;
absl::Duration sample_sum_ TF_GUARDED_BY(mu_);
};
// Tracks statistics for a particular model and batch size.
//
// Thread-safe.
class BatchSizeStats {
public:
CostTracker& tpu_cost() { return tpu_cost_; };
private:
CostTracker tpu_cost_;
};
// Tracks statistics for a particular model.
//
// Here, "model" means a specific version of a model (we assume that version is
// encoded in the op_name). In rare cases, when a model version has multiple
// BatchFunction operation, we also treat each such operation as a separate
// model in this context (they should also have different op_names).
//
// Thread-safe.
class ModelBatchStats {
public:
// Returns a reference to the BatchSizeStats instance for the given batch
// size.
//
// The returned reference persist for as long as 'this' is alive.
BatchSizeStats& batch_size(int32_t batch_size) {
mutex_lock l(mu_);
return batch_size_stats_by_batch_size_[batch_size];
}
// Registers that the model server has processed a batch of size `size`
// non-padding tasks for this model, updating the current cumulative
// processed size.
void RegisterProcessedSize(int64_t size) {
cumulative_processed_size_.fetch_add(size, std::memory_order_relaxed);
}
// Returns the cumulative size processed by this model (the total
// count of individual unit-sized queries processed by the model).
int64_t cumulative_processed_size() const {
return cumulative_processed_size_.load(std::memory_order_relaxed);
}
// Returns the list of batch sizes for which this model has statistics.
//
// The returned list is not guaranteed to be sorted.
std::vector<int32_t> BatchSizes() const {
std::vector<int32_t> result;
mutex_lock l(mu_);
result.reserve(batch_size_stats_by_batch_size_.size());
for (const auto& [key, value] : batch_size_stats_by_batch_size_) {
result.push_back(key);
}
return result;
}
void SetNumBatchThreads(int64_t num_batch_threads) {
num_batch_threads_.store(num_batch_threads, std::memory_order_relaxed);
}
int64_t num_batch_threads() const {
return num_batch_threads_.load(std::memory_order_relaxed);
}
void SetBatchTimeoutMicros(int64_t batch_timeout_micros) {
batch_timeout_micros_.store(batch_timeout_micros,
std::memory_order_relaxed);
}
int64_t batch_timeout_micros() const {
return batch_timeout_micros_.load(std::memory_order_relaxed);
}
private:
mutable mutex mu_;
// The storage of all BatchSizeStats instances.
//
// The mutex only protects adding/finding element in the map. Access to
// elements themselves (after they were created) is not protected here. No
// element deletion is possible because we return references to items in this
// map and don't track their lifetime. We are using the node hash map so that
// elements, once created, are fixed in memory.
absl::node_hash_map<int32_t, BatchSizeStats> batch_size_stats_by_batch_size_
TF_GUARDED_BY(mu_);
// The total count of individual unit-sized queries processed by this model.
// Can be used to generate an internal load metric per model. See
// RegisterQuerySize for more details.
std::atomic<int64_t> cumulative_processed_size_ = 0;
// The number of batch threads assigned to this model.
std::atomic<int64_t> num_batch_threads_ = kNumBatchThreadsUnknown;
// The timeout in microseconds for this model (after which the current batch
// is sent to be processed by the TPU).
std::atomic<int64_t> batch_timeout_micros_ = kBatchTimeoutMicrosUnknown;
};
// Tracks batch statistics for all models.
//
// Thread-safe.
class BatchStatsRegistry {
public:
// Returns a reference to ModelBatchStats for the provided model_name and
// op_name.
//
// Upon invocation with a not-yet-seen arguments, creates an empty
// ModelBatchStats instance.
//
// The returned reference persist for as long as 'this' is alive.
ModelBatchStats& model(const std::string& model_name,
const std::string& op_name) {
std::tuple key(model_name, op_name);
mutex_lock l(mu_);
return model_batch_stats_by_model_and_op_names_[key];
}
// Returns a list of all model and op names.
//
// This is the set of model/op names tracked by this BatchStats instance.
// Note that the returned list is not guaranteed to be sorted.
std::vector<std::tuple<std::string, std::string>> ModelAndOpNames() const {
std::vector<std::tuple<std::string, std::string>> result;
mutex_lock l(mu_);
result.reserve(model_batch_stats_by_model_and_op_names_.size());
for (const auto& [key, value] : model_batch_stats_by_model_and_op_names_) {
result.push_back(key);
}
return result;
}
private:
mutable mutex mu_;
// The storage of all ModelBatchStats instances.
//
// The mutex only protects adding/finding element in the map. Access to
// elements themselves (after they were created) is not protected here. No
// element deletion is possible because we return references to items in this
// map and don't track their lifetime. We are using the node hash map for
// element pointer stability.
absl::node_hash_map<std::tuple<std::string, std::string>, ModelBatchStats>
model_batch_stats_by_model_and_op_names_ TF_GUARDED_BY(mu_);
};
// Returns the global instance of BatchStats, to use used for all production
// purposes (one should only instantiate individual classes from this file to
// test them).
inline BatchStatsRegistry& GlobalBatchStatsRegistry() {
static BatchStatsRegistry* instance = new BatchStatsRegistry();
return *instance;
}
} // namespace tensorflow::serving
#endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BATCH_STATS_H_
tensorflow/core/kernels/batching_util/batch_stats_test.cc
+153
View File
@@ -0,0 +1,153 @@
/* 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/kernels/batching_util/batch_stats.h"
#include <tuple>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/time/time.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow::serving {
namespace {
using ::testing::UnorderedElementsAre;
TEST(BatchStatsTest, GlobalBatchStatsRegistryAlwaysReturnsTheSameInstance) {
ASSERT_EQ(&GlobalBatchStatsRegistry(), &GlobalBatchStatsRegistry());
}
TEST(BatchStatsTest, BasicOperation) {
BatchStatsRegistry stats;
stats.model(/* model_name= */ "m", /* op_name= */ "o")
.batch_size(1)
.tpu_cost()
.Register(absl::Hours(5));
ASSERT_EQ(stats.model(/* model_name= */ "m", /* op_name= */ "o")
.batch_size(1)
.tpu_cost()
.mean(),
absl::Hours(5));
}
TEST(BatchStatsTest, ModelBatchStatsAreUniqueForEachModel) {
BatchStatsRegistry stats;
ASSERT_NE(&stats.model(/* model_name= */ "m", /* op_name= */ "o"),
&stats.model(/* model_name= */ "m", /* op_name= */ "o2"));
}
TEST(BatchStatsTest, BatchSizeStatsAreUniqueForEachBatchSize) {
ModelBatchStats stats;
ASSERT_NE(&stats.batch_size(1), &stats.batch_size(2));
}
TEST(BatchStatsTest, CostTrackerStartsWithNoMean) {
CostTracker tracker;
ASSERT_FALSE(tracker.mean().has_value());
}
TEST(BatchStatsTest, CostTrackerMeanIsCorrect) {
CostTracker tracker;
tracker.Register(absl::Hours(5));
tracker.Register(absl::Hours(7));
ASSERT_EQ(*tracker.mean(), absl::Hours(6));
}
TEST(BatchStatsTest, ProcessedSizeIsCorrect) {
ModelBatchStats stats;
stats.RegisterProcessedSize(5);
stats.RegisterProcessedSize(7);
ASSERT_EQ(stats.cumulative_processed_size(), 12);
}
TEST(BatchStatsTest, ModelOpNamesAreCorrect) {
BatchStatsRegistry stats;
// Register a cost for model "m" and op "o".
stats.model(/* model_name= */ "m", /* op_name= */ "o")
.batch_size(1)
.tpu_cost()
.Register(absl::Hours(5));
// Register a cost for model "m2" and op "o".
stats.model(/* model_name= */ "m2", /* op_name= */ "o")
.batch_size(1)
.tpu_cost()
.Register(absl::Hours(7));
// Register another cost for model "m" and op "o" (but different batch size).
stats.model(/* model_name= */ "m", /* op_name= */ "o")
.batch_size(2)
.tpu_cost()
.Register(absl::Hours(4));
// Register a cost for model "m" and op "o2".
stats.model(/* model_name= */ "m", /* op_name= */ "o2")
.batch_size(1)
.tpu_cost()
.Register(absl::Hours(1));
// Check that the model/op names are correct.
ASSERT_THAT(stats.ModelAndOpNames(),
UnorderedElementsAre(
std::tuple(/* model_name= */ "m", /* op_name= */ "o"),
std::tuple(/* model_name= */ "m", /* op_name= */ "o2"),
std::tuple(/* model_name= */ "m2", /* op_name= */ "o")));
}
TEST(BatchStatsTest, BatchSizesAreCorrect) {
ModelBatchStats stats;
// Register costs for batch sizes 1, 2, and 4.
stats.batch_size(1).tpu_cost().Register(absl::Hours(5));
stats.batch_size(4).tpu_cost().Register(absl::Hours(7));
stats.batch_size(1).tpu_cost().Register(absl::Hours(4));
stats.batch_size(2).tpu_cost().Register(absl::Hours(1));
// Check that the batch sizes are correct.
ASSERT_THAT(stats.BatchSizes(), UnorderedElementsAre(1, 2, 4));
}
TEST(BatchStatsTest, BatchTimeoutIsCorrect) {
ModelBatchStats stats;
// Originally the batch timeout is -1 if unassigned.
ASSERT_EQ(stats.batch_timeout_micros(), -1);
// Assign a batch timeout of 100 microseconds.
stats.SetBatchTimeoutMicros(100);
ASSERT_EQ(stats.batch_timeout_micros(), 100);
}
TEST(BatchStatsTest, NumBatchThreadsIsCorrect) {
ModelBatchStats stats;
// Originally the number of batch threads is -1 if unassigned.
ASSERT_EQ(stats.num_batch_threads(), -1);
// Assign a number of per-model batch threads.
stats.SetNumBatchThreads(16);
ASSERT_EQ(stats.num_batch_threads(), 16);
}
} // namespace
} // namespace tensorflow::serving
tensorflow/core/kernels/batching_util/bounded_executor.cc
+104
View File
@@ -0,0 +1,104 @@
/* 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/kernels/batching_util/bounded_executor.h"
#include <algorithm>
#include <atomic>
#include "absl/functional/bind_front.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/threadpool.h"
namespace tensorflow {
namespace serving {
absl::StatusOr<std::unique_ptr<BoundedExecutor>> BoundedExecutor::Create(
const Options& options) {
if (options.env == nullptr) {
return absl::InvalidArgumentError("options.env must not be nullptr");
}
if (options.num_threads <= 0) {
return absl::InvalidArgumentError("options.num_threads must be positive");
}
return absl::WrapUnique(new BoundedExecutor(options));
}
BoundedExecutor::BoundedExecutor(const Options& options) : options_(options) {
InitWorker();
}
void BoundedExecutor::InitWorker() {
for (int i = 0; i < options_.num_threads; i++) {
std::unique_ptr<Thread> thread = absl::WrapUnique(
options_.env->StartThread(options_.thread_options, options_.thread_name,
[this]() { this->Run(); }));
threads_.push_back(std::move(thread));
}
}
BoundedExecutor::~BoundedExecutor() {
{
mutex_lock l(work_queue_mu_);
// Enqueue an empty task (nullptr) to signal exit.
// This way, each thread blocks on waiting a task, and exit run-loop
// if task is nullptr.
for (int i = 0; i < NumThreads(); i++) {
work_queue_.push_back(nullptr);
work_queue_cv_.notify_one();
}
}
// Each thread will be joined in its destructor.
threads_.clear();
}
void BoundedExecutor::Schedule(std::function<void()> func) {
// use DCHECK so as not to introduce CHECK in prod code.
DCHECK(func != nullptr) << "func is nullptr";
mutex_lock l(work_queue_mu_);
work_queue_.push_back(std::move(func));
work_queue_cv_.notify_one();
}
int BoundedExecutor::NumThreads() const { return options_.num_threads; }
int BoundedExecutor::CurrentThreadId() const { return -1; }
void BoundedExecutor::Run() {
while (true) {
std::function<void()> func = nullptr;
{
mutex_lock l(work_queue_mu_);
while (work_queue_.empty()) {
work_queue_cv_.wait(l);
}
func = std::move(work_queue_.front());
work_queue_.pop_front();
}
// Exit run-loop when func is nullptr.
if (func != nullptr) {
func();
} else {
break;
}
}
}
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/bounded_executor.h
+80
View File
@@ -0,0 +1,80 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BOUNDED_EXECUTOR_H_
#define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BOUNDED_EXECUTOR_H_
#include <string>
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/platform/threadpool.h"
#include "tensorflow/core/platform/threadpool_interface.h"
namespace tensorflow {
namespace serving {
// BoundedExecutor has a bounded number of threads and unlimited queue length,
// scheduled tasks are executed in a FIFO way.
class BoundedExecutor : public thread::ThreadPoolInterface {
public:
struct Options {
Env* env = Env::Default();
ThreadOptions thread_options;
std::string thread_name;
int num_threads = -1;
};
static absl::StatusOr<std::unique_ptr<BoundedExecutor>> Create(
const Options& options);
// Destructor. All threads will be joined.
~BoundedExecutor() override;
// Enqueue a function to be executed.
//
// Callers are responsible to guarantee `func` is not nullptr.
void Schedule(std::function<void()> func) override;
// Returns the number of threads.
int NumThreads() const override;
int CurrentThreadId() const override;
private:
explicit BoundedExecutor(const Options& options);
// Starts N workers (N == num_threads), polling tasks from `work_queue_`.
void InitWorker();
// A loop to fetch task from `work_queue_` and execute task.
void Run();
const Options& options_;
mutex work_queue_mu_;
std::deque<std::function<void()>> work_queue_ TF_GUARDED_BY(work_queue_mu_);
condition_variable work_queue_cv_ TF_GUARDED_BY(work_queue_mu_);
// A fixed number of threads.
std::vector<std::unique_ptr<Thread>> threads_;
BoundedExecutor(const BoundedExecutor&) = delete;
void operator=(const BoundedExecutor&) = delete;
};
} // namespace serving
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_BOUNDED_EXECUTOR_H_
tensorflow/core/kernels/batching_util/bounded_executor_test.cc
+152
View File
@@ -0,0 +1,152 @@
/* 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/kernels/batching_util/bounded_executor.h"
#include "absl/functional/bind_front.h"
#include "absl/synchronization/notification.h"
#include "absl/time/time.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/threadpool_interface.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
namespace tensorflow {
namespace serving {
namespace {
// Tracks the number of concurrently running tasks.
class TaskTracker {
public:
// Creates a functor that invokes Run() with the given arguments.
std::function<void()> MakeTask(int task_id, absl::Duration sleep_duration) {
return absl::bind_front(&TaskTracker::Run, this, task_id, sleep_duration);
}
// Updates run counts, sleeps for a short time and then returns.
// Exits early if fiber is cancelled.
void Run(int task_id, absl::Duration sleep_duration) {
LOG(INFO) << "Entering task " << task_id;
// Update run counters.
{
mutex_lock l(mutex_);
++task_count_;
++running_count_;
if (running_count_ > max_running_count_) {
max_running_count_ = running_count_;
}
}
// Use a sleep loop so we can quickly detect cancellation even when the
// total sleep time is very large.
Env::Default()->SleepForMicroseconds(
absl::ToInt64Microseconds(sleep_duration));
// Update run counters.
{
mutex_lock l(mutex_);
--running_count_;
}
LOG(INFO) << "Task " << task_id << " exiting.";
}
// Returns number of tasks that have been run.
int task_count() {
mutex_lock l(mutex_);
return task_count_;
}
// Returns number of tasks that are currently running.
int running_count() {
mutex_lock l(mutex_);
return running_count_;
}
// Returns the max number of tasks that have run concurrently.
int max_running_count() {
mutex_lock l(mutex_);
return max_running_count_;
}
private:
mutex mutex_;
int task_count_ = 0;
int running_count_ = 0;
int max_running_count_ = 0;
};
TEST(BoundedExecutorTest, InvalidEmptyEnv) {
BoundedExecutor::Options options;
options.num_threads = 2;
options.env = nullptr;
EXPECT_THAT(BoundedExecutor::Create(options),
absl_testing::StatusIs(error::INVALID_ARGUMENT,
"options.env must not be nullptr"));
}
TEST(BoundedExecutorTest, InvalidNumThreads) {
{
BoundedExecutor::Options options;
options.num_threads = 0;
EXPECT_THAT(BoundedExecutor::Create(options),
absl_testing::StatusIs(error::INVALID_ARGUMENT,
"options.num_threads must be positive"));
}
{
BoundedExecutor::Options options;
options.num_threads = -1;
EXPECT_THAT(BoundedExecutor::Create(options),
absl_testing::StatusIs(error::INVALID_ARGUMENT,
"options.num_threads must be positive"));
}
}
TEST(BoundedExecutorTest, AddRunsFunctionsEventually) {
BoundedExecutor::Options options;
options.num_threads = 2;
TF_ASSERT_OK_AND_ASSIGN(auto executor, BoundedExecutor::Create(options));
absl::Notification done0;
executor->Schedule([&done0] { done0.Notify(); });
absl::Notification done1;
executor->Schedule([&done1] { done1.Notify(); });
done0.WaitForNotification();
done1.WaitForNotification();
executor.reset();
}
TEST(BoundedExecutorTest, MaxInflightLimit) {
BoundedExecutor::Options options;
options.num_threads = 5;
TF_ASSERT_OK_AND_ASSIGN(auto executor, BoundedExecutor::Create(options));
const int num_tasks = 100;
TaskTracker task_tracker;
for (int i = 0; i < num_tasks; i++) {
executor->Schedule(task_tracker.MakeTask(i, absl::Seconds(1)));
}
executor.reset();
EXPECT_EQ(task_tracker.task_count(), num_tasks);
EXPECT_EQ(task_tracker.max_running_count(), options.num_threads);
EXPECT_EQ(task_tracker.running_count(), 0);
}
} // namespace
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/concat_split_util.h
+255
View File
@@ -0,0 +1,255 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_CONCAT_SPLIT_UTIL_H_
#define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_CONCAT_SPLIT_UTIL_H_
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/ops_util.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/kernels/concat_lib.h"
#include "tensorflow/core/kernels/split_lib.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace concat_split_util {
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
// Concatenates 'inputs' into a single tensor along the zeroth dimension.
// Requires that all elements of 'inputs' have element type T. Writes to
// 'output' using 'context' for the allocation to ensure proper device
// placement.
template <typename T>
absl::Status Concat(OpKernelContext* context,
const absl::Span<const Tensor> inputs, Tensor* output) {
const int input_dims = inputs[0].dims();
const TensorShape& input_shape = inputs[0].shape();
// Note that we reduce the concat of k-dimensional tensors into a two
// dimensional concat. Assuming the dimensions of any input tensor are
// {y0, y1,...,ym-1}, we flatten it to {1, y}, where y = Prod_i(yi).
std::vector<std::unique_ptr<typename TTypes<T, 2>::ConstMatrix>> inputs_flat;
inputs_flat.reserve(inputs.size());
int64_t output_dim0 = 0;
for (size_t i = 0; i < inputs.size(); ++i) {
const Tensor& input = inputs[i];
if (input.dims() != input_dims) {
return absl::InvalidArgumentError(
absl::StrCat("Ranks of all input tensors should match: shape[0] = ",
input_shape.DebugString(), " vs. shape[", i,
"] = ", input.shape().DebugString()));
}
for (int j = 1; j < input_dims; ++j) {
if (input.dim_size(j) != input_shape.dim_size(j)) {
return absl::InvalidArgumentError(
absl::StrCat("Dimensions of inputs should match: shape[0] = ",
input_shape.DebugString(), " vs. shape[", i,
"] = ", input.shape().DebugString()));
}
}
if (input.NumElements() > 0) {
inputs_flat.emplace_back(new typename TTypes<T, 2>::ConstMatrix(
input.shaped<T, 2>({1, input.NumElements()})));
}
output_dim0 += input.dim_size(0);
}
TensorShape output_shape(input_shape);
output_shape.set_dim(0, output_dim0);
AllocatorAttributes attr;
attr.set_on_host(true);
TF_RETURN_IF_ERROR(context->allocate_temp(DataTypeToEnum<T>::value,
output_shape, output, attr));
if (output->NumElements() > 0) {
auto output_flat = output->shaped<T, 2>({1, output->NumElements()});
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
if (std::is_same<Device, GPUDevice>::value) {
ConcatGPU<T>(context, inputs_flat, output, &output_flat);
return absl::OkStatus();
}
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
ConcatCPU<T>(context->device(), inputs_flat, &output_flat);
}
return absl::OkStatus();
}
// Same as 'Concat' above, but handles Tensor dtype deduction automatically.
inline absl::Status Concat(OpKernelContext* context,
const absl::Span<const Tensor> inputs,
Tensor* output) {
const DataType type = inputs[0].dtype();
absl::Status concat_status;
switch (type) {
#define CASE(type) \
case DataTypeToEnum<type>::value: \
concat_status = Concat<type>(context, inputs, output); \
break;
TF_CALL_ALL_TYPES(CASE) TF_CALL_float8_e4m3fn(CASE);
#undef CASE
default:
concat_status = absl::InvalidArgumentError(
absl::StrCat("Unsupported data type: ", type));
break;
}
return concat_status;
}
// The Split*() functions split 'input' with element type T into 'sizes.size()'
// tensors along the zeroth dimension, with the ith split having zeroth-
// dimension size 'sizes[i]'. They allocate the output tensors using 'context',
// for proper device placement.
// Handles special cases that are cheap. Sets 'done==true' iff it found an
// applicable special case and wrote to the outputs. Otherwise acts as a no-op.
template <typename T>
absl::Status SplitEasyCases(OpKernelContext* context, const Tensor& input,
const absl::Span<const int64_t> sizes,
std::vector<Tensor>* outputs, bool* done) {
*done = false;
int64_t total_size = 0;
for (const int64_t size : sizes) {
total_size += size;
}
if (total_size > input.shape().dim_size(0)) {
return absl::InvalidArgumentError(
"Sum of split sizes must not exceed dim0-size of input tensor");
}
// Special case 0: trivial 1-way split.
if (sizes.size() == 1 && sizes.at(0) == input.shape().dim_size(0)) {
outputs->push_back(input);
*done = true;
return absl::OkStatus();
}
// Special case 1: input is aligned.
if (IsInnerDimsSizeAligned<T>(input.shape())) {
int64_t position = 0;
for (const int64_t size : sizes) {
outputs->emplace_back(input.Slice(position, position + size));
position += size;
}
*done = true;
return absl::OkStatus();
}
return absl::OkStatus();
}
// Handles the general case, on CPU.
template <typename T>
absl::Status SplitCPU(OpKernelContext* context, const Tensor& input,
const absl::Span<const int64_t> sizes,
std::vector<Tensor>* outputs) {
int64_t suffix_dim_size = 1;
for (int i = 1; i < input.shape().dims(); ++i) {
suffix_dim_size *= input.shape().dim_size(i);
}
auto input_reshaped =
input.shaped<T, 2>({input.shape().dim_size(0), suffix_dim_size});
int64_t position = 0;
for (const int64_t size : sizes) {
TensorShape output_shape = input.shape();
output_shape.set_dim(0, size);
Tensor output;
AllocatorAttributes attr;
attr.set_on_host(true);
TF_RETURN_IF_ERROR(
context->allocate_temp(input.dtype(), output_shape, &output, attr));
auto output_shaped = output.shaped<T, 2>({size, suffix_dim_size});
Eigen::DSizes<Eigen::DenseIndex, 2> slice_indices{
static_cast<Eigen::DenseIndex>(position), 0};
Eigen::DSizes<Eigen::DenseIndex, 2> slice_sizes{
static_cast<Eigen::DenseIndex>(size),
static_cast<Eigen::DenseIndex>(suffix_dim_size)};
functor::Split<CPUDevice, T, 2>()(context->eigen_device<CPUDevice>(),
output_shaped, input_reshaped,
slice_indices, slice_sizes);
outputs->emplace_back(output);
position += size;
}
return absl::OkStatus();
}
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
// Handles the general case, on GPU.
template <typename T>
absl::Status SplitGPU(OpKernelContext* context, const Tensor& input,
const absl::Span<const int64_t>& sizes,
std::vector<Tensor>* outputs) {
// TODO(olston, apassos): Implement this.
LOG(FATAL) << "Not yet implemented"; // Crash ok
}
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
// The outer function that dispatches to the various Split*() functions above.
template <typename T>
absl::Status Split(OpKernelContext* context, const Tensor& input,
const absl::Span<const int64_t> sizes,
std::vector<Tensor>* outputs) {
bool easy_cases_done;
TF_RETURN_IF_ERROR(
SplitEasyCases<T>(context, input, sizes, outputs, &easy_cases_done));
if (easy_cases_done) {
return absl::OkStatus();
}
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
// TODO(olston, apassos): Handle non-CPU cases.
// return SplitGPU<T>(context, input, sizes, outputs);
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
return SplitCPU<T>(context, input, sizes, outputs);
}
// Same as 'Split' above, but handles Tensor dtype automatically.
inline absl::Status Split(OpKernelContext* context, const Tensor& input,
const absl::Span<const int64_t> sizes,
std::vector<Tensor>* outputs) {
const DataType type = input.dtype();
absl::Status split_status;
switch (type) {
#define CASE(type) \
case DataTypeToEnum<type>::value: \
split_status = Split<type>(context, input, sizes, outputs); \
break;
TF_CALL_ALL_TYPES(CASE) TF_CALL_float8_e4m3fn(CASE);
#undef CASE
default:
split_status = absl::InvalidArgumentError(
absl::StrCat("Unsupported data type: ", type));
break;
}
return split_status;
}
} // namespace concat_split_util
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_CONCAT_SPLIT_UTIL_H_
tensorflow/core/kernels/batching_util/fake_clock_env.cc
+92
View File
@@ -0,0 +1,92 @@
/* Copyright 2016 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/batching_util/fake_clock_env.h"
#include <string>
#include "absl/synchronization/notification.h"
namespace tensorflow {
namespace serving {
namespace test_util {
FakeClockEnv::FakeClockEnv(Env* wrapped) : EnvWrapper(wrapped) {}
void FakeClockEnv::AdvanceByMicroseconds(int micros) {
{
mutex_lock l(mu_);
current_time_ += micros;
for (auto it = sleeping_threads_.begin(); it != sleeping_threads_.end();) {
if (current_time_ >= it->wake_time) {
it->wake_notification->Notify();
it = sleeping_threads_.erase(it);
} else {
++it;
}
}
}
}
void FakeClockEnv::BlockUntilSleepingThread(uint64_t wake_time) {
for (;;) {
{
mutex_lock l(mu_);
for (auto it = sleeping_threads_.begin(); it != sleeping_threads_.end();
++it) {
if (it->wake_time == wake_time) {
return;
}
}
}
EnvWrapper::SleepForMicroseconds(100);
}
}
void FakeClockEnv::BlockUntilThreadsAsleep(int num_threads) {
for (;;) {
{
mutex_lock l(mu_);
if (num_threads <= sleeping_threads_.size()) {
return;
}
}
EnvWrapper::SleepForMicroseconds(100);
}
}
uint64_t FakeClockEnv::NowMicros() const {
{
mutex_lock l(mu_);
return current_time_;
}
}
void FakeClockEnv::SleepForMicroseconds(int64_t micros) {
if (micros == 0) {
return;
}
absl::Notification wake_notification;
{
mutex_lock l(mu_);
sleeping_threads_.push_back({current_time_ + micros, &wake_notification});
}
wake_notification.WaitForNotification();
}
} // namespace test_util
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/fake_clock_env.h
+77
View File
@@ -0,0 +1,77 @@
/* Copyright 2016 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_FAKE_CLOCK_ENV_H_
#define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_FAKE_CLOCK_ENV_H_
#include <functional>
#include <string>
#include <vector>
#include "absl/synchronization/notification.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace serving {
namespace test_util {
// An Env implementation with a fake clock for NowMicros() and
// SleepForMicroseconds(). The clock doesn't advance on its own; it advances via
// an explicit Advance() method.
// All other Env virtual methods pass through to a wrapped Env.
class FakeClockEnv : public EnvWrapper {
public:
explicit FakeClockEnv(Env* wrapped);
~FakeClockEnv() override = default;
// Advance the clock by a certain number of microseconds.
void AdvanceByMicroseconds(int micros);
// Blocks until there is a sleeping thread that is scheduled to wake up at
// the given (absolute) time.
void BlockUntilSleepingThread(uint64_t wake_time);
// Blocks until there are at least num_threads sleeping.
void BlockUntilThreadsAsleep(int num_threads);
// Methods that this class implements.
uint64_t NowMicros() const override;
void SleepForMicroseconds(int64_t micros) override;
private:
mutable mutex mu_;
uint64_t current_time_ TF_GUARDED_BY(mu_) = 0;
struct SleepingThread {
uint64_t wake_time;
absl::Notification* wake_notification;
};
std::vector<SleepingThread> sleeping_threads_ TF_GUARDED_BY(mu_);
FakeClockEnv(const FakeClockEnv&) = delete;
void operator=(const FakeClockEnv&) = delete;
};
} // namespace test_util
} // namespace serving
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_FAKE_CLOCK_ENV_H_
tensorflow/core/kernels/batching_util/input_split_metadata.cc
+99
View File
@@ -0,0 +1,99 @@
/* 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/kernels/batching_util/input_split_metadata.h"
#include <algorithm>
#include "absl/container/fixed_array.h"
#include "absl/strings/str_join.h"
namespace tensorflow {
namespace serving {
namespace internal {
namespace {
int compute_task_size_from_open_batch(int input_task_size,
int open_batch_remaining_slot,
int batch_size_limit) {
return (open_batch_remaining_slot > 0)
? (input_task_size + batch_size_limit - open_batch_remaining_slot)
: input_task_size;
}
int compute_head_task_size(int input_task_size, int open_batch_remaining_slot,
int batch_size_limit) {
if (open_batch_remaining_slot == 0) {
return std::min(input_task_size, batch_size_limit);
}
return std::min(open_batch_remaining_slot, input_task_size);
}
int compute_tail_task_size(int task_size_from_open_batch, int input_task_size,
int open_batch_remaining_slot,
int batch_size_limit) {
int tail_task_size;
if (input_task_size <= open_batch_remaining_slot) {
tail_task_size = input_task_size;
} else {
tail_task_size = task_size_from_open_batch % batch_size_limit;
if (tail_task_size == 0) {
tail_task_size = batch_size_limit;
}
}
return tail_task_size;
}
int compute_num_batches(int task_size_from_open_batch, int batch_size_limit) {
return (task_size_from_open_batch + batch_size_limit - 1) / batch_size_limit;
}
} // namespace
InputSplitMetadata::InputSplitMetadata(int input_task_size,
int open_batch_remaining_slot,
int batch_size_limit)
: task_sizes_(generate_task_sizes(
input_task_size, open_batch_remaining_slot, batch_size_limit)) {}
const absl::FixedArray<int>& InputSplitMetadata::task_sizes() const {
return task_sizes_;
}
std::string InputSplitMetadata::DebugString() const {
return absl::StrJoin(task_sizes_, ", ");
}
absl::FixedArray<int> InputSplitMetadata::generate_task_sizes(
int input_task_size, int open_batch_remaining_slot,
int batch_size_limit) const {
const int task_size_from_open_batch = compute_task_size_from_open_batch(
input_task_size, open_batch_remaining_slot, batch_size_limit);
const int num_batches =
compute_num_batches(task_size_from_open_batch, batch_size_limit);
absl::FixedArray<int> task_sizes(num_batches, batch_size_limit);
task_sizes.front() = compute_head_task_size(
input_task_size, open_batch_remaining_slot, batch_size_limit);
task_sizes.back() =
compute_tail_task_size(task_size_from_open_batch, input_task_size,
open_batch_remaining_slot, batch_size_limit);
return task_sizes;
}
} // namespace internal
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/input_split_metadata.h
+55
View File
@@ -0,0 +1,55 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_INPUT_SPLIT_METADATA_H_
#define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_INPUT_SPLIT_METADATA_H_
#include <algorithm>
#include "absl/container/fixed_array.h"
namespace tensorflow {
namespace serving {
namespace internal {
// InputSplitMetadata represents the task sizes of an batch-task after it's
// tailored according to queue status (`open_batch_remaining_slot` and
// `batch_size_limit`).
//
// This is an internal helper class, and the implementation is shared
// shared across different instantiations of internal::Queue<TaskType>
// in input-split mode (QueueOptions.enable_large_batch_splitting is true).
class InputSplitMetadata {
public:
InputSplitMetadata(int input_task_size, int open_batch_remaining_slot,
int batch_size_limit);
// Returns underlying task sizes.
const absl::FixedArray<int>& task_sizes() const;
// Serializes task split metadata into a string for debugging.
std::string DebugString() const;
private:
absl::FixedArray<int> generate_task_sizes(int input_task_size,
int open_batch_remaining_slot,
int batch_size_limit) const;
const absl::FixedArray<int> task_sizes_;
};
} // namespace internal
} // namespace serving
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_INPUT_SPLIT_METADATA_H_
tensorflow/core/kernels/batching_util/input_split_metadata_test.cc
+74
View File
@@ -0,0 +1,74 @@
/* 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/kernels/batching_util/input_split_metadata.h"
#include "absl/strings/str_join.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace serving {
namespace internal {
namespace {
TEST(InputSplitUtilTest, Basic) {
for (const auto& batch_task_param :
{std::tuple<int /* input_size */, int /* open_batch_remaining_slot */,
int /* batch_size_limit */, int /* expected_num_batches */,
int /* expected_num_new_batches */,
int /* expected_head_batch_task_size */,
int /* expected_tail_batch_task_size */>{5, 1, 1, 5, 4, 1,
1},
{10, 3, 4, 3, 2, 3, 3},
{20, 5, 6, 4, 3, 5, 3},
{30, 0, 11, 3, 3, 11, 8},
{5, 6, 8, 1, 0, 5, 5}}) {
const int input_size = std::get<0>(batch_task_param);
const int open_batch_remaining_slot = std::get<1>(batch_task_param);
const int batch_size_limit = std::get<2>(batch_task_param);
const int expected_num_batches = std::get<3>(batch_task_param);
const int expected_head_batch_task_size = std::get<5>(batch_task_param);
const int expected_tail_batch_task_size = std::get<6>(batch_task_param);
// The number of remaining slots should be smaller than or equal to
// batch_size_limit; whearas we allow one input (of `input_size`) to span
// over multiple batches.
ASSERT_LE(open_batch_remaining_slot, batch_size_limit);
InputSplitMetadata input_split_metadata(
input_size, open_batch_remaining_slot, batch_size_limit);
EXPECT_EQ(input_split_metadata.task_sizes().size(), expected_num_batches);
absl::FixedArray<int> expected_task_sizes(expected_num_batches);
for (int i = 0; i < expected_num_batches; i++) {
if (i == 0) {
expected_task_sizes[i] = expected_head_batch_task_size;
} else if (i == expected_num_batches - 1) {
expected_task_sizes[i] = expected_tail_batch_task_size;
} else {
expected_task_sizes[i] = batch_size_limit;
}
}
EXPECT_THAT(input_split_metadata.task_sizes(),
::testing::ElementsAreArray(expected_task_sizes));
EXPECT_EQ(input_split_metadata.DebugString(),
absl::StrJoin(expected_task_sizes, ", "));
}
}
} // namespace
} // namespace internal
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/periodic_function.cc
+104
View File
@@ -0,0 +1,104 @@
/* Copyright 2016 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/batching_util/periodic_function.h"
#include <algorithm>
#include <utility>
#include "absl/functional/any_invocable.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/logging.h"
namespace tensorflow {
namespace serving {
PeriodicFunction::PeriodicFunction(absl::AnyInvocable<void()> function,
const int64_t interval_micros,
const Options& options)
: function_(std::move(function)),
interval_micros_([interval_micros]() -> int64_t {
if (interval_micros < 0) {
const std::string error =
absl::StrCat(" The value of 'interval_micros' should be >= 0: ",
interval_micros, ". ");
DCHECK(false) << error;
LOG(WARNING) << error << "Resetting it to 0.";
return 0;
}
return interval_micros;
}()),
options_(options) {
thread_.reset(options_.env->StartThread(
options_.thread_options, options_.thread_name_prefix, [this]() {
// Record the starting time here instead of in RunLoop. That way, if
// there is a delay starting RunLoop, that does not affect the timing
// of
// the first function. (Such a delay can often happen in tests where
// the test simulates a large time delay immediately after calling
// Start.)
RunLoop(options_.env->NowMicros());
}));
}
PeriodicFunction::~PeriodicFunction() {
NotifyStop();
// Waits for thread_ to complete and clean up.
thread_.reset();
}
void PeriodicFunction::NotifyStop() {
if (!stop_thread_.HasBeenNotified()) {
stop_thread_.Notify();
}
}
void PeriodicFunction::RunLoop(const int64_t start) {
{
if (options_.startup_delay_micros > 0) {
const int64_t deadline = start + options_.startup_delay_micros;
options_.env->SleepForMicroseconds(deadline - start);
}
while (!stop_thread_.HasBeenNotified()) {
VLOG(3) << "Running function.";
const int64_t begin = options_.env->NowMicros();
function_();
// Take the max() here to guard against time going backwards which
// sometimes happens in multiproc machines.
const int64_t end =
std::max(static_cast<int64_t>(options_.env->NowMicros()), begin);
// The deadline is relative to when the last function started.
const int64_t deadline = begin + interval_micros_;
// We want to sleep until 'deadline'.
if (deadline > end) {
if (end > begin) {
VLOG(3) << "Reducing interval_micros from " << interval_micros_
<< " to " << (deadline - end);
}
options_.env->SleepForMicroseconds(deadline - end);
} else {
VLOG(3) << "Function took longer than interval_micros, so not sleeping";
}
}
}
}
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/periodic_function.h
+130
View File
@@ -0,0 +1,130 @@
/* Copyright 2016 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.
==============================================================================*/
// PeriodicFunction will periodically call the given function with a specified
// period in a background thread. After Start() returns, the thread is
// guaranteed to have started. The destruction of the class causes the
// background thread to be destroyed as well. Start() should not be called more
// than once.
//
// PeriodicFunction runs the function as soon as any previous run both is
// complete and was started more than "interval_micros" earlier. Thus, runs are
// both serialized, and normally have a period of "interval_micros" if no run
// exceeds the time.
//
// Note that, if the function takes longer than two interval_micross to finish,
// then PeriodicFunction will "skip" at least one call to the function. For
// instance, if the period is 50ms and the function starts runs at time 0 for
// 150ms, then the function will immediately start executing again at time 150,
// but there will be no function runs corresponding to times 50 or 100. This is
// especially important to remember when using an environment with a simulated
// clock: advancing simulated time atomically over N interval_micross will not
// cause the function to be called N times.
//
// This object is thread-safe.
//
// Example:
//
// class Foo {
// public:
// Foo() : periodic_function_([this]() { Bar(); },
// 1000 /* 1000us == 1ms*/) {
// }
//
// private:
// void Bar() { ... }
//
// PeriodicFunction periodic_function_;
// };
#ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_PERIODIC_FUNCTION_H_
#define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_PERIODIC_FUNCTION_H_
#include <functional>
#include <memory>
#include <string>
#include "absl/functional/any_invocable.h"
#include "absl/synchronization/notification.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace serving {
namespace internal {
class PeriodicFunctionTestAccess;
}
class PeriodicFunction {
public:
// Provides the ability to customize several aspects of the PeriodicFunction.
// Passed to constructor of PeriodicFunction.
struct Options {
Options() {}
// Any standard thread options, such as stack size, should
// be passed via "thread_options".
ThreadOptions thread_options;
// Specifies the thread name prefix (see the description in class
// Thread).
std::string thread_name_prefix = "periodic_function";
// The environment to use. Does not take ownership, but must remain alive
// for as long as the PeriodicFunction exists.
Env* env = Env::Default();
// Specifies the length of sleep before the first invocation of the
// function.
// This can be used for adding a random jitter to avoid synchronous behavior
// across multiple periodic functions.
int64_t startup_delay_micros = 0;
};
// Also starts the background thread which will be calling the function.
PeriodicFunction(absl::AnyInvocable<void()> function, int64_t interval_micros,
const Options& options = Options());
~PeriodicFunction();
private:
friend class internal::PeriodicFunctionTestAccess;
// Notifies the background thread to stop.
void NotifyStop();
// (Blocking.) Loops forever calling "function_" every "interval_micros_".
void RunLoop(int64_t start);
absl::AnyInvocable<void()> function_; // Actual client function
const int64_t interval_micros_; // Interval between calls.
const Options options_;
// Used to notify the thread to stop.
absl::Notification stop_thread_;
// Thread for running "function_"
std::unique_ptr<Thread> thread_ = nullptr;
PeriodicFunction(const PeriodicFunction&) = delete;
void operator=(const PeriodicFunction&) = delete;
};
} // namespace serving
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_PERIODIC_FUNCTION_H_
tensorflow/core/kernels/batching_util/periodic_function_test.cc
+226
View File
@@ -0,0 +1,226 @@
/* Copyright 2016 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/batching_util/periodic_function.h"
#include <memory>
#include <string>
#include "absl/synchronization/notification.h"
#include "tensorflow/core/kernels/batching_util/fake_clock_env.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace serving {
namespace internal {
class PeriodicFunctionTestAccess {
public:
explicit PeriodicFunctionTestAccess(PeriodicFunction* periodic_function)
: periodic_function_(periodic_function) {}
void NotifyStop() { periodic_function_->NotifyStop(); }
private:
PeriodicFunction* const periodic_function_;
};
} // namespace internal
namespace {
using test_util::FakeClockEnv;
void StopPeriodicFunction(PeriodicFunction* periodic_function,
FakeClockEnv* fake_clock_env,
const uint64_t pf_interval_micros) {
fake_clock_env->BlockUntilThreadsAsleep(1);
internal::PeriodicFunctionTestAccess(periodic_function).NotifyStop();
fake_clock_env->AdvanceByMicroseconds(pf_interval_micros);
}
TEST(PeriodicFunctionTest, ObeyInterval) {
const int64_t kPeriodMicros = 2;
const int kCalls = 10;
int actual_calls = 0;
{
FakeClockEnv fake_clock_env(Env::Default());
PeriodicFunction::Options options;
options.env = &fake_clock_env;
PeriodicFunction periodic_function([&actual_calls]() { ++actual_calls; },
kPeriodMicros, options);
for (int i = 0; i < kCalls; ++i) {
fake_clock_env.BlockUntilThreadsAsleep(1);
fake_clock_env.AdvanceByMicroseconds(kPeriodMicros);
}
StopPeriodicFunction(&periodic_function, &fake_clock_env, kPeriodMicros);
}
// The function gets called kCalls+1 times: once at time 0, once at time
// kPeriodMicros, once at time kPeriodMicros*2, up to once at time
// kPeriodMicros*kCalls.
ASSERT_EQ(actual_calls, kCalls + 1);
}
TEST(PeriodicFunctionTest, ObeyStartupDelay) {
const int64_t kDelayMicros = 10;
const int64_t kPeriodMicros = kDelayMicros / 10;
int actual_calls = 0;
{
PeriodicFunction::Options options;
options.startup_delay_micros = kDelayMicros;
FakeClockEnv fake_clock_env(Env::Default());
options.env = &fake_clock_env;
PeriodicFunction periodic_function([&actual_calls]() { ++actual_calls; },
kPeriodMicros, options);
// Wait for the thread to start up.
fake_clock_env.BlockUntilThreadsAsleep(1);
// Function shouldn't have been called yet.
EXPECT_EQ(0, actual_calls);
// Give enough time for startup delay to expire.
fake_clock_env.AdvanceByMicroseconds(kDelayMicros);
StopPeriodicFunction(&periodic_function, &fake_clock_env, kDelayMicros);
}
// Function should have been called at least once.
EXPECT_EQ(1, actual_calls);
}
// Test for race in calculating the first time the callback should fire.
TEST(PeriodicFunctionTest, StartupDelayRace) {
const int64_t kDelayMicros = 10;
const int64_t kPeriodMicros = kDelayMicros / 10;
mutex mu;
int counter = 0;
std::unique_ptr<absl::Notification> listener(new absl::Notification);
FakeClockEnv fake_clock_env(Env::Default());
PeriodicFunction::Options options;
options.env = &fake_clock_env;
options.startup_delay_micros = kDelayMicros;
PeriodicFunction periodic_function(
[&mu, &counter, &listener]() {
mutex_lock l(mu);
counter++;
listener->Notify();
},
kPeriodMicros, options);
fake_clock_env.BlockUntilThreadsAsleep(1);
fake_clock_env.AdvanceByMicroseconds(kDelayMicros);
listener->WaitForNotification();
{
mutex_lock l(mu);
EXPECT_EQ(1, counter);
// A notification can only be notified once.
listener = std::make_unique<absl::Notification>();
}
fake_clock_env.BlockUntilThreadsAsleep(1);
fake_clock_env.AdvanceByMicroseconds(kPeriodMicros);
listener->WaitForNotification();
{
mutex_lock l(mu);
EXPECT_EQ(2, counter);
}
StopPeriodicFunction(&periodic_function, &fake_clock_env, kPeriodMicros);
}
// If this test hangs forever, its probably a deadlock caused by setting the
// PeriodicFunction's interval to 0ms.
TEST(PeriodicFunctionTest, MinInterval) {
PeriodicFunction periodic_function(
[]() { Env::Default()->SleepForMicroseconds(20 * 1000); }, 0);
}
class PeriodicFunctionWithFakeClockEnvTest : public ::testing::Test {
protected:
const int64_t kPeriodMicros = 50;
PeriodicFunctionWithFakeClockEnvTest()
: fake_clock_env_(Env::Default()),
counter_(0),
pf_(
[this]() {
mutex_lock l(counter_mu_);
++counter_;
},
kPeriodMicros, GetPeriodicFunctionOptions()) {}
PeriodicFunction::Options GetPeriodicFunctionOptions() {
PeriodicFunction::Options options;
options.thread_name_prefix = "ignore";
options.env = &fake_clock_env_;
return options;
}
void SetUp() override {
// Note: counter_ gets initially incremented at time 0.
ASSERT_TRUE(AwaitCount(1));
}
void TearDown() override {
StopPeriodicFunction(&pf_, &fake_clock_env_, kPeriodMicros);
}
// The FakeClockEnv tests below advance simulated time and then expect the
// PeriodicFunction thread to run its function. This method helps the tests
// wait for the thread to execute, and then check the count matches the
// expectation.
bool AwaitCount(int expected_counter) {
fake_clock_env_.BlockUntilThreadsAsleep(1);
{
mutex_lock lock(counter_mu_);
return counter_ == expected_counter;
}
}
FakeClockEnv fake_clock_env_;
mutex counter_mu_;
int counter_;
PeriodicFunction pf_;
};
TEST_F(PeriodicFunctionWithFakeClockEnvTest, FasterThanRealTime) {
fake_clock_env_.AdvanceByMicroseconds(kPeriodMicros / 2);
for (int i = 2; i < 7; ++i) {
fake_clock_env_.AdvanceByMicroseconds(
kPeriodMicros); // advance past a tick
EXPECT_TRUE(AwaitCount(i));
}
}
TEST_F(PeriodicFunctionWithFakeClockEnvTest, SlowerThanRealTime) {
Env::Default()->SleepForMicroseconds(
125 * 1000); // wait for any unexpected breakage
EXPECT_TRUE(AwaitCount(1));
}
TEST(PeriodicFunctionDeathTest, BadInterval) {
EXPECT_DEBUG_DEATH(PeriodicFunction periodic_function([]() {}, -1),
".* should be >= 0");
EXPECT_DEBUG_DEATH(PeriodicFunction periodic_function(
[]() {}, -1, PeriodicFunction::Options()),
".* should be >= 0");
}
} // namespace
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/serial_device_batch_scheduler.h
+552
View File
@@ -0,0 +1,552 @@
/* Copyright 2017 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_SERIAL_DEVICE_BATCH_SCHEDULER_H_
#define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_SERIAL_DEVICE_BATCH_SCHEDULER_H_
#include <algorithm>
#include <functional>
#include <memory>
#include <random>
#include <unordered_map>
#include <vector>
#include "tensorflow/core/kernels/batching_util/batch_scheduler.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/platform/cpu_info.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace serving {
namespace internal {
template <typename TaskType>
class SDBSBatch;
template <typename TaskType>
class SDBSQueue;
} // namespace internal
// EXPERIMENTAL: API MAY BE SUBJECTED TO SUDDEN CHANGES.
//
// Shared batch scheduler designed for batches which are processed by a serial
// device (e.g. GPU, TPU). When batch processing involves a mix of
// parallelizable cpu work and non-parallelizable on-device work, overall
// latency can be minimized by producing batches at a (load dependent) rate
// which keeps the serial device uniformly busy.
//
// SerialDeviceBatchScheduler (SDBS) controls the batching rate by limiting the
// allowed number of concurrently processed batches. Too large a limit causes
// batches to pile up behind the serial device, adding to the overall batch
// latency. Too small a limit underutilizes the serial device and harms latency
// by forcing batches to wait longer to be processed. Feedback from the device
// (i.e. avg number of batches directly pending on the device) is used to set
// the correct limit.
//
// SDBS groups requests into per model batches which are processed when a batch
// processing thread becomes available. SDBS prioritizes batches primarily by
// age (i.e. the batch's oldest request) along with a configurable preference
// for scheduling larger batches first.
template <typename TaskType>
class SerialDeviceBatchScheduler : public std::enable_shared_from_this<
SerialDeviceBatchScheduler<TaskType>> {
public:
~SerialDeviceBatchScheduler();
struct Options {
// The name to use for the pool of batch threads.
std::string thread_pool_name = {"batch_threads"};
// Maximum number of batch processing threads.
int64_t num_batch_threads = port::NumSchedulableCPUs();
// Although batch selection is primarily based on age, this parameter
// specifies a preference for larger batches. A full batch will be
// scheduled before an older, nearly empty batch as long as the age gap is
// less than full_batch_scheduling_boost_micros. The optimal value for this
// parameter should be of order the batch processing latency, but must be
// chosen carefully, as too large a value will harm tail latency.
int64_t full_batch_scheduling_boost_micros = 0;
// The environment to use (typically only overridden by test code).
Env* env = Env::Default();
// Initial limit for number of batches being concurrently processed.
int64_t initial_in_flight_batches_limit = 3;
// Returns the current number of batches directly waiting to be processed
// by the serial device (i.e. GPU, TPU).
std::function<int64_t()> get_pending_on_serial_device;
// Desired average number of batches directly waiting to be processed by the
// serial device. Small numbers of O(1) should deliver the best latency.
double target_pending = 2;
// Number of batches between potential adjustments of
// in_flight_batches_limit. Larger numbers will reduce noise, but will be
// less responsive to sudden changes in workload.
int64_t batches_to_average_over = 1000;
};
// Ownership is shared between the caller of Create() and any queues created
// via AddQueue().
static absl::Status Create(
const Options& options,
std::shared_ptr<SerialDeviceBatchScheduler<TaskType>>* scheduler);
struct QueueOptions {
// Maximum size of each batch.
int max_batch_size = 1000;
// Maximum number of enqueued (i.e. non-scheduled) batches.
int max_enqueued_batches = 10;
};
using BatchProcessor = std::function<void(std::unique_ptr<Batch<TaskType>>)>;
// Adds queue (and its callback) to be managed by this scheduler.
absl::Status AddQueue(const QueueOptions& options,
BatchProcessor process_batch_callback,
std::unique_ptr<BatchScheduler<TaskType>>* queue);
double in_flight_batches_limit() {
mutex_lock l(mu_);
return in_flight_batches_limit_;
}
double recent_low_traffic_ratio() {
mutex_lock l(mu_);
return recent_low_traffic_ratio_;
}
private:
// access to AddBatch(), RemoveQueue(), env().
friend class internal::SDBSQueue<TaskType>;
explicit SerialDeviceBatchScheduler(const Options& options);
// Continuously retrieves and processes batches.
void ProcessBatches();
// Notifies scheduler of non-empty batch which is eligible for processing.
void AddBatch(const internal::SDBSBatch<TaskType>* batch);
// Removes queue from scheduler.
void RemoveQueue(const internal::SDBSQueue<TaskType>* queue);
Env* env() const { return options_.env; }
const Options options_;
// Collection of batches added by AddBatch. Owned by scheduler until they are
// released for processing.
std::vector<const internal::SDBSBatch<TaskType>*> batches_ TF_GUARDED_BY(mu_);
// Unowned queues and callbacks added by AddQueue.
std::unordered_map<const internal::SDBSQueue<TaskType>*, BatchProcessor>
queues_and_callbacks_ TF_GUARDED_BY(mu_);
// Responsible for running the batch processing callbacks.
std::unique_ptr<thread::ThreadPool> batch_thread_pool_;
// Limit on number of batches which can be concurrently processed.
int64_t in_flight_batches_limit_ TF_GUARDED_BY(mu_);
// Number of batch processing threads.
int64_t processing_threads_ TF_GUARDED_BY(mu_) = 0;
// Number of batches processed since the last in_flight_batches_limit_
// adjustment.
int64_t batch_count_ TF_GUARDED_BY(mu_) = 0;
// Number of times since the last in_flight_batches_limit_ adjustment when a
// processing thread was available but there were no batches to process.
int64_t no_batch_count_ TF_GUARDED_BY(mu_) = 0;
// Sum of batches pending on the serial device since the last
// in_flight_batches_limit_ adjustment.
int64_t pending_sum_ = 0;
// Sum of batch latencies since the last in_flight_batches_limit_ adjustment.
int64_t batch_latency_sum_ = 0;
// Average period between which two consecutive batches begin processing.
int64_t batch_period_micros_ = 0;
// Moving average tracking the fraction of recent in_flight_batches_limit_
// adjustments where the external traffic was not high enough to provide
// useful feedback for an adjustment.
double recent_low_traffic_ratio_ = 0;
mutex mu_;
SerialDeviceBatchScheduler(const SerialDeviceBatchScheduler&) = delete;
void operator=(const SerialDeviceBatchScheduler&) = delete;
};
//////////////////////////////////////////////////////////
// Implementation details follow. API users need not read.
namespace internal {
// Consolidates tasks into batches, passing them off to the
// SerialDeviceBatchScheduler for processing.
template <typename TaskType>
class SDBSQueue : public BatchScheduler<TaskType> {
public:
using QueueOptions =
typename SerialDeviceBatchScheduler<TaskType>::QueueOptions;
SDBSQueue(std::shared_ptr<SerialDeviceBatchScheduler<TaskType>> scheduler,
const QueueOptions& options);
~SDBSQueue() override;
// Adds task to current batch. Fails if the task size is larger than the batch
// size or if the current batch is full and this queue's number of outstanding
// batches is at its maximum.
absl::Status Schedule(std::unique_ptr<TaskType>* task) override;
// Number of tasks waiting to be scheduled.
size_t NumEnqueuedTasks() const override;
// Number of size 1 tasks which could currently be scheduled without failing.
size_t SchedulingCapacity() const override;
// Notifies queue that a batch is about to be scheduled; the queue should not
// place any more tasks in this batch.
void ReleaseBatch(const SDBSBatch<TaskType>* batch);
size_t max_task_size() const override { return options_.max_batch_size; }
private:
std::shared_ptr<SerialDeviceBatchScheduler<TaskType>> scheduler_;
const QueueOptions options_;
// Owned by scheduler_.
SDBSBatch<TaskType>* current_batch_ TF_GUARDED_BY(mu_) = nullptr;
int64_t num_enqueued_batches_ TF_GUARDED_BY(mu_) = 0;
int64_t num_enqueued_tasks_ TF_GUARDED_BY(mu_) = 0;
mutable mutex mu_;
SDBSQueue(const SDBSQueue&) = delete;
void operator=(const SDBSQueue&) = delete;
};
// Batch which remembers when and by whom it was created.
template <typename TaskType>
class SDBSBatch : public Batch<TaskType> {
public:
SDBSBatch(SDBSQueue<TaskType>* queue, int64_t creation_time_micros)
: queue_(queue), creation_time_micros_(creation_time_micros) {}
~SDBSBatch() override {}
SDBSQueue<TaskType>* queue() const { return queue_; }
int64_t creation_time_micros() const { return creation_time_micros_; }
private:
SDBSQueue<TaskType>* queue_;
const int64_t creation_time_micros_;
SDBSBatch(const SDBSBatch&) = delete;
void operator=(const SDBSBatch&) = delete;
};
} // namespace internal
// ---------------- SerialDeviceBatchScheduler ----------------
template <typename TaskType>
absl::Status SerialDeviceBatchScheduler<TaskType>::Create(
const Options& options,
std::shared_ptr<SerialDeviceBatchScheduler<TaskType>>* scheduler) {
if (options.num_batch_threads < 1) {
return errors::InvalidArgument("num_batch_threads must be positive; was ",
options.num_batch_threads);
}
if (options.initial_in_flight_batches_limit < 1) {
return errors::InvalidArgument(
"initial_in_flight_batches_limit must be positive; was ",
options.initial_in_flight_batches_limit);
}
if (options.initial_in_flight_batches_limit > options.num_batch_threads) {
return errors::InvalidArgument(
"initial_in_flight_batches_limit (",
options.initial_in_flight_batches_limit,
") should not be larger than num_batch_threads (",
options.num_batch_threads, ")");
}
if (options.full_batch_scheduling_boost_micros < 0) {
return errors::InvalidArgument(
"full_batch_scheduling_boost_micros can't be negative; was ",
options.full_batch_scheduling_boost_micros);
}
if (options.batches_to_average_over < 1) {
return errors::InvalidArgument(
"batches_to_average_over should be "
"greater than or equal to 1; was ",
options.batches_to_average_over);
}
if (options.target_pending <= 0) {
return errors::InvalidArgument(
"target_pending should be larger than zero; was ",
options.target_pending);
}
if (!options.get_pending_on_serial_device) {
return absl::InvalidArgumentError(
"get_pending_on_serial_device must be "
"specified");
}
scheduler->reset(new SerialDeviceBatchScheduler<TaskType>(options));
return absl::OkStatus();
}
template <typename TaskType>
SerialDeviceBatchScheduler<TaskType>::SerialDeviceBatchScheduler(
const Options& options)
: options_(options),
in_flight_batches_limit_(options.initial_in_flight_batches_limit),
processing_threads_(options.initial_in_flight_batches_limit) {
batch_thread_pool_ = std::make_unique<thread::ThreadPool>(
env(), options.thread_pool_name, options.num_batch_threads);
for (int i = 0; i < processing_threads_; i++) {
batch_thread_pool_->Schedule(
std::bind(&SerialDeviceBatchScheduler<TaskType>::ProcessBatches, this));
}
}
template <typename TaskType>
SerialDeviceBatchScheduler<TaskType>::~SerialDeviceBatchScheduler() {
// Signal processing threads to exit.
{
mutex_lock l(mu_);
processing_threads_ = 0;
}
// Hangs until all threads finish.
batch_thread_pool_.reset();
}
template <typename TaskType>
absl::Status SerialDeviceBatchScheduler<TaskType>::AddQueue(
const QueueOptions& options, BatchProcessor process_batch_callback,
std::unique_ptr<BatchScheduler<TaskType>>* queue) {
if (options.max_batch_size <= 0) {
return errors::InvalidArgument("max_batch_size must be positive; was ",
options.max_batch_size);
}
if (options.max_enqueued_batches <= 0) {
return errors::InvalidArgument(
"max_enqueued_batches must be positive; was ",
options.max_enqueued_batches);
}
internal::SDBSQueue<TaskType>* SDBS_queue_raw;
queue->reset(SDBS_queue_raw = new internal::SDBSQueue<TaskType>(
this->shared_from_this(), options));
mutex_lock l(mu_);
queues_and_callbacks_[SDBS_queue_raw] = process_batch_callback;
return absl::OkStatus();
}
template <typename TaskType>
void SerialDeviceBatchScheduler<TaskType>::AddBatch(
const internal::SDBSBatch<TaskType>* batch) {
mutex_lock l(mu_);
batches_.push_back(batch);
}
template <typename TaskType>
void SerialDeviceBatchScheduler<TaskType>::RemoveQueue(
const internal::SDBSQueue<TaskType>* queue) {
mutex_lock l(mu_);
queues_and_callbacks_.erase(queue);
}
template <typename TaskType>
void SerialDeviceBatchScheduler<TaskType>::ProcessBatches() {
const int64_t kIdleThreadSleepTimeMicros = 1000;
const double kMaxNoBatchRatio = .1;
const double kLowTrafficMovingAverageFactor = .1;
for (;;) {
mu_.lock();
if (processing_threads_ < 1 ||
processing_threads_ > in_flight_batches_limit_) {
processing_threads_--;
mu_.unlock();
break;
}
if (batches_.empty()) {
no_batch_count_++;
int64_t sleep_time = batch_period_micros_ ? batch_period_micros_
: kIdleThreadSleepTimeMicros;
mu_.unlock();
env()->SleepForMicroseconds(sleep_time);
continue;
}
auto best_it = batches_.begin();
double best_score =
(*best_it)->creation_time_micros() -
options_.full_batch_scheduling_boost_micros * (*best_it)->size() /
static_cast<double>((*best_it)->queue()->max_task_size());
for (auto it = batches_.begin() + 1; it != batches_.end(); it++) {
const double score =
(*it)->creation_time_micros() -
options_.full_batch_scheduling_boost_micros * (*it)->size() /
static_cast<double>((*it)->queue()->max_task_size());
if (score < best_score) {
best_score = score;
best_it = it;
}
}
const internal::SDBSBatch<TaskType>* batch = *best_it;
batches_.erase(best_it);
// Queue may destroy itself after ReleaseBatch is called.
batch->queue()->ReleaseBatch(batch);
auto callback = queues_and_callbacks_[batch->queue()];
mu_.unlock();
int64_t start_time = env()->NowMicros();
callback(std::unique_ptr<Batch<TaskType>>(
const_cast<internal::SDBSBatch<TaskType>*>(batch)));
int64_t end_time = env()->NowMicros();
mu_.lock();
batch_count_++;
batch_latency_sum_ += end_time - start_time;
pending_sum_ += options_.get_pending_on_serial_device();
if (batch_count_ == options_.batches_to_average_over) {
recent_low_traffic_ratio_ *= (1 - kLowTrafficMovingAverageFactor);
// Only adjust in_flight_batches_limit_ if external load is large enough
// to consistently provide batches. Otherwise we would (mistakenly) assume
// that the device is underutilized because in_flight_batches_limit_ is
// too small.
if (no_batch_count_ < kMaxNoBatchRatio * batch_count_) {
double avg_pending = pending_sum_ / static_cast<double>(batch_count_);
// Avg processing time / # of concurrent batches gives the avg period
// between which two consecutive batches begin processing. Used to set a
// reasonable sleep time for idle batch processing threads.
batch_period_micros_ =
batch_latency_sum_ / batch_count_ / in_flight_batches_limit_;
// When the processing pipeline is consistently busy, the average number
// of pending batches differs from in_flight_batches_limit_ by a
// load-dependent offset. Adjust in_flight_batches_limit_to maintain
// the desired target pending.
in_flight_batches_limit_ +=
std::round(options_.target_pending - avg_pending);
in_flight_batches_limit_ =
std::max(in_flight_batches_limit_, int64_t{1});
in_flight_batches_limit_ =
std::min(in_flight_batches_limit_, options_.num_batch_threads);
// Add extra processing threads if necessary.
if (processing_threads_ > 0 &&
processing_threads_ < in_flight_batches_limit_) {
int extra_threads = in_flight_batches_limit_ - processing_threads_;
for (int i = 0; i < extra_threads; i++) {
batch_thread_pool_->Schedule(std::bind(
&SerialDeviceBatchScheduler<TaskType>::ProcessBatches, this));
}
processing_threads_ = in_flight_batches_limit_;
}
} else {
recent_low_traffic_ratio_ += kLowTrafficMovingAverageFactor;
}
batch_count_ = 0;
no_batch_count_ = 0;
pending_sum_ = 0;
batch_latency_sum_ = 0;
}
mu_.unlock();
}
}
// ---------------- SDBSQueue ----------------
namespace internal {
template <typename TaskType>
SDBSQueue<TaskType>::SDBSQueue(
std::shared_ptr<SerialDeviceBatchScheduler<TaskType>> scheduler,
const QueueOptions& options)
: scheduler_(scheduler), options_(options) {}
template <typename TaskType>
SDBSQueue<TaskType>::~SDBSQueue() {
// Wait until last batch has been scheduled.
const int kSleepMicros = 1000;
for (;;) {
{
mutex_lock l(mu_);
if (num_enqueued_batches_ == 0) {
break;
}
}
scheduler_->env()->SleepForMicroseconds(kSleepMicros);
}
scheduler_->RemoveQueue(this);
}
template <typename TaskType>
absl::Status SDBSQueue<TaskType>::Schedule(std::unique_ptr<TaskType>* task) {
SDBSBatch<TaskType>* new_batch = nullptr;
size_t size = (*task)->size();
if (size > options_.max_batch_size) {
return errors::InvalidArgument("Task size ", size,
" is larger than maximum batch size ",
options_.max_batch_size);
}
{
mutex_lock l(mu_);
// Current batch is full, create another if allowed.
if (current_batch_ &&
current_batch_->size() + size > options_.max_batch_size) {
if (num_enqueued_batches_ >= options_.max_enqueued_batches) {
return absl::UnavailableError("The batch scheduling queue is full");
}
current_batch_->Close();
current_batch_ = nullptr;
}
if (!current_batch_) {
num_enqueued_batches_++;
current_batch_ = new_batch =
new SDBSBatch<TaskType>(this, scheduler_->env()->NowMicros());
}
current_batch_->AddTask(std::move(*task));
num_enqueued_tasks_++;
}
// AddBatch must be called outside of lock, since it may call ReleaseBatch.
if (new_batch != nullptr) scheduler_->AddBatch(new_batch);
return absl::OkStatus();
}
template <typename TaskType>
void SDBSQueue<TaskType>::ReleaseBatch(const SDBSBatch<TaskType>* batch) {
mutex_lock l(mu_);
num_enqueued_batches_--;
num_enqueued_tasks_ -= batch->num_tasks();
if (batch == current_batch_) {
current_batch_->Close();
current_batch_ = nullptr;
}
}
template <typename TaskType>
size_t SDBSQueue<TaskType>::NumEnqueuedTasks() const {
mutex_lock l(mu_);
return num_enqueued_tasks_;
}
template <typename TaskType>
size_t SDBSQueue<TaskType>::SchedulingCapacity() const {
mutex_lock l(mu_);
const int current_batch_capacity =
current_batch_ ? options_.max_batch_size - current_batch_->size() : 0;
const int spare_batches =
options_.max_enqueued_batches - num_enqueued_batches_;
return spare_batches * options_.max_batch_size + current_batch_capacity;
}
} // namespace internal
} // namespace serving
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_SERIAL_DEVICE_BATCH_SCHEDULER_H_
tensorflow/core/kernels/batching_util/serial_device_batch_scheduler_test.cc
+397
View File
@@ -0,0 +1,397 @@
/* Copyright 2017 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/batching_util/serial_device_batch_scheduler.h"
#include "absl/synchronization/notification.h"
#include "tensorflow/core/kernels/batching_util/fake_clock_env.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace serving {
namespace anonymous {
class FakeTask : public BatchTask {
public:
explicit FakeTask(size_t size) : size_(size) {}
~FakeTask() override = default;
size_t size() const override { return size_; }
private:
const size_t size_;
FakeTask(const FakeTask&) = delete;
void operator=(const FakeTask&) = delete;
};
// Creates a FakeTask of size 'task_size', and calls 'scheduler->Schedule()' on
// that task. Returns the resulting status.
absl::Status ScheduleTask(size_t task_size,
BatchScheduler<FakeTask>* scheduler) {
std::unique_ptr<FakeTask> task(new FakeTask(task_size));
absl::Status status = scheduler->Schedule(&task);
// Schedule() should have consumed 'task' iff it returned Status::OK.
CHECK_EQ(status.ok(), task == nullptr);
return status;
}
// Creates a thread that waits on 'start' and then advances the fake clock in
// 'env' in a loop until 'stop' is notified. Useful for allowing objects that
// use the clock to be destroyed.
std::unique_ptr<Thread> CreateFakeClockAdvancerThread(
test_util::FakeClockEnv* env, absl::Notification* start,
absl::Notification* stop) {
return std::unique_ptr<Thread>(Env::Default()->StartThread(
{}, "FakeClockAdvancerThread", [env, start, stop] {
start->WaitForNotification();
while (!stop->HasBeenNotified()) {
env->AdvanceByMicroseconds(10);
Env::Default()->SleepForMicroseconds(10);
}
}));
}
TEST(SerialDeviceBatchSchedulerTest, BadOptions) {
using Scheduler = SerialDeviceBatchScheduler<FakeTask>;
std::shared_ptr<Scheduler> scheduler;
Scheduler::Options default_options;
default_options.get_pending_on_serial_device = []() { return 0; };
Scheduler::Options options = default_options;
options.num_batch_threads = 0;
EXPECT_FALSE(Scheduler::Create(options, &scheduler).ok());
options = default_options;
options.initial_in_flight_batches_limit = 0;
EXPECT_FALSE(Scheduler::Create(options, &scheduler).ok());
options = default_options;
options.num_batch_threads = 5;
options.initial_in_flight_batches_limit = 8;
EXPECT_FALSE(Scheduler::Create(options, &scheduler).ok());
options = default_options;
options.batches_to_average_over = -5;
EXPECT_FALSE(Scheduler::Create(options, &scheduler).ok());
options = default_options;
options.target_pending = 0;
EXPECT_FALSE(Scheduler::Create(options, &scheduler).ok());
options = Scheduler::Options();
EXPECT_FALSE(Scheduler::Create(options, &scheduler).ok());
}
TEST(SerialDeviceBatchSchedulerTest, InFlightBatchesLimit) {
SerialDeviceBatchScheduler<FakeTask>::Options options;
options.num_batch_threads = 3;
options.initial_in_flight_batches_limit = 2;
options.batches_to_average_over = 1000;
options.get_pending_on_serial_device = []() { return 0; };
mutex mu;
int processed_batches = 0;
absl::Notification finish_processing;
auto queue_callback = [&mu, &processed_batches, &finish_processing](
std::unique_ptr<Batch<FakeTask>> batch) {
ASSERT_TRUE(batch->IsClosed());
EXPECT_GT(batch->num_tasks(), 0);
mu.lock();
int batch_num = ++processed_batches;
mu.unlock();
if (batch_num == 2) {
// Give third batch a chance to process if it's going to.
Env::Default()->SleepForMicroseconds(1000);
finish_processing.Notify();
}
if (batch_num == 3) {
ASSERT_TRUE(finish_processing.HasBeenNotified());
}
finish_processing.WaitForNotification();
};
std::shared_ptr<SerialDeviceBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(
SerialDeviceBatchScheduler<FakeTask>::Create(options, &scheduler));
std::unique_ptr<BatchScheduler<FakeTask>> queue1;
std::unique_ptr<BatchScheduler<FakeTask>> queue2;
std::unique_ptr<BatchScheduler<FakeTask>> queue3;
TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue1));
TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue2));
TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue3));
// Create 3 batches, only 2 should be processed concurrently.
TF_ASSERT_OK(ScheduleTask(100, queue1.get()));
TF_ASSERT_OK(ScheduleTask(100, queue2.get()));
TF_ASSERT_OK(ScheduleTask(100, queue3.get()));
}
TEST(SerialDeviceBatchSchedulerTest, PendingOnSerialDevice) {
mutex mu;
int pending;
SerialDeviceBatchScheduler<FakeTask>::Options options;
options.num_batch_threads = 3;
options.initial_in_flight_batches_limit = 1;
options.batches_to_average_over = 1;
options.target_pending = 3;
options.get_pending_on_serial_device = [&mu, &pending]() {
mutex_lock l(mu);
return pending;
};
std::shared_ptr<SerialDeviceBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(
SerialDeviceBatchScheduler<FakeTask>::Create(options, &scheduler));
int processed_batches = 0;
absl::Notification start_processing;
auto queue_callback = [&mu, &processed_batches, &start_processing, &pending,
&scheduler](std::unique_ptr<Batch<FakeTask>> batch) {
// Be careful with mutex mu to avoid potential deadlock with mutex mu_
// held in ProcessBatch() and in_flight_batches_limit().
int batch_num;
{
mutex_lock l(mu);
batch_num = ++processed_batches;
}
switch (batch_num) {
case 1:
start_processing.WaitForNotification();
{
mutex_lock l(mu);
pending = 3;
}
break;
case 2:
// Either low traffic or pending at target --> no adjustment.
CHECK_EQ(scheduler->in_flight_batches_limit(), 1);
{
mutex_lock l(mu);
pending = 1;
}
break;
case 3:
// Small pending --> 2 additional threads added.
CHECK_EQ(scheduler->in_flight_batches_limit(), 3);
{
mutex_lock l(mu);
pending = 3;
}
break;
default:
break;
}
};
std::unique_ptr<BatchScheduler<FakeTask>> queue;
TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue));
// Create 3 batches.
for (int i = 0; i < 3; i++) {
TF_ASSERT_OK(ScheduleTask(800, queue.get()));
}
start_processing.Notify();
}
TEST(SerialDeviceBatchSchedulerTest, FullBatchSchedulingBoostMicros) {
test_util::FakeClockEnv env(Env::Default());
absl::Notification start_teardown, stop_teardown;
std::unique_ptr<Thread> teardown_thread =
CreateFakeClockAdvancerThread(&env, &start_teardown, &stop_teardown);
{
SerialDeviceBatchScheduler<FakeTask>::Options options;
options.env = &env;
options.initial_in_flight_batches_limit = 1;
options.batches_to_average_over = 1000;
options.full_batch_scheduling_boost_micros = 10;
options.get_pending_on_serial_device = []() { return 0; };
mutex mu;
int processed_batches = 0;
auto queue_callback =
[&mu, &processed_batches](std::unique_ptr<Batch<FakeTask>> batch) {
ASSERT_TRUE(batch->IsClosed());
mutex_lock l(mu);
processed_batches++;
switch (processed_batches) {
case 1:
EXPECT_EQ(1000, batch->size());
break;
case 2:
EXPECT_EQ(100, batch->size());
break;
case 3:
EXPECT_EQ(80, batch->size());
break;
default:
EXPECT_TRUE(false) << "Should only have 3 batches";
}
};
std::shared_ptr<SerialDeviceBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(
SerialDeviceBatchScheduler<FakeTask>::Create(options, &scheduler));
// Make sure batch processing thread has gone to sleep.
Env::Default()->SleepForMicroseconds(1000);
SerialDeviceBatchScheduler<FakeTask>::QueueOptions queue_options;
std::unique_ptr<BatchScheduler<FakeTask>> queue1;
std::unique_ptr<BatchScheduler<FakeTask>> queue2;
std::unique_ptr<BatchScheduler<FakeTask>> queue3;
queue_options.max_batch_size = 1000;
TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue1));
queue_options.max_batch_size = 1000;
TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue2));
queue_options.max_batch_size = 100;
TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue3));
TF_ASSERT_OK(ScheduleTask(100, queue1.get()));
// First batch - creation time: 0, fullness: 0.1, sched score: -1
env.AdvanceByMicroseconds(3);
TF_ASSERT_OK(ScheduleTask(1000, queue2.get()));
// Second batch - creation time: 3, fullness: 1, sched score: -7
env.AdvanceByMicroseconds(5);
TF_ASSERT_OK(ScheduleTask(80, queue3.get()));
// Third batch - creation time: 8, fullness: .8, sched score: 0
// Release the batch processing thread.
env.AdvanceByMicroseconds(1000);
start_teardown.Notify();
}
stop_teardown.Notify();
}
TEST(SerialDeviceBatchSchedulerTest, DeleteQueue) {
SerialDeviceBatchScheduler<FakeTask>::Options options;
options.initial_in_flight_batches_limit = 1;
options.batches_to_average_over = 1000;
options.get_pending_on_serial_device = []() { return 0; };
mutex mu;
int processed_batches = 0;
absl::Notification finish_processing;
auto queue_callback = [&mu, &processed_batches, &finish_processing](
std::unique_ptr<Batch<FakeTask>> batch) {
ASSERT_TRUE(batch->IsClosed());
EXPECT_GT(batch->num_tasks(), 0);
finish_processing.WaitForNotification();
mu.lock();
processed_batches++;
mu.unlock();
};
std::shared_ptr<SerialDeviceBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(
SerialDeviceBatchScheduler<FakeTask>::Create(options, &scheduler));
std::unique_ptr<BatchScheduler<FakeTask>> queue;
TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue));
// Enqueue 2 tasks, should result in 2 batches.
for (int i = 0; i < 2; i++) {
TF_ASSERT_OK(ScheduleTask(800, queue.get()));
}
std::unique_ptr<Thread> queue_deleter(Env::Default()->StartThread(
{}, "QueueDeleterThread",
[&queue, &mu, &processed_batches, scheduler]() mutable {
// Delete queue, should be kept alive until empty.
queue.reset();
{
mutex_lock l(mu);
// queue may be destroyed before 2nd batch finishes processing.
EXPECT_GT(processed_batches, 0);
}
// Delete scheduler, should be kept alive until all batches processed.
scheduler.reset();
mutex_lock l(mu);
EXPECT_EQ(processed_batches, 2);
}));
// Release reference to scheduler, queue and callback above should keep alive.
scheduler.reset();
// Give queue_deleter thread time to delete queue.
Env::Default()->SleepForMicroseconds(1000);
finish_processing.Notify();
}
TEST(SerialDeviceBatchSchedulerTest, DeleteScheduler) {
SerialDeviceBatchScheduler<FakeTask>::Options options;
options.initial_in_flight_batches_limit = 1;
options.batches_to_average_over = 1000;
options.get_pending_on_serial_device = []() { return 0; };
mutex mu;
int processed_batches = 0;
absl::Notification start_processing;
absl::Notification finish_processing;
auto queue_callback =
[&mu, &processed_batches, &start_processing,
&finish_processing](std::unique_ptr<Batch<FakeTask>> batch) {
ASSERT_TRUE(batch->IsClosed());
EXPECT_GT(batch->num_tasks(), 0);
start_processing.WaitForNotification();
mutex_lock l(mu);
processed_batches++;
if (processed_batches == 2) {
finish_processing.Notify();
}
};
std::shared_ptr<SerialDeviceBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(
SerialDeviceBatchScheduler<FakeTask>::Create(options, &scheduler));
std::unique_ptr<BatchScheduler<FakeTask>> queue;
TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue));
// Enqueue 2 tasks, should result in 2 batches.
for (int i = 0; i < 2; i++) {
TF_ASSERT_OK(ScheduleTask(800, queue.get()));
}
// Delete scheduler, should be kept alive until queues are empty.
scheduler.reset();
start_processing.Notify();
finish_processing.WaitForNotification();
}
TEST(SerialDeviceBatchSchedulerTest, QueueCapacityInfo) {
SerialDeviceBatchScheduler<FakeTask>::Options options;
options.initial_in_flight_batches_limit = 1;
options.batches_to_average_over = 1000;
options.full_batch_scheduling_boost_micros = 1000;
options.get_pending_on_serial_device = []() { return 0; };
mutex mu;
int processed_batches = 0;
absl::Notification finish_processing;
auto queue_callback = [&mu, &processed_batches, &finish_processing](
std::unique_ptr<Batch<FakeTask>> batch) {
ASSERT_TRUE(batch->IsClosed());
EXPECT_GT(batch->num_tasks(), 0);
mu.lock();
int batch_num = ++processed_batches;
mu.unlock();
if (batch_num == 1) {
finish_processing.WaitForNotification();
}
};
std::shared_ptr<SerialDeviceBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(
SerialDeviceBatchScheduler<FakeTask>::Create(options, &scheduler));
std::unique_ptr<BatchScheduler<FakeTask>> queue1;
std::unique_ptr<BatchScheduler<FakeTask>> queue2;
TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue1));
TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue2));
// Blocker task, should schedule first.
TF_ASSERT_OK(ScheduleTask(800, queue1.get()));
TF_ASSERT_OK(ScheduleTask(100, queue2.get()));
EXPECT_EQ(queue2->NumEnqueuedTasks(), 1);
EXPECT_EQ(queue2->SchedulingCapacity(), 9 * 1000 + 900);
// Enqueue 2 more tasks, should fall in same batch.
TF_ASSERT_OK(ScheduleTask(100, queue2.get()));
TF_ASSERT_OK(ScheduleTask(200, queue2.get()));
EXPECT_EQ(queue2->NumEnqueuedTasks(), 3);
EXPECT_EQ(queue2->SchedulingCapacity(), 9 * 1000 + 600);
// Enqueue 1 more task, should create new batch.
TF_ASSERT_OK(ScheduleTask(700, queue2.get()));
EXPECT_EQ(queue2->NumEnqueuedTasks(), 4);
EXPECT_EQ(queue2->SchedulingCapacity(), 8 * 1000 + 300);
finish_processing.Notify();
}
} // namespace anonymous
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/shared_batch_scheduler.h
+2142
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/kernels/batching_util/shared_batch_scheduler_test.cc
+4697
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/kernels/batching_util/threadsafe_status.cc
+51
View File
@@ -0,0 +1,51 @@
/* 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/batching_util/threadsafe_status.h"
#include "absl/base/thread_annotations.h"
#include "absl/status/status.h"
#include "absl/synchronization/mutex.h"
#include "tensorflow/core/platform/mutex.h"
namespace tensorflow {
const absl::Status& ThreadSafeStatus::status() const& {
tf_shared_lock lock(mutex_);
return status_;
}
absl::Status ThreadSafeStatus::status() && {
tf_shared_lock lock(mutex_);
return std::move(status_);
}
void ThreadSafeStatus::Update(const absl::Status& new_status) {
if (new_status.ok()) {
return;
}
mutex_lock lock(mutex_);
status_.Update(new_status);
}
void ThreadSafeStatus::Update(absl::Status&& new_status) {
if (new_status.ok()) {
return;
}
mutex_lock lock(mutex_);
status_.Update(std::forward<absl::Status>(new_status));
}
} // namespace tensorflow
tensorflow/core/kernels/batching_util/threadsafe_status.h
+57
View File
@@ -0,0 +1,57 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_THREADSAFE_STATUS_H_
#define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_THREADSAFE_STATUS_H_
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/thread_annotations.h"
namespace tensorflow {
// Wrapper class to allow both lock-free construction and concurrent updates on
// a 'status'.
//
// Example Usage:
// std::thread threads[2];
// ThreadSafeStatus thread_safe_status;
// threads[0] = std::thread([&]() {
// status.Update(errors::Internal("internal error"));
// });
// threads[1] = std::thread([&]() {
// status.Update(errors::InvalidArgument("invalid argument"));
// });
// threads[0].Join();
// threads[1].Join();
//
// NOTE:
// When updated in a multi-threading setup, only the first error is retained.
class ThreadSafeStatus {
public:
const absl::Status& status() const& TF_LOCKS_EXCLUDED(mutex_);
absl::Status status() && TF_LOCKS_EXCLUDED(mutex_);
// Retains the first error status: replaces the current status with
// `new_status` if `new_status` is not OK and the previous status is OK.
void Update(const absl::Status& new_status) TF_LOCKS_EXCLUDED(mutex_);
void Update(absl::Status&& new_status) TF_LOCKS_EXCLUDED(mutex_);
private:
mutable mutex mutex_;
absl::Status status_ TF_GUARDED_BY(mutex_);
};
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_THREADSAFE_STATUS_H_
tensorflow/core/kernels/batching_util/threadsafe_status_test.cc
+51
View File
@@ -0,0 +1,51 @@
/* 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/batching_util/threadsafe_status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
namespace tensorflow {
namespace {
TEST(ThreadSafeStatus, DefaultOk) {
ThreadSafeStatus status;
TF_EXPECT_OK(status.status());
}
TEST(ThreadSafeStatus, Update) {
ThreadSafeStatus status;
TF_EXPECT_OK(status.status());
status.Update(absl::FailedPreconditionError("original error"));
EXPECT_EQ(status.status().code(), error::FAILED_PRECONDITION);
status.Update(absl::OkStatus());
EXPECT_EQ(status.status().code(), error::FAILED_PRECONDITION);
status.Update(absl::InternalError("new error"));
EXPECT_EQ(status.status().code(), error::FAILED_PRECONDITION);
}
TEST(ThreadSafeStatus, Move) {
ThreadSafeStatus status;
TF_EXPECT_OK(std::move(status).status());
}
} // namespace
} // namespace tensorflow
tensorflow/core/kernels/batching_util/warmup.cc
+82
View File
@@ -0,0 +1,82 @@
/* Copyright 2023 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/batching_util/warmup.h"
#include <memory>
#include <optional>
#include <utility>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/synchronization/mutex.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tsl/platform/logging.h"
namespace tensorflow {
namespace serving {
void WarmupStateRegistry::Handle::Release() {
if (!key_.has_value()) {
return;
}
DCHECK(registry_);
registry_->Unregister(*key_);
}
absl::StatusOr<WarmupStateRegistry::Handle> WarmupStateRegistry::Register(
const Key& model_key, std::unique_ptr<PerModelData> per_model_data) {
absl::MutexLock l(mu_);
VLOG(1) << "Registering model " << model_key.name << ":" << model_key.version
<< " to warm-up registry";
if (!states_.insert(std::pair(model_key, std::move(per_model_data))).second) {
return absl::AlreadyExistsError(
absl::StrCat("Model ", model_key.name, ":", model_key.version,
" already exists in the warm-up registry"));
}
return Handle(model_key, this);
}
void WarmupStateRegistry::Unregister(const Key& model_key) {
absl::MutexLock l(mu_);
VLOG(1) << "Unregistering model " << model_key.name << ":"
<< model_key.version << " from warm-up registry";
states_.erase(model_key);
}
const WarmupStateRegistry::PerModelData* WarmupStateRegistry::Lookup(
const Key& model_key) {
absl::ReaderMutexLock l(mu_);
return states_.contains(model_key) ? states_[model_key].get() : nullptr;
}
WarmupStateRegistry& GetGlobalWarmupStateRegistry() {
static auto* const registry = new WarmupStateRegistry;
return *registry;
}
bool ShouldWarmupAllBatchSizes(const OpKernelContext* c) {
auto metadata = c->session_metadata();
if (metadata == nullptr || metadata->name().empty()) {
return false;
}
serving::WarmupStateRegistry::Key key(metadata->name(), metadata->version());
auto per_model_data = serving::GetGlobalWarmupStateRegistry().Lookup(key);
return per_model_data && per_model_data->warmup_all_batch_sizes;
}
} // namespace serving
} // namespace tensorflow
tensorflow/core/kernels/batching_util/warmup.h
+132
View File
@@ -0,0 +1,132 @@
/* Copyright 2023 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_WARMUP_H_
#define TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_WARMUP_H_
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include "absl/container/flat_hash_map.h"
#include "absl/hash/hash.h"
#include "absl/status/statusor.h"
#include "absl/synchronization/mutex.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tsl/platform/logging.h"
namespace tensorflow {
namespace serving {
// Global registry for model's warm-up states. Before a model executes warm-up
// requests, it is registered here so that the runtime can distinguish demand
// requests vs. warm-up requests and apply warm-up specific optimizations.
class WarmupStateRegistry {
public:
struct Key {
std::string name;
int64_t version;
Key(std::string name, int64_t version)
: name(std::move(name)), version(version) {}
template <typename H>
friend H AbslHashValue(H state, const Key& key) {
return H::combine(std::move(state), key.name, key.version);
}
friend bool operator==(const Key& x, const Key& y) {
return x.name == y.name && x.version == y.version;
}
};
// Data stored per key.
struct PerModelData {
// If true, supported batch ops will execute the model on dummy batches
// for all `allowed_batch_sizes` of that batch op. This removes the
// need to issue separate warmup requests for each batch size.
bool warmup_all_batch_sizes = false;
};
// RAII handle for registered models.
class Handle {
public:
Handle() = default;
Handle(const Handle& other) = delete;
Handle& operator=(const Handle& other) = delete;
Handle(Handle&& other)
: key_(std::move(other.key_)), registry_(other.registry_) {
other.key_.reset();
}
Handle& operator=(Handle&& other) {
if (key_.has_value()) {
Release();
}
key_ = std::move(other.key_);
other.key_.reset();
registry_ = other.registry_;
return *this;
}
~Handle() { Release(); }
void Release();
private:
friend class WarmupStateRegistry;
// Can only be constructed by `WarmupStateRegistry::Register()`.
Handle(const Key& key, WarmupStateRegistry* registry)
: key_(key), registry_(registry) {
DCHECK(registry_);
}
std::optional<Key> key_;
WarmupStateRegistry* registry_ = nullptr;
};
// Registers the given model to be in a warm-up state and associates the given
// metadata with the model. Returns an RAII handle that unregisters the model
// at its destruction.
absl::StatusOr<Handle> Register(const Key& model_key,
std::unique_ptr<PerModelData> per_model_data);
// Return model data. A nullptr indicates the key was not present.
const PerModelData* Lookup(const Key& model_key);
private:
friend class Handle;
void Unregister(const Key& model_key);
absl::Mutex mu_;
// Map of model names/versions to miscellaneous data.
absl::flat_hash_map<Key, std::unique_ptr<PerModelData>> states_
ABSL_GUARDED_BY(&mu_);
};
WarmupStateRegistry& GetGlobalWarmupStateRegistry();
// Utility function that returns whether or not to warmup all batch sizes,
// based on the state of WarmupStateRegistry.
bool ShouldWarmupAllBatchSizes(const OpKernelContext* c);
} // namespace serving
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BATCHING_UTIL_WARMUP_H_
tensorflow/core/kernels/batchtospace_op.cc
+312
View File
@@ -0,0 +1,312 @@
/* Copyright 2016 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.
==============================================================================*/
// See docs in ../ops/array_ops.cc.
#define EIGEN_USE_THREADS
#include <memory>
#include <string>
#include <utility>
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/kernels/spacetobatch_functor.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/overflow.h"
namespace tensorflow {
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
template <typename Device, typename T>
static void BatchToSpaceOpCompute(OpKernelContext* context,
const Tensor& orig_input_tensor,
const Tensor& orig_block_shape,
const Tensor& orig_crops) {
const int input_dims = orig_input_tensor.dims();
OP_REQUIRES(context, TensorShapeUtils::IsVector(orig_block_shape.shape()),
absl::InvalidArgumentError(
absl::StrCat("block_shape rank should be 1 instead of ",
orig_block_shape.dims())));
const int block_dims = orig_block_shape.dim_size(0);
OP_REQUIRES(context, orig_input_tensor.dims() >= 1 + block_dims,
absl::InvalidArgumentError(
absl::StrCat("input rank should be >= ", 1 + block_dims,
" instead of ", orig_input_tensor.dims())));
OP_REQUIRES(context,
TensorShapeUtils::IsMatrix(orig_crops.shape()) &&
block_dims == orig_crops.dim_size(0) &&
2 == orig_crops.dim_size(1),
absl::InvalidArgumentError(absl::StrCat(
"crops should have shape [", block_dims, ", 2] instead of ",
orig_crops.shape().DebugString())));
// To avoid out-of-bounds access in the case that the block_shape and/or
// crops tensors are concurrently modified, we must copy the values.
absl::InlinedVector<int64_t, 4UL> block_shape;
absl::InlinedVector<int64_t, 8UL> crops;
internal::spacetobatch::SubtleMustCopyFlat(orig_block_shape, &block_shape);
internal::spacetobatch::SubtleMustCopyFlat(orig_crops, &crops);
// Determine the length of the prefix of block dims that can be combined
// into the batch dimension due to having no padding and block_shape=1.
int removed_prefix_block_dims = 0;
for (; removed_prefix_block_dims < block_dims; ++removed_prefix_block_dims) {
const int dim = removed_prefix_block_dims;
if (crops[2 * dim] != 0 || crops[2 * dim + 1] != 0 ||
block_shape[dim] != 1) {
break;
}
}
// Determine the length of the suffix of block dims that can be combined
// into the depth dimension due to having no padding and block_shape=1.
int removed_suffix_block_dims = 0;
for (; removed_suffix_block_dims < block_dims - removed_prefix_block_dims;
++removed_suffix_block_dims) {
const int dim = block_dims - 1 - removed_suffix_block_dims;
if (crops[2 * dim] != 0 || crops[2 * dim + 1] != 0 ||
block_shape[dim] != 1) {
break;
}
}
// Compute the product of the block_shape values.
int64_t block_shape_product = 1;
for (int block_dim = 0; block_dim < block_dims; ++block_dim) {
const int64_t block_shape_val = block_shape[block_dim];
OP_REQUIRES(context, block_shape_val >= 1,
absl::InvalidArgumentError(absl::StrCat(
"All values in block_shape must be positive, got ",
block_shape_val, " at index ", block_dim)));
block_shape_product =
MultiplyWithoutOverflow(block_shape_product, block_shape_val);
OP_REQUIRES(context, block_shape_product > 0,
absl::InvalidArgumentError(absl::StrCat(
"Product of block sizes must be positive, got ",
block_shape_product)));
}
const int64_t orig_input_batch_size = orig_input_tensor.dim_size(0);
OP_REQUIRES(context, orig_input_batch_size % block_shape_product == 0,
absl::InvalidArgumentError(
absl::StrCat("Input batch dimension (", orig_input_batch_size,
") is not divisible by product of block sizes (",
block_shape_product, ")")));
const int internal_block_dims =
block_dims - removed_prefix_block_dims - removed_suffix_block_dims;
OP_REQUIRES(context, internal_block_dims <= kMaxSpaceToBatchBlockDims,
absl::InvalidArgumentError(absl::StrCat(
"Maximum number of non-combined block dimensions is ",
internal_block_dims, " but must not exceed ",
kMaxSpaceToBatchBlockDims)));
if (internal_block_dims == 0) {
context->set_output(0, orig_input_tensor);
return;
}
// For the purpose of computing the result, the input will be treated as
// having this shape, of rank 2 + internal_block_dims.
TensorShape internal_input_shape;
// For the purpose of computing the result, the output will be treated as
// having this shape, of rank 2 + internal_block_dims.
TensorShape internal_output_shape;
// The actual output shape exposed to callers.
TensorShape external_output_shape;
OP_REQUIRES_OK(context, external_output_shape.AddDimWithStatus(
orig_input_batch_size / block_shape_product));
int64_t input_batch_size = orig_input_batch_size;
for (int block_dim = 0; block_dim < removed_prefix_block_dims; ++block_dim) {
const int64_t size = orig_input_tensor.dim_size(block_dim + 1);
input_batch_size *= size;
OP_REQUIRES_OK(context, external_output_shape.AddDimWithStatus(size));
}
OP_REQUIRES_OK(context,
internal_input_shape.AddDimWithStatus(input_batch_size));
OP_REQUIRES_OK(context, internal_output_shape.AddDimWithStatus(
input_batch_size / block_shape_product));
for (int block_dim = removed_prefix_block_dims;
block_dim < block_dims - removed_suffix_block_dims; ++block_dim) {
const int64_t crop_start = crops[2 * block_dim],
crop_end = crops[2 * block_dim + 1];
OP_REQUIRES(context, crop_start >= 0 && crop_end >= 0,
absl::InvalidArgumentError("Crops must be non-negative"));
const int64_t input_size = orig_input_tensor.dim_size(block_dim + 1);
const int64_t block_shape_value = block_shape[block_dim];
const int64_t input_size_block =
MultiplyWithoutOverflow(input_size, block_shape_value);
OP_REQUIRES(
context, input_size_block >= 0,
absl::InvalidArgumentError(absl::StrCat(
"Overflow when multiplying input_size and block_shape_value: ",
input_size, " * ", block_shape_value)));
const int64_t crop_sum = AddWithoutOverflow(crop_start, crop_end);
OP_REQUIRES(
context, crop_sum >= 0,
absl::InvalidArgumentError(absl::StrCat(
"Overflow when adding crops: ", crop_start, " + ", crop_end)));
const int64_t cropped_size = input_size_block - crop_sum;
OP_REQUIRES(context, cropped_size >= 0,
absl::InvalidArgumentError(
absl::StrCat("cropped_shape[", block_dim,
"]=", cropped_size, " must be non-negative")));
OP_REQUIRES_OK(context, internal_input_shape.AddDimWithStatus(input_size));
OP_REQUIRES_OK(context,
internal_output_shape.AddDimWithStatus(cropped_size));
OP_REQUIRES_OK(context,
external_output_shape.AddDimWithStatus(cropped_size));
}
int64_t depth = 1;
for (int dim = block_dims - removed_suffix_block_dims + 1; dim < input_dims;
++dim) {
const int64_t size = orig_input_tensor.dim_size(dim);
OP_REQUIRES_OK(context, external_output_shape.AddDimWithStatus(size));
depth *= size;
}
OP_REQUIRES_OK(context, internal_input_shape.AddDimWithStatus(depth));
OP_REQUIRES_OK(context, internal_output_shape.AddDimWithStatus(depth));
// Allocate output tensor.
Tensor* output_tensor = nullptr;
OP_REQUIRES_OK(context, context->allocate_output(0, external_output_shape,
&output_tensor));
const int64_t* internal_crops = &crops[2 * removed_prefix_block_dims];
const int64_t* internal_block_shape = &block_shape[removed_prefix_block_dims];
switch (internal_block_dims) {
#define TF_BATCHTOSPACE_BLOCK_DIMS_CASE(NUM_BLOCK_DIMS) \
case NUM_BLOCK_DIMS: { \
OP_REQUIRES_OK( \
context, \
(functor::SpaceToBatchFunctor<Device, T, NUM_BLOCK_DIMS, true>()( \
context->eigen_device<Device>(), \
output_tensor->shaped<T, NUM_BLOCK_DIMS + 2>( \
internal_output_shape.dim_sizes()), \
internal_block_shape, internal_crops, \
orig_input_tensor.shaped<T, NUM_BLOCK_DIMS + 2>( \
internal_input_shape.dim_sizes())))); \
} break; \
/**/
TF_SPACETOBATCH_FOR_EACH_NUM_BLOCK_DIMS(TF_BATCHTOSPACE_BLOCK_DIMS_CASE)
#undef TF_BATCHTOSPACE_BLOCK_DIMS_CASE
}
}
template <typename Device, typename T>
class BatchToSpaceNDOp : public OpKernel {
public:
explicit BatchToSpaceNDOp(OpKernelConstruction* context)
: OpKernel(context) {}
void Compute(OpKernelContext* context) override {
const Tensor& orig_input_tensor = context->input(0);
const Tensor& orig_block_shape = context->input(1);
const Tensor& orig_crops = context->input(2);
BatchToSpaceOpCompute<Device, T>(context, orig_input_tensor,
orig_block_shape, orig_crops);
}
};
template <typename Device, typename T>
class BatchToSpaceOp : public OpKernel {
public:
explicit BatchToSpaceOp(OpKernelConstruction* context) : OpKernel(context) {
OP_REQUIRES_OK(context, context->GetAttr("block_size", &block_size_));
OP_REQUIRES(context, block_size_ > 1,
absl::InvalidArgumentError(
absl::StrCat("Block size should be > 1: ", block_size_)));
block_shape_ = Tensor(tensorflow::DT_INT64, TensorShape({2}));
auto block_shape_vec = block_shape_.vec<int64_t>();
block_shape_vec(0) = block_size_;
block_shape_vec(1) = block_size_;
}
void Compute(OpKernelContext* context) override {
const Tensor& in0 = context->input(0);
const Tensor& in1 = context->input(1);
const int dims = in0.dims();
// Check on the input dimensions first.
// The input is presumed to be [batch, height, width, depth]
static const int kRequiredDims = 4;
OP_REQUIRES(
context, kRequiredDims == dims,
absl::InvalidArgumentError(absl::StrCat(
"Input rank should be: ", kRequiredDims, "instead of: ", dims)));
BatchToSpaceOpCompute<Device, T>(context, in0, block_shape_, in1);
}
private:
int block_size_;
Tensor block_shape_;
};
#define REGISTER(T) \
REGISTER_KERNEL_BUILDER(Name("BatchToSpaceND") \
.Device(DEVICE_CPU) \
.TypeConstraint<T>("T") \
.HostMemory("block_shape") \
.HostMemory("crops"), \
BatchToSpaceNDOp<CPUDevice, T>); \
REGISTER_KERNEL_BUILDER(Name("BatchToSpace") \
.Device(DEVICE_CPU) \
.TypeConstraint<T>("T") \
.HostMemory("crops"), \
BatchToSpaceOp<CPUDevice, T>);
TF_CALL_REAL_NUMBER_TYPES(REGISTER);
#undef REGISTER
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define REGISTER(T) \
REGISTER_KERNEL_BUILDER(Name("BatchToSpaceND") \
.Device(DEVICE_GPU) \
.TypeConstraint<T>("T") \
.HostMemory("block_shape") \
.HostMemory("crops"), \
BatchToSpaceNDOp<GPUDevice, T>); \
REGISTER_KERNEL_BUILDER(Name("BatchToSpace") \
.Device(DEVICE_GPU) \
.TypeConstraint<T>("T") \
.HostMemory("crops"), \
BatchToSpaceOp<GPUDevice, T>);
TF_CALL_GPU_NUMBER_TYPES(REGISTER);
#undef REGISTER
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // end namespace tensorflow
tensorflow/core/kernels/bcast_ops.cc
+185
View File
@@ -0,0 +1,185 @@
/* Copyright 2015 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/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/bcast.h"
namespace tensorflow {
// Given shapes of two tensors, computes the broadcast shape.
template <typename T>
class BCastArgsOp : public OpKernel {
public:
explicit BCastArgsOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
void Compute(OpKernelContext* ctx) override {
OP_REQUIRES(
ctx, ctx->num_inputs() == 2,
absl::UnimplementedError("Broadcast for n-ary operations (n > 2)"));
absl::InlinedVector<BCast::Vec, 4UL> shapes;
for (int i = 0; i < ctx->num_inputs(); ++i) {
const Tensor& in = ctx->input(i);
OP_REQUIRES(
ctx, TensorShapeUtils::IsVector(in.shape()),
absl::InvalidArgumentError(absl::StrCat(
"In[", i, "] must be a vector.", in.shape().DebugString())));
BCast::Vec vec;
for (int64_t i = 0; i < in.NumElements(); ++i) {
vec.push_back(in.vec<T>()(i));
}
shapes.push_back(vec);
}
BCast bcast(shapes[0], shapes[1]);
OP_REQUIRES(ctx, bcast.IsValid(),
absl::InvalidArgumentError(absl::StrCat(
"Incompatible shapes: [", absl::StrJoin(shapes[0], ","),
"] vs. [", absl::StrJoin(shapes[1], ","), "]")));
Output(ctx, 0, bcast.output_shape());
}
bool IsExpensive() override { return false; }
private:
void Output(OpKernelContext* ctx, int idx, const BCast::Vec& v) {
const int64_t len = v.size();
Tensor* o = nullptr;
OP_REQUIRES_OK(ctx, ctx->allocate_output(idx, TensorShape({len}), &o));
for (int64_t i = 0; i < len; ++i) {
o->flat<T>()(i) = static_cast<T>(v[i]);
}
}
BCastArgsOp(const BCastArgsOp&) = delete;
void operator=(const BCastArgsOp&) = delete;
};
// Given shapes of two tensors, computes the reduction indices for the
// gradient computation.
//
// TODO(zhifengc):
// 1. Adds support for n-ary (n >= 2).
template <typename T>
class BCastGradArgsOp : public OpKernel {
public:
explicit BCastGradArgsOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
void Compute(OpKernelContext* ctx) override {
OP_REQUIRES(
ctx, ctx->num_inputs() == 2,
absl::UnimplementedError("Broadcast for n-ary operations (n > 2)"));
absl::InlinedVector<BCast::Vec, 4UL> shapes;
for (int i = 0; i < ctx->num_inputs(); ++i) {
const Tensor& in = ctx->input(i);
OP_REQUIRES(
ctx, TensorShapeUtils::IsVector(in.shape()),
absl::InvalidArgumentError(absl::StrCat(
"In[", i, "] must be a vector.", in.shape().DebugString())));
BCast::Vec vec;
for (int64_t i = 0; i < in.NumElements(); ++i) {
vec.push_back(in.vec<T>()(i));
}
shapes.push_back(vec);
}
BCast bcast(shapes[0], shapes[1]);
OP_REQUIRES(ctx, bcast.IsValid(),
absl::InvalidArgumentError(absl::StrCat(
"Incompatible shapes: [", absl::StrJoin(shapes[0], ","),
"] vs. [", absl::StrJoin(shapes[1], ","), "]")));
Output(ctx, 0, bcast.grad_x_reduce_idx());
Output(ctx, 1, bcast.grad_y_reduce_idx());
}
bool IsExpensive() override { return false; }
private:
void Output(OpKernelContext* ctx, int idx, const BCast::Vec& v) {
const int64_t len = v.size();
Tensor* o = nullptr;
OP_REQUIRES_OK(ctx, ctx->allocate_output(idx, TensorShape({len}), &o));
for (int64_t i = 0; i < len; ++i) {
o->flat<T>()(i) = static_cast<T>(v[i]);
}
}
BCastGradArgsOp(const BCastGradArgsOp&) = delete;
void operator=(const BCastGradArgsOp&) = delete;
};
REGISTER_KERNEL_BUILDER(Name("BroadcastArgs")
.Device(DEVICE_CPU)
.TypeConstraint<int32_t>("T")
.HostMemory("s0")
.HostMemory("s1")
.HostMemory("r0"),
BCastArgsOp<int32_t>);
REGISTER_KERNEL_BUILDER(Name("BroadcastArgs")
.Device(DEVICE_CPU)
.TypeConstraint<int64_t>("T")
.HostMemory("s0")
.HostMemory("s1")
.HostMemory("r0"),
BCastArgsOp<int64_t>);
REGISTER_KERNEL_BUILDER(Name("BroadcastArgs")
.Device(DEVICE_DEFAULT)
.TypeConstraint<int32_t>("T")
.HostMemory("s0")
.HostMemory("s1")
.HostMemory("r0"),
BCastArgsOp<int32_t>);
REGISTER_KERNEL_BUILDER(Name("BroadcastArgs")
.Device(DEVICE_DEFAULT)
.TypeConstraint<int64_t>("T")
.HostMemory("s0")
.HostMemory("s1")
.HostMemory("r0"),
BCastArgsOp<int64_t>);
REGISTER_KERNEL_BUILDER(Name("BroadcastGradientArgs")
.Device(DEVICE_CPU)
.TypeConstraint<int32_t>("T")
.HostMemory("s0")
.HostMemory("s1")
.HostMemory("r0")
.HostMemory("r1"),
BCastGradArgsOp<int32_t>);
REGISTER_KERNEL_BUILDER(Name("BroadcastGradientArgs")
.Device(DEVICE_CPU)
.TypeConstraint<int64_t>("T")
.HostMemory("s0")
.HostMemory("s1")
.HostMemory("r0")
.HostMemory("r1"),
BCastGradArgsOp<int64_t>);
REGISTER_KERNEL_BUILDER(Name("BroadcastGradientArgs")
.Device(DEVICE_DEFAULT)
.TypeConstraint<int32_t>("T")
.HostMemory("s0")
.HostMemory("s1")
.HostMemory("r0")
.HostMemory("r1"),
BCastGradArgsOp<int32_t>);
REGISTER_KERNEL_BUILDER(Name("BroadcastGradientArgs")
.Device(DEVICE_DEFAULT)
.TypeConstraint<int64_t>("T")
.HostMemory("s0")
.HostMemory("s1")
.HostMemory("r0")
.HostMemory("r1"),
BCastGradArgsOp<int64_t>);
} // end namespace tensorflow
tensorflow/core/kernels/betainc_op.cc
+170
View File
@@ -0,0 +1,170 @@
/* Copyright 2016 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.
==============================================================================*/
// See docs in ../ops/nn_ops.cc.
#define EIGEN_USE_THREADS
// TODO(b/31098934): Figure out why this is necessary here but not in
// any other place, e.g., the cwise lgamma ops.
#define EIGEN_HAS_C99_MATH 1
#include "tensorflow/core/kernels/betainc_op.h"
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/numeric_op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/util/bcast.h"
namespace tensorflow {
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
template <typename Device, typename T>
class BetaincOp : public OpKernel {
public:
explicit BetaincOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
void Compute(OpKernelContext* ctx) override {
const Tensor& a = ctx->input(0);
const Tensor& b = ctx->input(1);
const Tensor& x = ctx->input(2);
const TensorShape& a_shape = a.shape();
const TensorShape& b_shape = b.shape();
const TensorShape& x_shape = x.shape();
if (a_shape.dims() > 0 && b_shape.dims() > 0) {
OP_REQUIRES(ctx, a_shape == b_shape,
absl::InvalidArgumentError(absl::StrCat(
"Shapes of a and b are inconsistent: ",
a_shape.DebugString(), " vs. ", b_shape.DebugString())));
}
if (a_shape.dims() > 0 && x_shape.dims() > 0) {
OP_REQUIRES(ctx, a_shape == x_shape,
absl::InvalidArgumentError(absl::StrCat(
"Shapes of a and x are inconsistent: ",
a_shape.DebugString(), " vs. ", x_shape.DebugString())));
}
if (b_shape.dims() > 0 && x_shape.dims() > 0) {
OP_REQUIRES(ctx, b_shape == x_shape,
absl::InvalidArgumentError(absl::StrCat(
"Shapes of b and x are inconsistent: ",
b_shape.DebugString(), " vs. ", x_shape.DebugString())));
}
TensorShape merged_shape(a_shape);
if (b_shape.dims() > 0) merged_shape = b_shape;
if (x_shape.dims() > 0) merged_shape = x_shape;
Tensor* output = nullptr;
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, merged_shape, &output));
if (a_shape == b_shape && a_shape == x_shape) {
functor::Betainc<Device, T, 1> functor;
functor(ctx->eigen_device<Device>(), a.flat<T>(), b.flat<T>(),
x.flat<T>(), output->flat<T>());
return;
}
auto merged_shape_vec = BCast::FromShape(merged_shape);
BCast a_shaper(BCast::FromShape(a_shape), merged_shape_vec);
BCast b_shaper(BCast::FromShape(b_shape), merged_shape_vec);
BCast x_shaper(BCast::FromShape(x_shape), merged_shape_vec);
int ndims = static_cast<int>(a_shaper.x_reshape().size());
switch (ndims) {
#define CASE(NDIM) \
case NDIM: { \
functor::Betainc<Device, T, NDIM> functor; \
auto a_value = a.shaped<T, NDIM>(a_shaper.x_reshape()); \
auto b_value = b.shaped<T, NDIM>(b_shaper.x_reshape()); \
auto x_value = x.shaped<T, NDIM>(x_shaper.x_reshape()); \
functor.BCast(ctx->eigen_device<Device>(), a_value, \
BCast::ToIndexArray<NDIM>(a_shaper.x_bcast()), b_value, \
BCast::ToIndexArray<NDIM>(b_shaper.x_bcast()), x_value, \
BCast::ToIndexArray<NDIM>(x_shaper.x_bcast()), \
output->shaped<T, NDIM>(a_shaper.y_reshape())); \
return; \
}
CASE(1);
CASE(2);
default: {
ctx->SetStatus(absl::InvalidArgumentError(
absl::StrCat("Broadcasting rank not supported: ", ndims)));
return;
}
}
}
};
#define REGISTER_KERNELS(type) \
REGISTER_KERNEL_BUILDER( \
Name("Betainc").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
BetaincOp<CPUDevice, type>);
REGISTER_KERNELS(float);
REGISTER_KERNELS(double);
#undef REGISTER_KERNELS
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
// Forward declarations of the functor specializations for GPU.
namespace functor {
#define DECLARE_GPU_SPEC_NDIM(T, NDIM) \
template <> \
void Betainc<GPUDevice, T, NDIM>::operator()( \
const GPUDevice& d, typename TTypes<T, NDIM>::ConstTensor a, \
typename TTypes<T, NDIM>::ConstTensor b, \
typename TTypes<T, NDIM>::ConstTensor x, \
typename TTypes<T, NDIM>::Tensor output); \
template <> \
void Betainc<GPUDevice, T, NDIM>::BCast( \
const GPUDevice& d, typename TTypes<T, NDIM>::ConstTensor a, \
const typename Eigen::array<Eigen::DenseIndex, NDIM>& bcast_a, \
typename TTypes<T, NDIM>::ConstTensor b, \
const typename Eigen::array<Eigen::DenseIndex, NDIM>& bcast_b, \
typename TTypes<T, NDIM>::ConstTensor x, \
const typename Eigen::array<Eigen::DenseIndex, NDIM>& bcast_x, \
typename TTypes<T, NDIM>::Tensor output); \
extern template struct Betainc<GPUDevice, T, NDIM>;
#define DECLARE_GPU_SPEC(T) \
DECLARE_GPU_SPEC_NDIM(T, 1) \
DECLARE_GPU_SPEC_NDIM(T, 2)
DECLARE_GPU_SPEC(float);
DECLARE_GPU_SPEC(double);
#undef DECLARE_GPU_SPEC
#undef DECLARE_GPU_SPEC_NDIM
} // namespace functor
// Registration of the GPU implementations.
#define REGISTER_GPU_KERNELS(type) \
REGISTER_KERNEL_BUILDER( \
Name("Betainc").Device(DEVICE_GPU).TypeConstraint<type>("T"), \
BetaincOp<GPUDevice, type>);
REGISTER_GPU_KERNELS(float);
REGISTER_GPU_KERNELS(double);
#undef REGISTER_GPU_KERNELS
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow
tensorflow/core/kernels/betainc_op.h
+51
View File
@@ -0,0 +1,51 @@
/* Copyright 2016 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BETAINC_OP_H_
#define TENSORFLOW_CORE_KERNELS_BETAINC_OP_H_
// Functor definition for BetaincOp, must be compilable by nvcc.
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/tensor_types.h"
namespace tensorflow {
namespace functor {
// Functor used by BetaincOp to do the computations.
template <typename Device, typename T, int NDIM>
struct Betainc {
void operator()(const Device& d, typename TTypes<T, NDIM>::ConstTensor a,
typename TTypes<T, NDIM>::ConstTensor b,
typename TTypes<T, NDIM>::ConstTensor x,
typename TTypes<T, NDIM>::Tensor output) {
output.device(d) = Eigen::betainc(a, b, x);
}
void BCast(const Device& d, typename TTypes<T, NDIM>::ConstTensor a,
const typename Eigen::array<Eigen::DenseIndex, NDIM>& bcast_a,
typename TTypes<T, NDIM>::ConstTensor b,
const typename Eigen::array<Eigen::DenseIndex, NDIM>& bcast_b,
typename TTypes<T, NDIM>::ConstTensor x,
const typename Eigen::array<Eigen::DenseIndex, NDIM>& bcast_x,
typename TTypes<T, NDIM>::Tensor output) {
output.device(d) = Eigen::betainc(
a.broadcast(bcast_a), b.broadcast(bcast_b), x.broadcast(bcast_x));
}
};
} // namespace functor
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BETAINC_OP_H_
tensorflow/core/kernels/betainc_op_gpu.cu.cc
+45
View File
@@ -0,0 +1,45 @@
/* Copyright 2016 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.
==============================================================================*/
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
#define EIGEN_USE_GPU
#include <stdio.h>
#include "tensorflow/core/kernels/betainc_op.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor_types.h"
namespace tensorflow {
typedef Eigen::GpuDevice GPUDevice;
// Definition of the GPU implementations declared in betainc_op.cc.
#define DEFINE_GPU_KERNELS_NDIM(T, NDIM) \
template struct functor::Betainc<GPUDevice, T, NDIM>;
#define DEFINE_GPU_KERNELS(T) \
DEFINE_GPU_KERNELS_NDIM(T, 1) \
DEFINE_GPU_KERNELS_NDIM(T, 2)
DEFINE_GPU_KERNELS(float);
DEFINE_GPU_KERNELS(double);
} // end namespace tensorflow
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
tensorflow/core/kernels/bias_op.cc
+558
View File
@@ -0,0 +1,558 @@
/* Copyright 2015 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.
==============================================================================*/
// See docs in ../ops/nn_ops.cc.
#define EIGEN_USE_THREADS
#include "tensorflow/core/kernels/bias_op.h"
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/bounds_check.h"
#include "tensorflow/core/framework/numeric_op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/kernels/redux_functor.h"
#include "tensorflow/core/profiler/lib/scoped_annotation.h"
#include "tensorflow/core/util/determinism.h"
#include "tensorflow/core/util/tensor_format.h"
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#include "xla/stream_executor/event_based_timer.h"
#include "tensorflow/core/kernels/bias_op_gpu.h"
#include "tensorflow/core/platform/stream_executor.h"
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
namespace tensorflow {
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
namespace {
void GetBiasValueDims(const Tensor& value_tensor, TensorFormat data_format,
int32_t* batch, int32_t* height, int32_t* width,
int32_t* depth, int32_t* channel) {
*batch = 1;
*height = 1;
*width = 1;
*depth = 1;
*channel = 1;
if (data_format == FORMAT_NHWC) {
int32_t channel_dim = value_tensor.dims() - 1;
*channel = static_cast<int32_t>(value_tensor.dim_size(channel_dim));
for (int32_t i = 0; i < channel_dim; i++) {
*batch *= static_cast<int32_t>(value_tensor.dim_size(i));
}
} else if (data_format == FORMAT_NCHW) {
*batch = static_cast<int32_t>(value_tensor.dim_size(0));
*channel = static_cast<int32_t>(value_tensor.dim_size(1));
*height = static_cast<int32_t>(value_tensor.dim_size(2));
if (value_tensor.dims() > 3) {
*width = static_cast<int32_t>(value_tensor.dim_size(3));
}
if (value_tensor.dims() > 4) {
*depth = static_cast<int32_t>(value_tensor.dim_size(4));
}
}
}
template <class T>
struct AccumulatorType {
typedef T type;
};
// float is faster on the CPU than half, and also more precise,
// so use float for the temporary accumulators.
template <>
struct AccumulatorType<Eigen::half> {
typedef float type;
};
} // namespace
template <typename Device, typename T>
class BiasOp : public BinaryOp<T> {
public:
explicit BiasOp(OpKernelConstruction* context) : BinaryOp<T>(context) {
std::string data_format;
if (context->GetAttr("data_format", &data_format).ok()) {
OP_REQUIRES(context, FormatFromString(data_format, &data_format_),
absl::InvalidArgumentError("Invalid data format"));
} else {
data_format_ = FORMAT_NHWC;
}
}
void Compute(OpKernelContext* context) override {
const Tensor& input = context->input(0);
const Tensor& bias = context->input(1);
OP_REQUIRES(context, TensorShapeUtils::IsMatrixOrHigher(input.shape()),
errors::InvalidArgument("Input tensor must be at least 2D: ",
input.shape()));
OP_REQUIRES(context, TensorShapeUtils::IsVector(bias.shape()),
errors::InvalidArgument("Biases must be 1D: ", bias.shape()));
// Added by intel_tf to support NCHW on CPU regardless of MKL used or not.
int channel_dim;
if (data_format_ == FORMAT_NCHW) {
channel_dim = 1; // NCHW always have channel dim in 1 (with 3, 4, 5
// dimensions data).
} else {
channel_dim = input.shape().dims() - 1; // End of code by intel_tf.
}
OP_REQUIRES(context,
bias.shape().dim_size(0) == input.shape().dim_size(channel_dim),
errors::InvalidArgument(
"Must provide as many biases as the last dimension "
"of the input tensor: ",
bias.shape(), " vs. ", input.shape()));
Tensor* output = nullptr;
OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
{0}, 0, input.shape(), &output));
if (input.NumElements() == 0) return;
functor::Bias<Device, T> functor;
const Device& d = context->eigen_device<Device>();
if (data_format_ == FORMAT_NCHW && input.shape().dims() > 2) {
functor(d, input.flat_inner_outer_dims<T, 2>(1),
bias.flat_outer_dims<T, 2>(),
output->flat_inner_outer_dims<T, 2>(1));
} else {
functor(d, input.flat<T>(), bias.vec<T>(), output->flat<T>());
}
}
private:
TensorFormat data_format_;
};
#define REGISTER_KERNEL(type) \
REGISTER_KERNEL_BUILDER( \
Name("BiasAdd").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
BiasOp<CPUDevice, type>); \
REGISTER_KERNEL_BUILDER( \
Name("BiasAddV1").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
BiasOp<CPUDevice, type>);
TF_CALL_NUMBER_TYPES(REGISTER_KERNEL);
#undef REGISTER_KERNEL
template <typename Device, typename T>
class BiasGradOp : public OpKernel {
public:
explicit BiasGradOp(OpKernelConstruction* context) : OpKernel(context) {
std::string data_format;
if (context->GetAttr("data_format", &data_format).ok()) {
OP_REQUIRES(context, FormatFromString(data_format, &data_format_),
absl::InvalidArgumentError("Invalid data format"));
} else {
data_format_ = FORMAT_NHWC;
}
}
void Compute(OpKernelContext* context) override {
const Tensor& output_backprop = context->input(0);
OP_REQUIRES(context,
TensorShapeUtils::IsMatrixOrHigher(output_backprop.shape()),
errors::InvalidArgument("Input tensor must be at least 2D: ",
output_backprop.shape()));
OP_REQUIRES(context,
FastBoundsCheck(output_backprop.NumElements(),
std::numeric_limits<int32_t>::max()),
absl::InvalidArgumentError(
"BiasGrad requires tensor size <= int32 max"));
int channel_dim;
if (data_format_ == FORMAT_NCHW) {
channel_dim = 1;
} else {
channel_dim = output_backprop.shape().dims() - 1;
}
Tensor* output = nullptr;
TensorShape output_shape{output_backprop.shape().dim_size(channel_dim)};
OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output));
if (output_backprop.NumElements() == 0) {
// Eigen often crashes by design on empty tensors, but setZero is safe
output->template flat<T>().setZero();
} else {
// Added by intel_tf to support NCHW on CPU regardless of MKL used or not.
using AccumT = typename AccumulatorType<T>::type;
if (data_format_ == FORMAT_NCHW) {
const functor::ReduceMiddleDimensions<
T, AccumT, T, Eigen::internal::scalar_sum_op<AccumT>,
Eigen::internal::SumReducer<T>>
redux;
auto flat_outer = output_backprop.flat_outer_dims<T, 3>();
redux(context->eigen_device<Device>(), flat_outer.dimensions(),
output_backprop, output, 1);
} else {
const functor::ReduceOuterDimensions<
T, AccumT, T, Eigen::internal::scalar_sum_op<AccumT>>
redux;
auto flat_inner = output_backprop.flat_inner_dims<T, 2>();
redux(context->eigen_device<Device>(), flat_inner.dimensions(),
output_backprop, output);
}
}
}
private:
TensorFormat data_format_;
};
// Registration of the GPU implementations.
#define REGISTER_KERNEL(type) \
REGISTER_KERNEL_BUILDER( \
Name("BiasAddGrad").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
BiasGradOp<CPUDevice, type>);
TF_CALL_NUMBER_TYPES(REGISTER_KERNEL);
#undef REGISTER_KERNEL
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
template <typename T>
class BiasOp<GPUDevice, T> : public BinaryOp<T> {
public:
typedef GPUDevice Device;
explicit BiasOp(OpKernelConstruction* context) : BinaryOp<T>(context) {
std::string data_format;
if (context->GetAttr("data_format", &data_format).ok()) {
OP_REQUIRES(context, FormatFromString(data_format, &data_format_),
absl::InvalidArgumentError("Invalid data format"));
} else {
data_format_ = FORMAT_NHWC;
}
}
void Compute(OpKernelContext* context) override {
const Tensor& input = context->input(0);
const Tensor& bias = context->input(1);
OP_REQUIRES(context, TensorShapeUtils::IsMatrixOrHigher(input.shape()),
absl::InvalidArgumentError(
absl::StrCat("Input tensor must be at least 2D: ",
input.shape().DebugString())));
OP_REQUIRES(context, TensorShapeUtils::IsVector(bias.shape()),
absl::InvalidArgumentError(absl::StrCat(
"Biases must be 1D: ", bias.shape().DebugString())));
int32_t batch, height, width, depth, channel;
GetBiasValueDims(input, data_format_, &batch, &height, &width, &depth,
&channel);
OP_REQUIRES(context, bias.shape().dim_size(0) == channel,
absl::InvalidArgumentError(absl::StrCat(
"Must provide as many biases as the channel dimension "
"of the input tensor: ",
bias.shape().DebugString(), " vs. ", channel, " in ",
input.shape().DebugString())));
Tensor* output = nullptr;
OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
{0}, 0, input.shape(), &output));
if (input.NumElements() > 0) {
OP_REQUIRES_OK(context, BiasGPU<T>::compute(
context->template eigen_device<Device>(),
input.flat<T>().data(), bias.flat<T>().data(),
output->flat<T>().data(), batch, width,
height, depth, channel, data_format_));
}
}
private:
TensorFormat data_format_;
};
// Registration of the GPU implementations.
#define REGISTER_GPU_KERNEL(type) \
REGISTER_KERNEL_BUILDER( \
Name("BiasAdd").Device(DEVICE_GPU).TypeConstraint<type>("T"), \
BiasOp<GPUDevice, type>); \
REGISTER_KERNEL_BUILDER( \
Name("BiasAddV1").Device(DEVICE_GPU).TypeConstraint<type>("T"), \
BiasOp<GPUDevice, type>);
TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNEL);
REGISTER_GPU_KERNEL(int32_t);
#undef REGISTER_GPU_KERNEL
struct BiasGradAutotuneGroup {
static std::string name() { return "BiasGrad"; }
};
class BiasAddGradGPUConfig {
public:
BiasAddGradGPUConfig() : mode_(BiasAddGradGPUMode::kReduction) {}
std::string ToString() const {
if (mode_ == BiasAddGradGPUMode::kNative) {
return "native CUDA kernel.";
}
if (mode_ == BiasAddGradGPUMode::kReduction) {
return "cub reduction kernel.";
}
return "unknown kernel.";
}
BiasAddGradGPUMode get_mode() const { return mode_; }
void set_mode(BiasAddGradGPUMode val) { mode_ = val; }
bool operator==(const BiasAddGradGPUConfig& other) const {
return this->mode_ == other.get_mode();
}
bool operator!=(const BiasAddGradGPUConfig& other) const {
return !(*this == other);
}
private:
BiasAddGradGPUMode mode_;
};
// Encapsulate all the shape information that is used in bias add grad
// operations.
class BiasAddParams {
public:
// We use a list to maintain both the shape value and the order (data format).
using SpatialArray = absl::InlinedVector<int64_t, 4UL>;
BiasAddParams(const SpatialArray& in_shape, TensorFormat data_format,
DataType dtype, int device_id)
: in_shape_(in_shape),
data_format_(data_format),
dtype_(dtype),
device_id_(device_id) {
for (int64_t val : in_shape_) {
hash_code_ = Hash64Combine(hash_code_, val);
}
hash_code_ = Hash64Combine(hash_code_, data_format);
hash_code_ = Hash64Combine(hash_code_, dtype);
hash_code_ = Hash64Combine(hash_code_, device_id);
}
bool operator==(const BiasAddParams& other) const {
return this->get_data_as_tuple() == other.get_data_as_tuple();
}
bool operator!=(const BiasAddParams& other) const {
return !(*this == other);
}
uint64_t hash() const { return hash_code_; }
std::string ToString() const {
// clang-format off
return strings::StrCat(
"(", absl::StrJoin(in_shape_, ", "), "), ",
data_format_, ", ", dtype_, ", ", device_id_);
// clang-format on
}
protected:
using ParamsDataType = std::tuple<SpatialArray, TensorFormat, DataType, int>;
ParamsDataType get_data_as_tuple() const {
return std::make_tuple(in_shape_, data_format_, dtype_, device_id_);
}
uint64_t hash_code_ = 0;
private:
SpatialArray in_shape_;
TensorFormat data_format_;
DataType dtype_;
int device_id_;
};
typedef AutotuneSingleton<BiasGradAutotuneGroup, BiasAddParams,
BiasAddGradGPUConfig>
AutotuneBiasGrad;
template <typename T>
class BiasGradOp<GPUDevice, T> : public OpKernel {
public:
typedef GPUDevice Device;
explicit BiasGradOp(OpKernelConstruction* context) : OpKernel(context) {
std::string data_format;
if (context->GetAttr("data_format", &data_format).ok()) {
OP_REQUIRES(context, FormatFromString(data_format, &data_format_),
absl::InvalidArgumentError("Invalid data format"));
} else {
data_format_ = FORMAT_NCHW;
}
}
void ComputeWithCustomKernel(OpKernelContext* context,
const Tensor& output_backprop, int32_t batch,
int32_t width, int32_t height, int32_t depth,
int32_t channel, Tensor* output) {
OP_REQUIRES_OK(context, BiasGradGPU<T>::compute(
context->template eigen_device<Device>(),
output_backprop.template flat<T>().data(),
output->flat<T>().data(), batch, width, height,
depth, channel, data_format_));
}
void ComputeWithReduceSum(OpKernelContext* context,
const Tensor& output_backprop, int32_t batch,
int32_t width, int32_t height, int32_t depth,
int32_t channel, Tensor* output) {
if (data_format_ == FORMAT_NCHW) {
int32_t row_count = batch * channel;
int32_t col_count = height * width * depth;
Tensor temp_grad_outputs;
// For 'NCHW' format, we perform reduction twice: first HW, then N.
TensorShape temp_grad_output_shape{row_count, col_count};
OP_REQUIRES_OK(context, context->allocate_temp(DataTypeToEnum<T>::value,
temp_grad_output_shape,
&temp_grad_outputs));
BiasGradGPU<T>::DoRowReduction(
context, temp_grad_outputs.flat<T>().data(),
output_backprop.template flat<T>().data(), row_count, col_count);
row_count = batch;
col_count = channel;
BiasGradGPU<T>::DoColReduction(context, output->flat<T>().data(),
temp_grad_outputs.flat<T>().data(),
row_count, col_count);
} else {
// For 'NHWC', we simply apply reduction once on NHW.
int32_t row_count = batch * height * width * depth;
int32_t col_count = channel;
BiasGradGPU<T>::DoColReduction(
context, const_cast<T*>(output->flat<T>().data()),
reinterpret_cast<const T*>(output_backprop.template flat<T>().data()),
row_count, col_count);
}
}
void Compute(OpKernelContext* context) override {
const Tensor& output_backprop = context->input(0);
OP_REQUIRES(context,
TensorShapeUtils::IsMatrixOrHigher(output_backprop.shape()),
absl::InvalidArgumentError(
absl::StrCat("Input tensor must be at least 2D: ",
output_backprop.shape().DebugString())));
int32_t batch, height, width, depth, channel;
GetBiasValueDims(output_backprop, data_format_, &batch, &height, &width,
&depth, &channel);
Tensor* output = nullptr;
TensorShape output_shape{channel};
OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output));
if (channel == 0) return;
auto* stream = context->op_device_context()->stream();
OP_REQUIRES(context, stream,
absl::InternalError("No GPU stream available."));
stream_executor::DeviceAddressBase output_ptr(
output->flat<T>().data(), output->NumElements() * sizeof(T));
OP_REQUIRES_OK(context, stream->MemZero(&output_ptr,
output->NumElements() * sizeof(T)));
if (output_backprop.NumElements() <= 0) return;
if (OpDeterminismRequired()) {
// ComputeWithReduceSum is the only deterministic algorithm.
ComputeWithReduceSum(context, output_backprop, batch, width, height,
depth, channel, output);
return;
}
int device_id = stream->parent()->device_ordinal();
DataType dtype = output_backprop.dtype();
BiasAddParams bias_parameters = {
{batch, height * width * depth, channel},
data_format_,
dtype,
device_id,
};
// Autotune two algorithm: customized
BiasAddGradGPUConfig algo_config;
if (!AutotuneBiasGrad::GetInstance()->Find(bias_parameters, &algo_config)) {
tsl::profiler::ScopedAnnotation trace("bias_grad_autotuning");
BiasGradGPUProfileResult best_result;
// Initialize the timer.
absl::StatusOr<std::unique_ptr<stream_executor::EventBasedTimer>> timer =
stream->CreateEventBasedTimer(false);
OP_REQUIRES_OK(context, timer.status());
ComputeWithCustomKernel(context, output_backprop, batch, width, height,
depth, channel, output);
absl::StatusOr<absl::Duration> bias_duration =
timer.value()->GetElapsedDuration();
OP_REQUIRES_OK(context, bias_duration.status());
int64_t elapsed_microseconds = absl::ToInt64Microseconds(*bias_duration);
VLOG(1) << "BiasAddGrad " << bias_parameters.ToString()
<< " Native algo latency: " << elapsed_microseconds << "us";
if (elapsed_microseconds < best_result.elapsed_time()) {
best_result.set_algorithm(BiasAddGradGPUMode::kNative);
best_result.set_elapsed_time(elapsed_microseconds);
}
// Try reduction and profile.
absl::StatusOr<std::unique_ptr<stream_executor::EventBasedTimer>>
reduction_timer = stream->CreateEventBasedTimer(false);
OP_REQUIRES_OK(context, reduction_timer.status());
ComputeWithReduceSum(context, output_backprop, batch, width, height,
depth, channel, output);
absl::StatusOr<absl::Duration> reduction_duration =
reduction_timer.value()->GetElapsedDuration();
OP_REQUIRES_OK(context, reduction_duration.status());
elapsed_microseconds += absl::ToInt64Microseconds(*reduction_duration);
VLOG(1) << "BiasAddGrad " << bias_parameters.ToString()
<< " Reduction algo latency: " << elapsed_microseconds;
if (elapsed_microseconds < best_result.elapsed_time()) {
best_result.set_algorithm(BiasAddGradGPUMode::kReduction);
best_result.set_elapsed_time(elapsed_microseconds);
}
algo_config.set_mode(best_result.algorithm());
AutotuneBiasGrad::GetInstance()->Insert(bias_parameters, algo_config);
// Results are already available during autotune, so no need to continue.
return;
}
// Choose the best algorithm based on autotune results.
if (algo_config.get_mode() == BiasAddGradGPUMode::kReduction) {
ComputeWithReduceSum(context, output_backprop, batch, width, height,
depth, channel, output);
} else {
// Default to the customized kernel.
ComputeWithCustomKernel(context, output_backprop, batch, width, height,
depth, channel, output);
}
}
private:
TensorFormat data_format_;
};
// Registration of the GPU implementations.
#define REGISTER_GPU_KERNEL(type) \
REGISTER_KERNEL_BUILDER( \
Name("BiasAddGrad").Device(DEVICE_GPU).TypeConstraint<type>("T"), \
BiasGradOp<GPUDevice, type>);
TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNEL);
#undef REGISTER_GPU_KERNEL
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow
tensorflow/core/kernels/bias_op.h
+60
View File
@@ -0,0 +1,60 @@
/* Copyright 2015 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BIAS_OP_H_
#define TENSORFLOW_CORE_KERNELS_BIAS_OP_H_
// Functor definition for BiasOp, must be compilable by nvcc.
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/tensor_types.h"
namespace tensorflow {
namespace functor {
// Functor used by BiasOp to do the computations.
template <typename Device, typename T>
struct Bias {
// Add "bias" to "input", repeating "bias".
void operator()(const Device& d, typename TTypes<T>::ConstFlat input,
typename TTypes<T>::ConstVec bias,
typename TTypes<T>::Flat output) {
const Eigen::Index rest_size = input.size() / bias.dimension(0);
Eigen::DSizes<Eigen::Index, 1> bcast(rest_size);
MaybeWith32BitIndexing<Device>(
[&](auto input32, auto bias32, auto output32, const auto& bcast32) {
output32.device(d) = input32 + bias32.broadcast(bcast32);
},
input, bias, output, bcast);
}
// NCHW layout, repeating on the first dimension, broadcasting on the last
// dimension.
void operator()(const Device& d, typename TTypes<T>::ConstMatrix input,
typename TTypes<T>::ConstMatrix bias1, // shape [C, 1].
typename TTypes<T>::Matrix output) {
const Eigen::Index rest_size = input.dimension(0) / bias1.dimension(0);
Eigen::DSizes<Eigen::Index, 2> bcast(rest_size, input.dimension(1));
MaybeWith32BitIndexing<Device>(
[&](auto input32, auto bias32, auto output32, const auto& bcast32) {
output32.device(d) = input32 + bias32.broadcast(bcast32);
},
input, bias1, output, bcast);
}
};
} // namespace functor
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BIAS_OP_H_
tensorflow/core/kernels/bias_op_gpu.cu.cc
+339
View File
@@ -0,0 +1,339 @@
/* Copyright 2015 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.
==============================================================================*/
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define EIGEN_USE_GPU
#include "tensorflow/core/kernels/bias_op_gpu.h"
#include <algorithm>
#include <limits>
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/kernels/bias_op.h"
#include "tensorflow/core/kernels/reduction_gpu_kernels.cu.h"
#include "tensorflow/core/kernels/reduction_ops_common.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/gpu_kernel_helper.h"
#include "tensorflow/core/util/overflow.h"
namespace tensorflow {
typedef Eigen::GpuDevice GPUDevice;
// There are no native fp16 atomics (we simulate them using 32-bit atomics),
// so fp16 sums are done in fp32 internally. (We don't have a lot of shared
// memory traffic; BiasGradNCHW_SharedAtomics in particular works almost
// entirely on a local variable.)
template <class T>
struct AccumulatorType {
typedef T type;
};
template <>
struct AccumulatorType<Eigen::half> {
typedef float type;
};
template <>
struct AccumulatorType<Eigen::bfloat16> {
typedef float type;
};
// Definition of the GPU implementations declared in bias_op.cc.
template <typename T>
__global__ void BiasNHWCKernel(int32_t nthreads, const T* __restrict__ input,
const T* __restrict__ bias,
T* __restrict__ output, int32_t bias_size) {
GPU_1D_KERNEL_LOOP(index, nthreads) {
int32_t bias_offset = index % bias_size;
output[index] = ldg(input + index) + ldg(bias + bias_offset);
}
}
template <typename T>
__global__ void BiasNCHWKernel(int32_t nthreads, const T* __restrict__ input,
const T* __restrict__ bias,
T* __restrict__ output, int32_t bias_size,
int32_t image_size) {
GPU_1D_KERNEL_LOOP(index, nthreads) {
int32_t index2 = index / image_size;
int32_t bias_offset = index2 % bias_size;
output[index] = ldg(input + index) + ldg(bias + bias_offset);
}
}
// Add "bias" to "input", broadcasting it on all dimensions but the bias
// dimension.
template <typename T>
absl::Status BiasGPU<T>::compute(const GPUDevice& d, const T* input,
const T* bias, T* output, int32_t batch,
int32_t height, int32_t width, int depth,
int32_t channel, TensorFormat data_format) {
const int32_t bias_size = channel;
int64_t image_size_64 =
MultiplyWithoutOverflow(MultiplyWithoutOverflow(height, width), depth);
int64_t total_count_64 = MultiplyWithoutOverflow(
MultiplyWithoutOverflow(batch, bias_size), image_size_64);
if (total_count_64 < 0 ||
total_count_64 > std::numeric_limits<int32_t>::max() ||
image_size_64 < 0 ||
image_size_64 > std::numeric_limits<int32_t>::max()) {
return absl::InternalError("BiasGPU: dimensions exceed int32 bounds");
}
const int32_t image_size = image_size_64;
const int32_t total_count = total_count_64;
if (total_count == 0) {
return absl::OkStatus();
}
if (data_format == FORMAT_NHWC) {
GpuLaunchConfig config =
GetGpuLaunchConfig(total_count, d, BiasNHWCKernel<T>, 0, 0);
TF_CHECK_OK(GpuLaunchKernel(BiasNHWCKernel<T>, config.block_count,
config.thread_per_block, 0, d.stream(),
config.virtual_thread_count, input, bias,
output, bias_size));
} else {
GpuLaunchConfig config =
GetGpuLaunchConfig(total_count, d, BiasNCHWKernel<T>, 0, 0);
TF_CHECK_OK(GpuLaunchKernel(BiasNCHWKernel<T>, config.block_count,
config.thread_per_block, 0, d.stream(),
config.virtual_thread_count, input, bias,
output, bias_size, image_size));
}
return absl::OkStatus();
}
// A naive implementation that is functional on all cases.
template <typename T>
__global__ void BiasGradNHWC_Naive(int32_t nthreads,
const T* __restrict__ output_backprop,
T* __restrict__ bias_backprop,
int32_t bias_size) {
GPU_1D_KERNEL_LOOP(index, nthreads) {
int32_t bias_offset = index % bias_size;
GpuAtomicAdd(bias_backprop + bias_offset, ldg(output_backprop + index));
}
}
// A naive implementation that is functional on all cases.
template <typename T>
__global__ void BiasGradNCHW_Naive(int32_t nthreads,
const T* __restrict__ output_backprop,
T* __restrict__ bias_backprop,
int32_t bias_size, int32_t image_size) {
GPU_1D_KERNEL_LOOP(index, nthreads) {
int32_t index2 = index / image_size;
int32_t bias_offset = index2 % bias_size;
GpuAtomicAdd(bias_backprop + bias_offset, ldg(output_backprop + index));
}
}
template <typename T>
__global__ void BiasGradNHWC_SharedAtomics(
int32_t nthreads, const T* __restrict__ output_backprop,
T* __restrict__ bias_backprop, int32_t bias_size) {
typedef typename AccumulatorType<T>::type AccT;
GPU_DYNAMIC_SHARED_MEM_DECL(8, char, s_buf);
AccT* s_data = reinterpret_cast<AccT*>(s_buf);
for (int32_t index = threadIdx.x; index < bias_size; index += blockDim.x) {
s_data[index] = AccT(0);
}
__syncthreads();
for (int32_t index = blockIdx.x * blockDim.x + threadIdx.x; index < nthreads;
index += blockDim.x * gridDim.x) {
int32_t bias_offset = index % bias_size;
GpuAtomicAdd(s_data + bias_offset, AccT(ldg(output_backprop + index)));
}
__syncthreads();
for (int32_t index = threadIdx.x; index < bias_size; index += blockDim.x) {
GpuAtomicAdd(bias_backprop + index, T(s_data[index]));
}
}
template <typename T>
__global__ void BiasGradNCHW_SharedAtomics(
const T* __restrict__ output_backprop, T* __restrict__ bias_backprop,
int32_t batch, int32_t bias_size, int32_t image_size, int group_size) {
// Initialize the shared memory.
typedef typename AccumulatorType<T>::type AccT;
const int32_t kSDataSize = 32;
__shared__ AccT s_data[kSDataSize];
for (int32_t index = threadIdx.x; index < kSDataSize; index += blockDim.x) {
s_data[index] = AccT(0);
}
__syncthreads();
// Accumulate all the values within this thread. They all have the same bias
// index.
int32_t bias_index = blockIdx.x % bias_size;
int32_t group_index = blockIdx.x / bias_size;
int32_t total_count = batch * image_size;
AccT sum(0);
for (int32_t index = group_index * blockDim.x + threadIdx.x;
index < total_count; index += blockDim.x * group_size) {
int32_t image_offset = index % image_size;
int32_t batch = index / image_size;
T val = ldg(output_backprop +
(batch * bias_size + bias_index) * image_size + image_offset);
sum += AccT(val);
}
// Write the accumulated sum in this thread to the shared memory. Each thread
// shifts their write location to avoid bank conflict.
int bias_offset = threadIdx.x % 32;
GpuAtomicAdd(s_data + bias_offset, sum);
__syncthreads();
// Accumulate the results in the shared memory into the first element.
// No syncthreads is needed since this is only in the same warp.
int32_t thread_index = threadIdx.x;
#if GOOGLE_CUDA
if (thread_index < 32) {
AccT data = s_data[thread_index];
for (int32_t delta = warpSize / 2; delta > 0; delta /= 2) {
data += GpuShuffleXorSync(kCudaWarpAll, data, delta);
}
if (thread_index == 0) {
GpuAtomicAdd(bias_backprop + bias_index, T(data));
}
}
#elif TENSORFLOW_USE_ROCM
if (thread_index < 16) s_data[thread_index] += s_data[thread_index + 16];
if (thread_index < 8) s_data[thread_index] += s_data[thread_index + 8];
if (thread_index < 4) s_data[thread_index] += s_data[thread_index + 4];
if (thread_index < 2) s_data[thread_index] += s_data[thread_index + 2];
if (thread_index < 1) s_data[thread_index] += s_data[thread_index + 1];
// The first thread writes out the accumulated result to the global location.
if (thread_index == 0) {
GpuAtomicAdd(bias_backprop + bias_index, T(s_data[0]));
}
#endif
}
template <typename T>
absl::Status BiasGradGPU<T>::compute(const GPUDevice& d,
const T* output_backprop, T* bias_backprop,
int32_t batch, int32_t height,
int32_t width, int32_t depth,
int32_t channel,
TensorFormat data_format) {
const int32_t bias_size = channel;
int64_t image_size_64 =
MultiplyWithoutOverflow(MultiplyWithoutOverflow(height, width), depth);
int64_t total_count_64 = MultiplyWithoutOverflow(
MultiplyWithoutOverflow(batch, bias_size), image_size_64);
if (total_count_64 < 0 ||
total_count_64 > std::numeric_limits<int32_t>::max() ||
image_size_64 < 0 ||
image_size_64 > std::numeric_limits<int32_t>::max()) {
return absl::InternalError("BiasGradGPU: dimensions exceed int32 bounds");
}
const int32_t image_size = image_size_64;
const int32_t total_count = total_count_64;
if (total_count == 0) {
return absl::OkStatus();
}
static constexpr int32_t kWarpSize = 32;
GpuLaunchConfig config = GetGpuLaunchConfig(total_count, d);
const int max_shared_memory_size = d.sharedMemPerBlock() / 2;
int32_t shared_memory_size = 0;
if (data_format == FORMAT_NHWC) {
shared_memory_size = bias_size * sizeof(typename AccumulatorType<T>::type);
}
// Check if we have enough shared memory.
if (shared_memory_size <= max_shared_memory_size) {
if (data_format == FORMAT_NHWC) {
TF_CHECK_OK(GpuLaunchKernel(BiasGradNHWC_SharedAtomics<T>,
config.block_count, config.thread_per_block,
shared_memory_size, d.stream(), total_count,
output_backprop, bias_backprop, bias_size));
} else {
// Round up the block count to multiple of bias_size.
int group_size = (config.block_count + bias_size - 1) / bias_size;
config.block_count = group_size * bias_size;
if (config.thread_per_block < kWarpSize) {
config.thread_per_block = kWarpSize;
}
TF_CHECK_OK(GpuLaunchKernel(BiasGradNCHW_SharedAtomics<T>,
config.block_count, config.thread_per_block,
0, d.stream(), output_backprop, bias_backprop,
batch, bias_size, image_size, group_size));
}
} else {
// Note that even if we don't have enough shared memory to fit the entire
// output block, it is possible to process one group of elements at a time.
// But for now, we simply fall back to the naive implementation.
if (data_format == FORMAT_NHWC) {
TF_CHECK_OK(GpuLaunchKernel(
BiasGradNHWC_Naive<T>, config.block_count, config.thread_per_block, 0,
d.stream(), total_count, output_backprop, bias_backprop, bias_size));
} else {
TF_CHECK_OK(GpuLaunchKernel(BiasGradNCHW_Naive<T>, config.block_count,
config.thread_per_block, 0, d.stream(),
total_count, output_backprop, bias_backprop,
bias_size, image_size));
}
}
return absl::OkStatus();
}
template <typename T>
void BiasGradGPU<T>::DoRowReduction(OpKernelContext* context, T* output,
const T* input, int rows, int cols) {
typedef const Eigen::array<TTypes<float>::Tensor::Index, 1>& ReductionAxes;
Constants<GPUDevice> constants;
gpuprim::Sum op;
functor::ReduceImpl<T, gpuprim::Sum, T*, const T*, ReductionAxes>(
context, output, input, 2, rows, cols, 1, 1, constants.kOne, op);
}
template <typename T>
void BiasGradGPU<T>::DoColReduction(OpKernelContext* context, T* output,
const T* input, int rows, int cols) {
typedef const Eigen::array<TTypes<float>::Tensor::Index, 1>& ReductionAxes;
Constants<GPUDevice> constants;
gpuprim::Sum op;
functor::ReduceImpl<T, gpuprim::Sum, T*, const T*, ReductionAxes>(
context, output, input, 2, rows, cols, 1, 1, constants.kZero, op);
}
#define DEFINE_GPU_SPECS(T) \
template struct BiasGPU<T>; \
template struct BiasGradGPU<T>;
TF_CALL_GPU_NUMBER_TYPES(DEFINE_GPU_SPECS);
// No BiasGrad kernel for int32.
template struct BiasGPU<int32_t>;
} // end namespace tensorflow
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
tensorflow/core/kernels/bias_op_gpu.h
+83
View File
@@ -0,0 +1,83 @@
/* Copyright 2015 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BIAS_OP_GPU_H_
#define TENSORFLOW_CORE_KERNELS_BIAS_OP_GPU_H_
#define EIGEN_USE_GPU
#include "absl/status/status.h"
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/kernels/gpu_utils.h"
#include "tensorflow/core/util/tensor_format.h"
namespace tensorflow {
typedef Eigen::GpuDevice GPUDevice;
template <typename T>
struct BiasGPU {
static absl::Status compute(const GPUDevice& d, const T* input, const T* bias,
T* output, int32_t batch, int32_t height,
int32_t width, int32_t depth, int32_t channel,
TensorFormat data_format);
};
template <typename T>
struct BiasGradGPU {
static absl::Status compute(const GPUDevice& device, const T* output_backprop,
T* bias_backprop, int32_t batch, int32_t height,
int32_t width, int32_t depth, int32_t channel,
TensorFormat data_format);
static void DoRowReduction(OpKernelContext* context, T* output,
const T* input, int rows, int cols);
static void DoColReduction(OpKernelContext* context, T* output,
const T* input, int rows, int cols);
};
enum class BiasAddGradGPUMode {
kInvalid = 0,
kNative = 1,
kReduction = 2,
};
// Describe the BiasGradGPU result from a perf experiment.
//
// Arguments:
// algorithm: returns the method to use for bias add grad.
// elapsed_time; returns the measured elapsed time in microseconds.
class BiasGradGPUProfileResult {
public:
bool is_valid() const {
return (algorithm_ != BiasAddGradGPUMode::kInvalid &&
elapsed_time_ != std::numeric_limits<float>::max());
}
BiasAddGradGPUMode algorithm() const { return algorithm_; }
void set_algorithm(BiasAddGradGPUMode val) { algorithm_ = val; }
uint64_t elapsed_time() const { return elapsed_time_; }
void set_elapsed_time(uint64_t val) { elapsed_time_ = val; }
private:
BiasAddGradGPUMode algorithm_ = BiasAddGradGPUMode::kInvalid;
uint64_t elapsed_time_ = std::numeric_limits<uint64_t>::max();
};
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BIAS_OP_GPU_H_
tensorflow/core/kernels/bias_op_test.cc
+101
View File
@@ -0,0 +1,101 @@
/* Copyright 2019 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/bias_op.h"
#include <random>
#include "tensorflow/core/common_runtime/kernel_benchmark_testlib.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
namespace tensorflow {
static Graph* BiasAdd(int d0, int d1, int d2, int d3) {
auto* g = new Graph(OpRegistry::Global());
Tensor input(DT_FLOAT, TensorShape({d0, d1, d2, d3}));
Tensor bias(DT_FLOAT, TensorShape({d3}));
input.flat<float>().setRandom();
bias.flat<float>().setRandom();
test::graph::Binary(g, "BiasAdd", test::graph::Constant(g, input),
test::graph::Constant(g, bias));
return g;
}
static Graph* BiasAddGrad(int d0, int d1, int d2, int d3) {
auto* g = new Graph(OpRegistry::Global());
Tensor out_backprop(DT_FLOAT, TensorShape({d0, d1, d2, d3}));
out_backprop.flat<float>().setRandom();
test::graph::Unary(g, "BiasAddGrad", test::graph::Constant(g, out_backprop));
return g;
}
#define BM_BiasAddNHWC(N, W, H, C, DEVICE) \
static void BM_BiasAddNHWC##_##N##_##H##_##W##_##C##_##DEVICE( \
::testing::benchmark::State& state) { \
test::Benchmark(#DEVICE, BiasAdd(N, H, W, C), /*old_benchmark_api=*/false) \
.Run(state); \
state.SetItemsProcessed(static_cast<int64_t>(state.iterations()) * N * H * \
W * C); \
} \
BENCHMARK(BM_BiasAddNHWC##_##N##_##H##_##W##_##C##_##DEVICE)->UseRealTime();
#define BM_BiasAddGradNHWC(N, W, H, C, DEVICE) \
static void BM_BiasAddGradNHWC##_##N##_##H##_##W##_##C##_##DEVICE( \
::testing::benchmark::State& state) { \
test::Benchmark(#DEVICE, BiasAddGrad(N, H, W, C), \
/*old_benchmark_api=*/false) \
.Run(state); \
state.SetItemsProcessed(static_cast<int64_t>(state.iterations()) * N * H * \
W * C); \
} \
BENCHMARK(BM_BiasAddGradNHWC##_##N##_##H##_##W##_##C##_##DEVICE) \
->UseRealTime();
// CPU
BM_BiasAddNHWC(32, 32, 32, 128, cpu);
BM_BiasAddNHWC(32, 32, 32, 256, cpu);
BM_BiasAddNHWC(32, 32, 32, 512, cpu);
BM_BiasAddNHWC(32, 32, 32, 1024, cpu);
BM_BiasAddNHWC(32, 64, 64, 128, cpu);
BM_BiasAddNHWC(32, 64, 64, 256, cpu);
BM_BiasAddNHWC(32, 64, 64, 512, cpu);
BM_BiasAddNHWC(32, 64, 64, 1024, cpu);
BM_BiasAddGradNHWC(32, 32, 32, 128, cpu);
BM_BiasAddGradNHWC(32, 32, 32, 256, cpu);
BM_BiasAddGradNHWC(32, 32, 32, 512, cpu);
BM_BiasAddGradNHWC(32, 32, 32, 1024, cpu);
BM_BiasAddGradNHWC(32, 64, 64, 128, cpu);
BM_BiasAddGradNHWC(32, 64, 64, 256, cpu);
BM_BiasAddGradNHWC(32, 64, 64, 512, cpu);
BM_BiasAddGradNHWC(32, 64, 64, 1024, cpu);
#ifdef GOOGLE_CUDA
BM_BiasAddGradNHWC(32, 32, 32, 128, gpu);
BM_BiasAddGradNHWC(32, 32, 32, 256, gpu);
BM_BiasAddGradNHWC(32, 32, 32, 512, gpu);
BM_BiasAddGradNHWC(32, 32, 32, 1024, gpu);
BM_BiasAddGradNHWC(32, 64, 64, 128, gpu);
BM_BiasAddGradNHWC(32, 64, 64, 256, gpu);
BM_BiasAddGradNHWC(32, 64, 64, 512, gpu);
BM_BiasAddGradNHWC(32, 64, 64, 1024, gpu);
#endif // GOOGLE_CUDA
} // end namespace tensorflow
tensorflow/core/kernels/bincount_op.cc
+565
View File
@@ -0,0 +1,565 @@
/* Copyright 2017 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.
==============================================================================*/
// See docs in ../ops/math_ops.cc.
#define EIGEN_USE_THREADS
#include "tensorflow/core/kernels/bincount_op.h"
#include <atomic>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/kernels/fill_functor.h"
#include "tensorflow/core/kernels/sparse_utils.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/determinism.h"
namespace tensorflow {
using thread::ThreadPool;
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
namespace functor {
template <typename Tidx, typename T>
struct BincountFunctor<CPUDevice, Tidx, T, true> {
static absl::Status Compute(OpKernelContext* context,
const typename TTypes<Tidx, 1>::ConstTensor& arr,
const typename TTypes<T, 1>::ConstTensor& weights,
typename TTypes<T, 1>::Tensor& output,
const Tidx num_bins) {
Tensor all_nonneg_t;
TF_RETURN_IF_ERROR(context->allocate_temp(
DT_BOOL, TensorShape({}), &all_nonneg_t, AllocatorAttributes()));
all_nonneg_t.scalar<bool>().device(context->eigen_cpu_device()) =
(arr >= Tidx(0)).all();
if (!all_nonneg_t.scalar<bool>()()) {
return absl::InvalidArgumentError("Input arr must be non-negative!");
}
// Allocate partial output bin sums for each worker thread. Worker ids in
// ParallelForWithWorkerId range from 0 to NumThreads() inclusive.
ThreadPool* thread_pool =
context->device()->tensorflow_cpu_worker_threads()->workers;
const int64_t num_threads = thread_pool->NumThreads() + 1;
Tensor partial_bins_t;
TF_RETURN_IF_ERROR(context->allocate_temp(
DT_BOOL, TensorShape({num_threads, num_bins}), &partial_bins_t));
auto partial_bins = partial_bins_t.matrix<bool>();
partial_bins.setZero();
thread_pool->ParallelForWithWorkerId(
arr.size(), thread::ThreadPool::SchedulingParams::Adaptive(8),
[&](int64_t start_ind, int64_t limit_ind, int64_t worker_id) {
for (int64_t i = start_ind; i < limit_ind; i++) {
Tidx value = arr(i);
if (value < num_bins) {
partial_bins(worker_id, value) = true;
}
}
});
// Sum the partial bins along the 0th axis.
Eigen::array<int, 1> reduce_dim{0};
output.device(context->eigen_cpu_device()) =
partial_bins.any(reduce_dim).cast<T>();
return absl::OkStatus();
}
};
template <typename Tidx, typename T>
struct BincountFunctor<CPUDevice, Tidx, T, false> {
static absl::Status Compute(OpKernelContext* context,
const typename TTypes<Tidx, 1>::ConstTensor& arr,
const typename TTypes<T, 1>::ConstTensor& weights,
typename TTypes<T, 1>::Tensor& output,
const Tidx num_bins) {
Tensor all_nonneg_t;
TF_RETURN_IF_ERROR(context->allocate_temp(
DT_BOOL, TensorShape({}), &all_nonneg_t, AllocatorAttributes()));
all_nonneg_t.scalar<bool>().device(context->eigen_cpu_device()) =
(arr >= Tidx(0)).all();
if (!all_nonneg_t.scalar<bool>()()) {
return absl::InvalidArgumentError("Input arr must be non-negative!");
}
// Allocate partial output bin sums for each worker thread. Worker ids in
// ParallelForWithWorkerId range from 0 to NumThreads() inclusive.
ThreadPool* thread_pool =
context->device()->tensorflow_cpu_worker_threads()->workers;
const int64_t num_threads = thread_pool->NumThreads() + 1;
const Tidx* arr_data = arr.data();
const std::ptrdiff_t arr_size = arr.size();
const T* weight_data = weights.data();
if (weights.size() && weights.size() != arr_size) {
return absl::InvalidArgumentError(
"Input indices and weights must have the same size.");
}
if (num_threads == 1) {
output.setZero();
T* output_data = output.data();
if (weights.size()) {
for (int64_t i = 0; i < arr_size; i++) {
const Tidx value = arr_data[i];
if (value < num_bins) {
output_data[value] += weight_data[i];
}
}
} else {
for (int64_t i = 0; i < arr_size; i++) {
const Tidx value = arr_data[i];
if (value < num_bins) {
// Complex numbers don't support "++".
output_data[value] += T(1);
}
}
}
} else {
Tensor partial_bins_t;
TF_RETURN_IF_ERROR(context->allocate_temp(
DataTypeToEnum<T>::value, TensorShape({num_threads, num_bins}),
&partial_bins_t));
auto partial_bins = partial_bins_t.matrix<T>();
partial_bins.setZero();
thread_pool->ParallelForWithWorkerId(
arr_size, thread::ThreadPool::SchedulingParams::Adaptive(8),
[&](int64_t start_ind, int64_t limit_ind, int64_t worker_id) {
if (weights.size()) {
for (int64_t i = start_ind; i < limit_ind; i++) {
Tidx value = arr_data[i];
if (value < num_bins) {
partial_bins(worker_id, value) += weight_data[i];
}
}
} else {
for (int64_t i = start_ind; i < limit_ind; i++) {
Tidx value = arr_data[i];
if (value < num_bins) {
// Complex numbers don't support "++".
partial_bins(worker_id, value) += T(1);
}
}
}
});
// Sum the partial bins along the 0th axis.
Eigen::array<int, 1> reduce_dim{0};
output.device(context->eigen_cpu_device()) = partial_bins.sum(reduce_dim);
}
return absl::OkStatus();
}
};
template <typename Tidx, typename T, bool binary_output>
struct BincountReduceFunctor<CPUDevice, Tidx, T, binary_output> {
static absl::Status Compute(OpKernelContext* context,
const typename TTypes<Tidx, 2>::ConstTensor& in,
const typename TTypes<T, 2>::ConstTensor& weights,
typename TTypes<T, 2>::Tensor& out,
const Tidx num_bins) {
std::atomic<int> err_neg_val = 0;
const int num_rows = out.dimension(0);
const int num_cols = in.dimension(1);
ThreadPool* thread_pool =
context->device()->tensorflow_cpu_worker_threads()->workers;
thread_pool->ParallelForWithWorkerId(
num_rows, thread::ThreadPool::SchedulingParams::Adaptive(8),
[&](int64_t start_row, int64_t end_row, int64_t worker_id) {
for (int64_t i = start_row; i < end_row; ++i) {
for (int64_t j = 0; j < num_cols; ++j) {
Tidx value = in(i, j);
if (value < 0) {
err_neg_val = value;
} else if (value < num_bins) {
if (binary_output) {
out(i, value) = T(1);
} else {
if (weights.size()) {
out(i, value) += weights(i, j);
} else {
out(i, value) += T(1);
}
}
}
}
}
});
if (err_neg_val < 0) {
return absl::InvalidArgumentError(absl::StrCat(
"Input 'in' must be non-negative! Negative input value found: ",
static_cast<int>(err_neg_val)));
}
return absl::OkStatus();
}
};
} // namespace functor
template <typename Device, typename T>
class BincountOp : public OpKernel {
public:
explicit BincountOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
void Compute(OpKernelContext* ctx) override {
const Tensor& arr_t = ctx->input(0);
const Tensor& size_tensor = ctx->input(1);
OP_REQUIRES(ctx, size_tensor.dims() == 0,
absl::InvalidArgumentError(absl::StrCat(
"Shape must be rank 0 but is rank ", size_tensor.dims())));
int32_t size = size_tensor.scalar<int32_t>()();
OP_REQUIRES(ctx, size >= 0,
absl::InvalidArgumentError(
absl::StrCat("size (", size, ") must be non-negative")));
const Tensor& weights_t = ctx->input(2);
const auto arr = arr_t.flat<int32_t>();
const auto weights = weights_t.flat<T>();
Tensor* output_t;
OP_REQUIRES_OK(ctx,
ctx->allocate_output(0, TensorShape({size}), &output_t));
auto output = output_t->flat<T>();
OP_REQUIRES_OK(ctx,
functor::BincountFunctor<Device, int32_t, T, false>::Compute(
ctx, arr, weights, output, size));
}
};
#define REGISTER_KERNELS(type) \
REGISTER_KERNEL_BUILDER( \
Name("Bincount").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
BincountOp<CPUDevice, type>)
TF_CALL_NUMBER_TYPES(REGISTER_KERNELS);
#undef REGISTER_KERNELS
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define REGISTER_KERNELS(type) \
REGISTER_KERNEL_BUILDER(Name("Bincount") \
.Device(DEVICE_GPU) \
.HostMemory("size") \
.TypeConstraint<type>("T"), \
BincountOp<GPUDevice, type>)
TF_CALL_int32(REGISTER_KERNELS);
TF_CALL_float(REGISTER_KERNELS);
#undef REGISTER_KERNELS
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
template <typename Device, typename Tidx, typename T>
class DenseBincountOp : public OpKernel {
public:
explicit DenseBincountOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("binary_output", &binary_output_));
if (std::is_same<Device, GPUDevice>::value) {
OP_REQUIRES(
ctx, !OpDeterminismRequired(),
absl::UnimplementedError(
"Determinism is not yet supported in GPU implementation of "
"DenseBincount."));
}
}
void Compute(OpKernelContext* ctx) override {
const Tensor& data = ctx->input(0);
OP_REQUIRES(ctx, data.dims() <= 2,
absl::InvalidArgumentError(absl::StrCat(
"Shape must be at most rank 2 but is rank ", data.dims())));
const Tensor& size_t = ctx->input(1);
const Tensor& weights = ctx->input(2);
OP_REQUIRES(ctx, size_t.dims() == 0,
absl::InvalidArgumentError(absl::StrCat(
"Shape must be rank 0 but is rank ", size_t.dims())));
OP_REQUIRES(ctx,
weights.shape() == data.shape() || weights.NumElements() == 0,
absl::InvalidArgumentError(absl::StrCat(
"`weights` must be the same shape as `arr` or a length-0 "
"`Tensor`, in which case it acts as all weights equal to "
"1. Received ",
weights.shape().DebugString())));
Tidx size = size_t.scalar<Tidx>()();
OP_REQUIRES(ctx, size >= 0,
absl::InvalidArgumentError(
absl::StrCat("size (", size, ") must be non-negative")));
Tensor* out_t;
functor::SetZeroFunctor<Device, T> fill;
if (data.dims() <= 1) {
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({size}), &out_t));
auto out = out_t->flat<T>();
fill(ctx->eigen_device<Device>(), out);
if (binary_output_) {
OP_REQUIRES_OK(
ctx, functor::BincountFunctor<Device, Tidx, T, true>::Compute(
ctx, data.flat<Tidx>(), weights.flat<T>(), out, size));
} else {
OP_REQUIRES_OK(
ctx, functor::BincountFunctor<Device, Tidx, T, false>::Compute(
ctx, data.flat<Tidx>(), weights.flat<T>(), out, size));
}
} else if (data.dims() == 2) {
const int64_t num_rows = data.dim_size(0);
auto weight_matrix =
(weights.NumElements() == 0)
? weights.shaped<T, 2>(absl::InlinedVector<int64_t, 2UL>(2, 0))
: weights.matrix<T>();
OP_REQUIRES_OK(
ctx, ctx->allocate_output(0, TensorShape({num_rows, size}), &out_t));
auto out = out_t->matrix<T>();
fill(ctx->eigen_device<Device>(), out_t->flat<T>());
if (binary_output_) {
OP_REQUIRES_OK(
ctx, functor::BincountReduceFunctor<Device, Tidx, T, true>::Compute(
ctx, data.matrix<Tidx>(), weight_matrix, out, size));
} else {
OP_REQUIRES_OK(
ctx,
functor::BincountReduceFunctor<Device, Tidx, T, false>::Compute(
ctx, data.matrix<Tidx>(), weight_matrix, out, size));
}
}
}
private:
bool binary_output_;
};
#define REGISTER_KERNELS(Tidx, T) \
REGISTER_KERNEL_BUILDER(Name("DenseBincount") \
.Device(DEVICE_CPU) \
.TypeConstraint<T>("T") \
.TypeConstraint<Tidx>("Tidx"), \
DenseBincountOp<CPUDevice, Tidx, T>);
#define REGISTER_CPU_KERNELS(T) \
REGISTER_KERNELS(int32, T); \
REGISTER_KERNELS(int64_t, T);
TF_CALL_NUMBER_TYPES(REGISTER_CPU_KERNELS);
#undef REGISTER_CPU_KERNELS
#undef REGISTER_KERNELS
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define REGISTER_KERNELS(Tidx, T) \
REGISTER_KERNEL_BUILDER(Name("DenseBincount") \
.Device(DEVICE_GPU) \
.HostMemory("size") \
.TypeConstraint<T>("T") \
.TypeConstraint<Tidx>("Tidx"), \
DenseBincountOp<GPUDevice, Tidx, T>);
#define REGISTER_GPU_KERNELS(T) \
REGISTER_KERNELS(int32, T); \
REGISTER_KERNELS(int64_t, T);
TF_CALL_int32(REGISTER_GPU_KERNELS);
TF_CALL_float(REGISTER_GPU_KERNELS);
#undef REGISTER_GPU_KERNELS
#undef REGISTER_KERNELS
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
template <typename Device, typename Tidx, typename T>
class SparseBincountOp : public OpKernel {
public:
explicit SparseBincountOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("binary_output", &binary_output_));
}
void Compute(OpKernelContext* ctx) override {
const Tensor& indices = ctx->input(0);
const Tensor& values = ctx->input(1);
const auto values_flat = values.flat<Tidx>();
const Tensor& dense_shape = ctx->input(2);
const Tensor& size_t = ctx->input(3);
const auto weights = ctx->input(4).flat<T>();
const int64_t weights_size = weights.size();
OP_REQUIRES(ctx, size_t.dims() == 0,
absl::InvalidArgumentError(absl::StrCat(
"Shape must be rank 0 but is rank ", size_t.dims())));
Tidx size = size_t.scalar<Tidx>()();
OP_REQUIRES(ctx, size >= 0,
absl::InvalidArgumentError(
absl::StrCat("size (", size, ") must be non-negative")));
OP_REQUIRES_OK(ctx, sparse_utils::ValidateSparseTensor<int64_t>(
indices, values, dense_shape,
sparse_utils::IndexValidation::kUnordered));
OP_REQUIRES(ctx, dense_shape.NumElements() > 0,
absl::InvalidArgumentError(absl::StrCat(
"dense_shape must have at least 1 dimension, got ",
dense_shape.NumElements())));
bool is_1d = dense_shape.NumElements() == 1;
Tensor* out_t;
functor::SetZeroFunctor<Device, T> fill;
if (is_1d) {
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({size}), &out_t));
auto out = out_t->flat<T>();
fill(ctx->eigen_device<Device>(), out);
if (binary_output_) {
OP_REQUIRES_OK(ctx,
functor::BincountFunctor<Device, Tidx, T, true>::Compute(
ctx, values_flat, weights, out, size));
} else {
OP_REQUIRES_OK(
ctx, functor::BincountFunctor<Device, Tidx, T, false>::Compute(
ctx, values_flat, weights, out, size));
}
} else {
const auto shape = dense_shape.flat<int64_t>();
const int64_t num_rows = shape(0);
OP_REQUIRES_OK(
ctx, ctx->allocate_output(0, TensorShape({num_rows, size}), &out_t));
const auto out = out_t->matrix<T>();
fill(ctx->eigen_device<Device>(), out_t->flat<T>());
const auto indices_mat = indices.matrix<int64_t>();
for (int64_t i = 0; i < indices_mat.dimension(0); ++i) {
const int64_t batch = indices_mat(i, 0);
const Tidx bin = values_flat(i);
OP_REQUIRES(
ctx, batch < out.dimension(0),
errors::InvalidArgument("Index out of bound. `batch` (", batch,
") must be less than the dimension size (",
out.dimension(0), ")."));
if (0 <= bin && bin < size) {
if (binary_output_) {
out(batch, bin) = T(1);
} else {
if (weights_size) {
out(batch, bin) += weights(i);
} else {
out(batch, bin) += T(1);
}
}
}
}
}
}
private:
bool binary_output_;
};
#define REGISTER_KERNELS(Tidx, T) \
REGISTER_KERNEL_BUILDER(Name("SparseBincount") \
.Device(DEVICE_CPU) \
.TypeConstraint<T>("T") \
.TypeConstraint<Tidx>("Tidx"), \
SparseBincountOp<CPUDevice, Tidx, T>);
#define REGISTER_CPU_KERNELS(T) \
REGISTER_KERNELS(int32, T); \
REGISTER_KERNELS(int64_t, T);
TF_CALL_NUMBER_TYPES(REGISTER_CPU_KERNELS);
#undef REGISTER_CPU_KERNELS
#undef REGISTER_KERNELS
template <typename Device, typename Tidx, typename T>
class RaggedBincountOp : public OpKernel {
public:
explicit RaggedBincountOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("binary_output", &binary_output_));
}
void Compute(OpKernelContext* ctx) override {
const auto splits = ctx->input(0).flat<int64_t>();
const auto values = ctx->input(1).flat<Tidx>();
const Tensor& size_t = ctx->input(2);
const auto weights = ctx->input(3).flat<T>();
const int64_t weights_size = weights.size();
OP_REQUIRES(ctx, size_t.dims() == 0,
absl::InvalidArgumentError(absl::StrCat(
"Shape must be rank 0 but is rank ", size_t.dims())));
Tidx size = size_t.scalar<Tidx>()();
OP_REQUIRES(ctx, size >= 0,
absl::InvalidArgumentError(
absl::StrCat("size (", size, ") must be non-negative")));
int num_rows = splits.size() - 1;
int num_values = values.size();
int batch_idx = 0;
OP_REQUIRES(ctx, splits.size() > 0,
absl::InvalidArgumentError("Splits must be non-empty"));
OP_REQUIRES(ctx, splits(0) == 0,
absl::InvalidArgumentError(absl::StrCat(
"Splits must start with 0, not with ", splits(0))));
OP_REQUIRES(ctx, splits(num_rows) == num_values,
absl::InvalidArgumentError(absl::StrCat(
"Splits must end with the number of values, got ",
splits(num_rows), " instead of ", num_values)));
Tensor* out_t;
OP_REQUIRES_OK(
ctx, ctx->allocate_output(0, TensorShape({num_rows, size}), &out_t));
functor::SetZeroFunctor<Device, T> fill;
fill(ctx->eigen_device<Device>(), out_t->flat<T>());
const auto out = out_t->matrix<T>();
for (int idx = 0; idx < num_values; ++idx) {
while (idx >= splits(batch_idx)) {
batch_idx++;
}
Tidx bin = values(idx);
OP_REQUIRES(ctx, bin >= 0,
absl::InvalidArgumentError("Input must be non-negative"));
if (bin < size) {
if (binary_output_) {
out(batch_idx - 1, bin) = T(1);
} else {
T value = (weights_size > 0) ? weights(idx) : T(1);
out(batch_idx - 1, bin) += value;
}
}
}
}
private:
bool binary_output_;
};
#define REGISTER_KERNELS(Tidx, T) \
REGISTER_KERNEL_BUILDER(Name("RaggedBincount") \
.Device(DEVICE_CPU) \
.TypeConstraint<T>("T") \
.TypeConstraint<Tidx>("Tidx"), \
RaggedBincountOp<CPUDevice, Tidx, T>);
#define REGISTER_CPU_KERNELS(T) \
REGISTER_KERNELS(int32, T); \
REGISTER_KERNELS(int64_t, T);
TF_CALL_NUMBER_TYPES(REGISTER_CPU_KERNELS);
#undef REGISTER_CPU_KERNELS
#undef REGISTER_KERNELS
} // end namespace tensorflow
tensorflow/core/kernels/bincount_op.h
+51
View File
@@ -0,0 +1,51 @@
/* Copyright 2015 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BINCOUNT_OP_H_
#define TENSORFLOW_CORE_KERNELS_BINCOUNT_OP_H_
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/core/errors.h"
namespace tensorflow {
namespace functor {
template <typename Device, typename Tidx, typename T, bool binary_count>
struct BincountFunctor {
static absl::Status Compute(OpKernelContext* context,
const typename TTypes<Tidx, 1>::ConstTensor& arr,
const typename TTypes<T, 1>::ConstTensor& weights,
typename TTypes<T, 1>::Tensor& output,
const Tidx num_bins);
};
template <typename Device, typename Tidx, typename T, bool binary_count>
struct BincountReduceFunctor {
static absl::Status Compute(OpKernelContext* context,
const typename TTypes<Tidx, 2>::ConstTensor& in,
const typename TTypes<T, 2>::ConstTensor& weights,
typename TTypes<T, 2>::Tensor& out,
const Tidx num_bins);
};
} // end namespace functor
} // end namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BINCOUNT_OP_H_
tensorflow/core/kernels/bincount_op_gpu.cu.cc
+255
View File
@@ -0,0 +1,255 @@
/* Copyright 2017 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.
==============================================================================*/
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define EIGEN_USE_GPU
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/kernels/bincount_op.h"
#include "tensorflow/core/kernels/gpu_prim.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/determinism.h"
#include "tensorflow/core/util/gpu_kernel_helper.h"
namespace tensorflow {
typedef Eigen::GpuDevice GPUDevice;
namespace functor {
template <typename Tidx, typename T>
struct BincountFunctor<GPUDevice, Tidx, T, false> {
static absl::Status Compute(OpKernelContext* context,
const typename TTypes<Tidx, 1>::ConstTensor& arr,
const typename TTypes<T, 1>::ConstTensor& weights,
typename TTypes<T, 1>::Tensor& output,
const Tidx num_bins) {
if (weights.size() != 0) {
return absl::UnimplementedError(
"Weights are not yet supported by the GPU implementation of Bincount."
" Please use unsorted_segment_sum instead or put Bincount inside"
" tf.function(jit_compile=True).");
}
if (output.size() == 0) {
return absl::OkStatus();
}
if (tensorflow::OpDeterminismRequired()) {
// TODO(reedwm): Is this really nondeterministic?
// DeviceHistogram::HistogramEven is called, and it is unclear
// if it is deterministic on floating-point inputs.
// See https://github.com/NVIDIA/cub/issues/471#issuecomment-1194682443.
return absl::UnimplementedError(
"Determinism is not yet supported in GPU implementation of "
"Bincount.");
}
// In case weight.size() == 0, use CUB
size_t temp_storage_bytes = 0;
const Tidx* d_samples = arr.data();
T* d_histogram = output.data();
int num_levels = output.size() + 1;
Tidx lower_level = Tidx(0);
Tidx upper_level = num_bins;
int num_samples = arr.size();
const gpuStream_t& stream = GetGpuStream(context);
// The first HistogramEven is to obtain the temp storage size required
// with d_temp_storage = NULL passed to the call.
auto err = gpuprim::DeviceHistogram::HistogramEven(
/* d_temp_storage */ NULL,
/* temp_storage_bytes */ temp_storage_bytes,
/* d_samples */ d_samples,
/* d_histogram */ d_histogram,
/* num_levels */ num_levels,
/* lower_level */ lower_level,
/* upper_level */ upper_level,
/* num_samples */ num_samples,
/* stream */ stream);
if (err != gpuSuccess) {
return errors::Internal(
"Could not launch HistogramEven to get temp storage: ",
GpuGetErrorString(err), ".");
}
Tensor temp_storage;
TF_RETURN_IF_ERROR(context->allocate_temp(
DataTypeToEnum<int8_t>::value,
TensorShape({static_cast<int64_t>(temp_storage_bytes)}),
&temp_storage));
void* d_temp_storage = temp_storage.flat<int8_t>().data();
// The second HistogramEven is to actual run with d_temp_storage
// allocated with temp_storage_bytes.
err = gpuprim::DeviceHistogram::HistogramEven(
/* d_temp_storage */ d_temp_storage,
/* temp_storage_bytes */ temp_storage_bytes,
/* d_samples */ d_samples,
/* d_histogram */ d_histogram,
/* num_levels */ num_levels,
/* lower_level */ lower_level,
/* upper_level */ upper_level,
/* num_samples */ num_samples,
/* stream */ stream);
if (err != gpuSuccess) {
return errors::Internal(
"Could not launch HistogramEven: ", GpuGetErrorString(err), ".");
}
return absl::OkStatus();
}
};
template <typename Tidx, typename T>
__global__ void BincountReduceKernel(const Tidx* in, T* out, const int nthreads,
const Tidx num_bins) {
GPU_1D_KERNEL_LOOP(index, nthreads) {
Tidx bin = ldg(in + index);
if (bin < num_bins) {
out[bin] = T(1);
}
}
}
template <typename Tidx, typename T>
struct BincountFunctor<GPUDevice, Tidx, T, true> {
static absl::Status Compute(OpKernelContext* context,
const typename TTypes<Tidx, 1>::ConstTensor& arr,
const typename TTypes<T, 1>::ConstTensor& weights,
typename TTypes<T, 1>::Tensor& output,
const Tidx num_bins) {
const int nthreads = arr.dimension(0);
auto d = context->eigen_gpu_device();
GpuLaunchConfig config = GetGpuLaunchConfig(nthreads, d);
return GpuLaunchKernel(BincountReduceKernel<Tidx, T>, config.block_count,
config.thread_per_block, 0, d.stream(), arr.data(),
output.data(), nthreads, num_bins);
}
};
template <typename Tidx, typename T, bool binary_count>
__global__ void BincountColReduceKernel(const Tidx* in, const T* weights,
const int weights_size, T* out,
const int num_rows, const int num_cols,
const Tidx num_bins) {
const int nthreads = num_rows * num_cols;
GPU_1D_KERNEL_LOOP(index, nthreads) {
Tidx bin = ldg(in + index);
if (bin < num_bins) {
int row = index / num_cols;
int offset = row * num_bins + bin;
if (binary_count) {
out[offset] = T(1);
} else {
T value = (weights_size == 0) ? T(1) : ldg(weights + index);
GpuAtomicAdd(out + offset, value);
}
}
}
}
template <typename Tidx, typename T, bool binary_count>
__global__ void BincountColReduceSharedKernel(const Tidx* in, const T* weights,
const int weights_size, T* out,
const int num_rows,
const int num_cols,
const Tidx num_bins) {
const int out_size = num_rows * num_bins;
GPU_DYNAMIC_SHARED_MEM_DECL(sizeof(T), unsigned char, shared_col_mem);
T* shared_col_bins = reinterpret_cast<T*>(shared_col_mem);
for (unsigned int binIdx = threadIdx.x; binIdx < out_size;
binIdx += blockDim.x) {
shared_col_bins[binIdx] = T(0);
}
__syncthreads();
const int nthreads = num_rows * num_cols;
GPU_1D_KERNEL_LOOP(index, nthreads) {
Tidx bin = ldg(in + index);
if (bin < num_bins) {
int row = index / num_cols;
int offset = row * num_bins + bin;
if (binary_count) {
shared_col_bins[offset] = T(1);
} else {
T value = (weights_size == 0) ? T(1) : ldg(weights + index);
GpuAtomicAdd(shared_col_bins + offset, value);
}
}
}
__syncthreads();
for (unsigned int binIdx = threadIdx.x; binIdx < out_size;
binIdx += blockDim.x) {
if (binary_count) {
// out[binIdx] = out[binIdx] & shared_col_bins[binIdx];
if (shared_col_bins[binIdx]) {
out[binIdx] = shared_col_bins[binIdx];
}
} else {
GpuAtomicAdd(out + binIdx, shared_col_bins[binIdx]);
}
}
}
template <typename Tidx, typename T, bool binary_count>
struct BincountReduceFunctor<GPUDevice, Tidx, T, binary_count> {
static absl::Status Compute(OpKernelContext* context,
const typename TTypes<Tidx, 2>::ConstTensor& in,
const typename TTypes<T, 2>::ConstTensor& weights,
typename TTypes<T, 2>::Tensor& out,
const Tidx num_bins) {
const int num_rows = in.dimension(0);
const int num_cols = in.dimension(1);
auto d = context->eigen_gpu_device();
GpuLaunchConfig config = GetGpuLaunchConfig(num_rows * num_cols, d);
// Use half of maximum shared memory, approximately 6 * 1024 inputs.
int smem_max = d.sharedMemPerBlock() / 2;
int smem_usage = out.size() * sizeof(T);
if (smem_usage < smem_max) {
return GpuLaunchKernel(
BincountColReduceSharedKernel<Tidx, T, binary_count>,
config.block_count, config.thread_per_block, smem_usage, d.stream(),
in.data(), weights.data(), weights.size(), out.data(), num_rows,
num_cols, num_bins);
}
return GpuLaunchKernel(
BincountColReduceKernel<Tidx, T, binary_count>, config.block_count,
config.thread_per_block, 0, d.stream(), in.data(), weights.data(),
weights.size(), out.data(), num_rows, num_cols, num_bins);
}
};
} // end namespace functor
#define REGISTER_GPU_SPEC(T) \
template struct functor::BincountFunctor<GPUDevice, int32, T, true>; \
template struct functor::BincountFunctor<GPUDevice, int64, T, true>; \
template struct functor::BincountFunctor<GPUDevice, int32, T, false>; \
template struct functor::BincountFunctor<GPUDevice, int64, T, false>; \
template struct functor::BincountReduceFunctor<GPUDevice, int32, T, true>; \
template struct functor::BincountReduceFunctor<GPUDevice, int64, T, true>; \
template struct functor::BincountReduceFunctor<GPUDevice, int32, T, false>; \
template struct functor::BincountReduceFunctor<GPUDevice, int64, T, false>;
TF_CALL_int32(REGISTER_GPU_SPEC);
TF_CALL_float(REGISTER_GPU_SPEC);
#undef REGISTER_GPU_SPEC
} // namespace tensorflow
#endif // GOOGLE_CUDA
tensorflow/core/kernels/bincount_op_test.cc
+79
View File
@@ -0,0 +1,79 @@
/* Copyright 2015 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/kernel_benchmark_testlib.h"
#include "tensorflow/core/framework/fake_input.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
namespace tensorflow {
static Graph* Bincount(int arr_size, int nbins) {
Graph* g = new Graph(OpRegistry::Global());
Tensor arr(DT_INT32, TensorShape({arr_size}));
arr.flat<int32_t>() = arr.flat<int32_t>().setRandom().abs();
Tensor size(DT_INT32, TensorShape({static_cast<int32_t>(1)}));
size.flat<int32_t>()(0) = static_cast<int32_t>(nbins);
Tensor weights(DT_INT32, TensorShape({0}));
Node* node;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "Bincount")
.Input(test::graph::Constant(g, arr))
.Input(test::graph::Constant(g, size))
.Input(test::graph::Constant(g, weights))
.Attr("T", DT_INT32)
.Finalize(g, &node));
return g;
}
#define BM_BincountDev(K, NBINS, type) \
static void BM_Bincount##_##type##_##K##_##NBINS( \
::testing::benchmark::State& state) { \
test::Benchmark(#type, Bincount(K * 1024, NBINS), \
/*old_benchmark_api=*/false) \
.Run(state); \
state.SetItemsProcessed(static_cast<int64_t>(state.iterations()) * K * \
1024); \
} \
BENCHMARK(BM_Bincount##_##type##_##K##_##NBINS);
BM_BincountDev(32, 1000, cpu);
BM_BincountDev(32, 2000, cpu);
BM_BincountDev(32, 5000, cpu);
BM_BincountDev(64, 1000, cpu);
BM_BincountDev(64, 2000, cpu);
BM_BincountDev(64, 5000, cpu);
BM_BincountDev(128, 1000, cpu);
BM_BincountDev(128, 2000, cpu);
BM_BincountDev(128, 5000, cpu);
BM_BincountDev(32, 1000, gpu);
BM_BincountDev(32, 2000, gpu);
BM_BincountDev(32, 5000, gpu);
BM_BincountDev(64, 1000, gpu);
BM_BincountDev(64, 2000, gpu);
BM_BincountDev(64, 5000, gpu);
BM_BincountDev(128, 1000, gpu);
BM_BincountDev(128, 2000, gpu);
BM_BincountDev(128, 5000, gpu);
} // end namespace tensorflow
tensorflow/core/kernels/broadcast_to_op.cc
+168
View File
@@ -0,0 +1,168 @@
/* Copyright 2017 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.
==============================================================================*/
#define EIGEN_USE_THREADS
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
#define EIGEN_USE_GPU
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#if !defined(PLUGGABLE_DEVICE_SUPPORTED_MACOS) && defined(__APPLE__) && \
!defined(ANDROID) && !defined(__ANDROID__) && \
(!defined(TARGET_OS_IOS) || !TARGET_OS_IOS)
#define PLUGGABLE_DEVICE_SUPPORTED_MACOS 1
#endif
#include "tensorflow/core/kernels/broadcast_to_op.h"
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_util.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/util/bcast.h"
namespace tensorflow {
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
template <typename Device, typename T>
class BroadcastToOp : public OpKernel {
public:
explicit BroadcastToOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
void Compute(OpKernelContext* ctx) override {
const Tensor& input_tensor = ctx->input(0);
const TensorShape& input_shape = input_tensor.shape();
const Tensor& shape_tensor = ctx->input(1);
TensorShape output_shape;
OP_REQUIRES_OK(ctx, tensor::MakeShape(shape_tensor, &output_shape));
// Handle copy.
if (output_shape == input_shape) {
ctx->set_output(0, input_tensor);
return;
}
OP_REQUIRES(ctx, input_shape.dims() <= output_shape.dims(),
absl::InvalidArgumentError(absl::StrCat(
"Rank of input (", input_shape.dims(),
") must be no greater than rank of output shape (",
output_shape.dims(), ").")));
Tensor* output_tensor = nullptr;
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, output_shape, &output_tensor));
// Handle broadcast from Scalar.
const Device& device = ctx->eigen_device<Device>();
if (input_shape.dims() == 0) {
functor::FillFunctor<Device, T>()(device, output_tensor->flat<T>(),
input_tensor.scalar<T>());
return;
}
// Check whether the broadcast is valid.
BCast bcast(BCast::FromShape(input_shape), BCast::FromShape(output_shape),
/*fewer_dims_optimization=*/true);
OP_REQUIRES(ctx, bcast.IsValid(),
absl::InvalidArgumentError(absl::StrCat(
"Incompatible shapes: ", input_shape.DebugString(), " vs. ",
output_shape.DebugString())));
OP_REQUIRES(ctx, BCast::ToShape(bcast.output_shape()) == output_shape,
errors::InvalidArgument("Unable to broadcast tensor of shape ",
input_shape, " to tensor of shape ",
output_shape));
// Handle empty case.
if (output_shape.num_elements() == 0) {
return;
}
functor::BroadcastTo<Device, T>()(device, ctx, *output_tensor, output_shape,
input_tensor, input_shape, bcast);
}
};
// As tensor::MakeShape is able to handle both DT_INT32 and DT_INT64,
// no need to have TypeConstraint for `Tidx`
#define REGISTER_KERNEL(type) \
REGISTER_KERNEL_BUILDER( \
Name("BroadcastTo").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
BroadcastToOp<CPUDevice, type>);
TF_CALL_ALL_TYPES(REGISTER_KERNEL);
TF_CALL_float8_e5m2(REGISTER_KERNEL);
TF_CALL_float8_e4m3fn(REGISTER_KERNEL);
#undef REGISTER_KERNEL
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
namespace functor {
#define DECLARE_GPU_TEMPLATE(Type) \
template <> \
void BroadcastTo<GPUDevice, Type>::operator()( \
const GPUDevice& d, OpKernelContext* ctx, Tensor& output, \
const TensorShape& output_shape, const Tensor& input, \
const TensorShape& input_shape, const BCast& bcast) const; \
extern template struct BroadcastTo<GPUDevice, Type>;
TF_CALL_GPU_ALL_TYPES(DECLARE_GPU_TEMPLATE);
TF_CALL_int64(DECLARE_GPU_TEMPLATE);
TF_CALL_float8_e5m2(DECLARE_GPU_TEMPLATE);
TF_CALL_float8_e4m3fn(DECLARE_GPU_TEMPLATE);
#undef DECLARE_GPU_KERNEL
} // namespace functor
#define REGISTER_KERNEL(type) \
REGISTER_KERNEL_BUILDER(Name("BroadcastTo") \
.Device(DEVICE_GPU) \
.TypeConstraint<type>("T") \
.HostMemory("shape"), \
BroadcastToOp<GPUDevice, type>);
TF_CALL_GPU_ALL_TYPES(REGISTER_KERNEL);
TF_CALL_int64(REGISTER_KERNEL);
TF_CALL_float8_e5m2(REGISTER_KERNEL);
TF_CALL_float8_e4m3fn(REGISTER_KERNEL);
#undef REGISTER_KERNEL
// A special GPU kernel for int32.
// TODO(b/25387198): Also enable int32 in device memory. This kernel
// registration requires all int32 inputs and outputs to be in host memory.
REGISTER_KERNEL_BUILDER(Name("BroadcastTo")
.Device(DEVICE_GPU)
.TypeConstraint<int32_t>("T")
.HostMemory("input")
.HostMemory("shape")
.HostMemory("output"),
BroadcastToOp<CPUDevice, int32_t>);
#endif
#if defined(PLUGGABLE_DEVICE_SUPPORTED_MACOS)
REGISTER_KERNEL_BUILDER(Name("BroadcastTo")
.Device(DEVICE_DEFAULT)
.TypeConstraint<int32>("T")
.HostMemory("input")
.HostMemory("shape")
.HostMemory("output"),
BroadcastToOp<CPUDevice, int32>);
#endif
} // namespace tensorflow
tensorflow/core/kernels/broadcast_to_op.h
+91
View File
@@ -0,0 +1,91 @@
/* Copyright 2015 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BROADCAST_TO_OP_H_
#define TENSORFLOW_CORE_KERNELS_BROADCAST_TO_OP_H_
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/kernels/fill_functor.h"
#include "tensorflow/core/util/bcast.h"
namespace tensorflow {
namespace functor {
template <typename Device, typename T>
struct BroadcastTo {
template <int NDIMS>
void DoBCast(
const Device &device, typename TTypes<T, NDIMS>::Tensor out,
typename TTypes<T, NDIMS>::ConstTensor in,
const typename Eigen::array<Eigen::DenseIndex, NDIMS> &bcast) const {
MaybeWith32BitIndexing<Device>(
[&](auto out32, auto in32, const auto &bcast32) {
out32.device(device) = in32.broadcast(bcast32);
},
out, in, bcast);
}
template <int NDIMS>
void ReshapeAndBCast(const Device &device, Tensor &output_tensor,
const Tensor &input_tensor, const BCast &bcast) const {
DoBCast<NDIMS>(
device, output_tensor.template shaped<T, NDIMS>(bcast.result_shape()),
input_tensor.template shaped<T, NDIMS>(bcast.x_reshape()),
BCast::ToIndexArrayType<Eigen::DenseIndex, NDIMS>(bcast.x_bcast()));
}
// PRECONDITION: rank(input_shape) > 0 &&
// rank(input_shape) <= rank(output_shape) &&
// output_shape.num_elements() > 0.
void operator()(const Device &device, OpKernelContext *ctx,
Tensor &output_tensor, const TensorShape &output_shape,
const Tensor &input_tensor, const TensorShape &input_shape,
const BCast &bcast) const {
const int ndims = bcast.y_reshape().size();
switch (ndims) {
case 1:
ReshapeAndBCast<1>(device, output_tensor, input_tensor, bcast);
break;
case 2:
ReshapeAndBCast<2>(device, output_tensor, input_tensor, bcast);
break;
case 3:
ReshapeAndBCast<3>(device, output_tensor, input_tensor, bcast);
break;
case 4:
ReshapeAndBCast<4>(device, output_tensor, input_tensor, bcast);
break;
case 5:
ReshapeAndBCast<5>(device, output_tensor, input_tensor, bcast);
break;
default:
ctx->SetStatus(absl::UnimplementedError(absl::StrCat(
"Broadcast between ", input_shape.DebugString(), " and ",
output_shape.DebugString(), " is not supported yet.")));
break;
}
}
};
} // namespace functor
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BROADCAST_TO_OP_H_
tensorflow/core/kernels/broadcast_to_op_gpu.cu.cc
+38
View File
@@ -0,0 +1,38 @@
/* Copyright 2018 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.
==============================================================================*/
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
#define EIGEN_USE_GPU
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/kernels/broadcast_to_op.h"
namespace tensorflow {
typedef Eigen::GpuDevice GPUDevice;
#define INSTANTIATE_GPU_KERNEL(Type) \
template class functor::BroadcastTo<GPUDevice, Type>;
TF_CALL_GPU_ALL_TYPES(INSTANTIATE_GPU_KERNEL);
TF_CALL_int64(INSTANTIATE_GPU_KERNEL);
TF_CALL_float8_e5m2(INSTANTIATE_GPU_KERNEL);
TF_CALL_float8_e4m3fn(INSTANTIATE_GPU_KERNEL);
#undef INSTANTIATE_GPU_KERNEL
} // namespace tensorflow
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
tensorflow/core/kernels/broadcast_to_op_test.cc
+91
View File
@@ -0,0 +1,91 @@
/* Copyright 2019 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/kernel_benchmark_testlib.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
namespace tensorflow {
template <typename InputShape>
static Graph* BroadcastTo(int dim0, int dim1, InputShape input_shape) {
Graph* g = new Graph(OpRegistry::Global());
Tensor input(DT_FLOAT, input_shape(dim0, dim1));
input.flat<float>() = input.flat<float>().setRandom();
Tensor shape(DT_INT32, TensorShape({2}));
shape.flat<int32_t>()(0) = dim0;
shape.flat<int32_t>()(1) = dim1;
Node* node;
TF_CHECK_OK(NodeBuilder(g->NewName("n"), "BroadcastTo")
.Input(test::graph::Constant(g, input))
.Input(test::graph::Constant(g, shape))
.Attr("T", DT_FLOAT)
.Attr("Tidx", DT_INT32)
.Finalize(g, &node));
return g;
}
#define BM_BroadcastTo_InnerDim(DIM0, DIM1, type) \
static void BM_BroadcastTo_Inner##_##type##_##DIM0##_##DIM1( \
::testing::benchmark::State& state) { \
test::Benchmark(#type, \
BroadcastTo(DIM0, DIM1, \
[](int dim0, int dim1) { \
return TensorShape({dim0, 1}); \
}), \
/*old_benchmark_api=*/false) \
.Run(state); \
state.SetItemsProcessed(static_cast<int64_t>(state.iterations()) * DIM0 * \
DIM1); \
} \
BENCHMARK(BM_BroadcastTo_Inner##_##type##_##DIM0##_##DIM1)->UseRealTime();
#define BM_BroadcastTo_OuterDim(DIM0, DIM1, type) \
static void BM_BroadcastTo_Outer##_##type##_##DIM0##_##DIM1( \
::testing::benchmark::State& state) { \
test::Benchmark(#type, \
BroadcastTo(DIM0, DIM1, \
[](int dim0, int dim1) { \
return TensorShape({1, dim1}); \
}), \
/*old_benchmark_api=*/false) \
.Run(state); \
state.SetItemsProcessed(static_cast<int64_t>(state.iterations()) * DIM0 * \
DIM1); \
} \
BENCHMARK(BM_BroadcastTo_Outer##_##type##_##DIM0##_##DIM1)->UseRealTime();
BM_BroadcastTo_InnerDim(64, 64, cpu);
BM_BroadcastTo_InnerDim(128, 128, cpu);
BM_BroadcastTo_InnerDim(256, 256, cpu);
BM_BroadcastTo_InnerDim(512, 512, cpu);
BM_BroadcastTo_InnerDim(1024, 1024, cpu);
BM_BroadcastTo_InnerDim(500, 20000, cpu);
BM_BroadcastTo_OuterDim(64, 64, cpu);
BM_BroadcastTo_OuterDim(128, 128, cpu);
BM_BroadcastTo_OuterDim(256, 256, cpu);
BM_BroadcastTo_OuterDim(512, 512, cpu);
BM_BroadcastTo_OuterDim(1024, 1024, cpu);
BM_BroadcastTo_OuterDim(500, 20000, cpu);
} // end namespace tensorflow
tensorflow/core/kernels/bucketize_op.cc
+109
View File
@@ -0,0 +1,109 @@
/* Copyright 2016 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.
==============================================================================*/
// See docs in ../ops/math_ops.cc.
#include "tensorflow/core/kernels/bucketize_op.h"
#include <algorithm>
#include <vector>
#include "absl/status/status.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
using CPUDevice = Eigen::ThreadPoolDevice;
using GPUDevice = Eigen::GpuDevice;
namespace functor {
template <typename T>
struct BucketizeFunctor<CPUDevice, T> {
// PRECONDITION: boundaries_vector must be sorted.
static absl::Status Compute(OpKernelContext* context,
const typename TTypes<T, 1>::ConstTensor& input,
const std::vector<float>& boundaries_vector,
typename TTypes<int32_t, 1>::Tensor& output) {
const int N = input.size();
for (int i = 0; i < N; i++) {
auto first_bigger_it = std::upper_bound(
boundaries_vector.begin(), boundaries_vector.end(), input(i));
output(i) = first_bigger_it - boundaries_vector.begin();
}
return absl::OkStatus();
}
};
} // namespace functor
template <typename Device, typename T>
class BucketizeOp : public OpKernel {
public:
explicit BucketizeOp(OpKernelConstruction* context) : OpKernel(context) {
OP_REQUIRES_OK(context, context->GetAttr("boundaries", &boundaries_));
OP_REQUIRES(context, std::is_sorted(boundaries_.begin(), boundaries_.end()),
absl::InvalidArgumentError("Expected sorted boundaries"));
}
void Compute(OpKernelContext* context) override {
const Tensor& input_tensor = context->input(0);
const auto input = input_tensor.flat<T>();
Tensor* output_tensor = nullptr;
OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
&output_tensor));
auto output = output_tensor->template flat<int32_t>();
if (input.size() > 0) {
OP_REQUIRES_OK(context, functor::BucketizeFunctor<Device, T>::Compute(
context, input, boundaries_, output));
}
}
private:
std::vector<float> boundaries_;
};
#define REGISTER_KERNEL(T) \
REGISTER_KERNEL_BUILDER( \
Name("Bucketize").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
BucketizeOp<CPUDevice, T>);
REGISTER_KERNEL(int32_t);
REGISTER_KERNEL(int64_t);
REGISTER_KERNEL(float);
REGISTER_KERNEL(double);
#undef REGISTER_KERNEL
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define REGISTER_KERNEL(T) \
REGISTER_KERNEL_BUILDER( \
Name("Bucketize").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
BucketizeOp<GPUDevice, T>);
REGISTER_KERNEL(int32_t);
REGISTER_KERNEL(int64_t);
REGISTER_KERNEL(float);
REGISTER_KERNEL(double);
#undef REGISTER_KERNEL
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow
tensorflow/core/kernels/bucketize_op.h
+42
View File
@@ -0,0 +1,42 @@
/* Copyright 2015 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_BUCKETIZE_OP_H_
#define TENSORFLOW_CORE_KERNELS_BUCKETIZE_OP_H_
#include <vector>
#include "absl/status/status.h"
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/core/errors.h"
namespace tensorflow {
namespace functor {
template <typename Device, typename T>
struct BucketizeFunctor {
static absl::Status Compute(OpKernelContext* context,
const typename TTypes<T, 1>::ConstTensor& input,
const std::vector<float>& boundaries_vector,
typename TTypes<int32_t, 1>::Tensor& output);
};
} // namespace functor
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_BUCKETIZE_OP_H_
tensorflow/core/kernels/bucketize_op_gpu.cu.cc
+128
View File
@@ -0,0 +1,128 @@
/* Copyright 2016 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.
==============================================================================*/
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define EIGEN_USE_GPU
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/kernels/bucketize_op.h"
#include "tensorflow/core/kernels/gpu_device_array.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/gpu_kernel_helper.h"
namespace tensorflow {
typedef Eigen::GpuDevice GPUDevice;
template <typename T, bool useSharedMem>
__global__ void BucketizeCustomKernel(
const int32_t size_in, const T* __restrict__ in,
const int32_t size_boundaries, GpuDeviceArrayStruct<float> boundaries_array,
int32_t* __restrict__ out) {
const float* boundaries = GetGpuDeviceArrayOnDevice(&boundaries_array);
GPU_DYNAMIC_SHARED_MEM_DECL(sizeof(float), unsigned char, shared_mem);
float* shared_mem_boundaries = reinterpret_cast<float*>(shared_mem);
if (useSharedMem) {
int32_t lidx = threadIdx.y * blockDim.x + threadIdx.x;
int32_t blockSize = blockDim.x * blockDim.y;
for (int32_t i = lidx; i < size_boundaries; i += blockSize) {
shared_mem_boundaries[i] = boundaries[i];
}
__syncthreads();
boundaries = shared_mem_boundaries;
}
GPU_1D_KERNEL_LOOP(i, size_in) {
T value = in[i];
int32_t bucket = 0;
int32_t count = size_boundaries;
while (count > 0) {
int32_t l = bucket;
int32_t step = count / 2;
l += step;
if (!(value < static_cast<T>(boundaries[l]))) {
bucket = ++l;
count -= step + 1;
} else {
count = step;
}
}
out[i] = bucket;
}
}
namespace functor {
template <typename T>
struct BucketizeFunctor<GPUDevice, T> {
// PRECONDITION: boundaries_vector must be sorted.
static absl::Status Compute(OpKernelContext* context,
const typename TTypes<T, 1>::ConstTensor& input,
const std::vector<float>& boundaries_vector,
typename TTypes<int32_t, 1>::Tensor& output) {
const GPUDevice& d = context->eigen_device<GPUDevice>();
GpuDeviceArrayOnHost<float> boundaries_array(context,
boundaries_vector.size());
TF_RETURN_IF_ERROR(boundaries_array.Init());
for (int i = 0; i < boundaries_vector.size(); ++i) {
boundaries_array.Set(i, boundaries_vector[i]);
}
TF_RETURN_IF_ERROR(boundaries_array.Finalize());
GpuLaunchConfig config = GetGpuLaunchConfig(input.size(), d);
int32_t shared_mem_size = sizeof(float) * boundaries_vector.size();
const int32_t kMaxSharedMemBytes = 16384;
if (shared_mem_size < d.sharedMemPerBlock() &&
shared_mem_size < kMaxSharedMemBytes) {
TF_CHECK_OK(GpuLaunchKernel(BucketizeCustomKernel<T, true>,
config.block_count, config.thread_per_block,
shared_mem_size, d.stream(), input.size(),
input.data(), boundaries_vector.size(),
boundaries_array.data(), output.data()));
} else {
TF_CHECK_OK(GpuLaunchKernel(
BucketizeCustomKernel<T, false>, config.block_count,
config.thread_per_block, 0, d.stream(), input.size(), input.data(),
boundaries_vector.size(), boundaries_array.data(), output.data()));
}
return absl::OkStatus();
}
};
} // namespace functor
#define REGISTER_GPU_SPEC(type) \
template struct functor::BucketizeFunctor<GPUDevice, type>;
REGISTER_GPU_SPEC(int32_t);
REGISTER_GPU_SPEC(int64_t);
REGISTER_GPU_SPEC(float);
REGISTER_GPU_SPEC(double);
#undef REGISTER_GPU_SPEC
} // namespace tensorflow
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
tensorflow/core/kernels/candidate_sampler_ops.cc
+285
View File
@@ -0,0 +1,285 @@
/* Copyright 2015 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.
==============================================================================*/
// See docs in ../ops/candidate_sampling_ops.cc.
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/types/span.h"
#define EIGEN_USE_THREADS
#include <cfloat>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/kernels/range_sampler.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/util/guarded_philox_random.h"
namespace tensorflow {
class BaseCandidateSamplerOp : public OpKernel {
public:
explicit BaseCandidateSamplerOp(OpKernelConstruction* context)
: OpKernel(context) {
OP_REQUIRES_OK(context, context->GetAttr("num_sampled", &num_sampled_));
OP_REQUIRES_OK(context, context->GetAttr("num_true", &num_true_));
OP_REQUIRES_OK(context, context->GetAttr("unique", &unique_));
OP_REQUIRES_OK(context, generator_.Init(context));
}
void Compute(OpKernelContext* context) override {
const Tensor& true_classes = context->input(0);
OP_REQUIRES(context, true_classes.dims() == 2,
absl::InvalidArgumentError("true_classes must be a matrix"));
const int32_t batch_size = true_classes.dim_size(0);
OP_REQUIRES(context, true_classes.dim_size(1) == num_true_,
absl::InvalidArgumentError(absl::StrCat(
"true_classes must have "
"num_true columns, expected: ",
true_classes.dim_size(1), " was: ", num_true_)));
CHECK(sampler_) << "CandidateSamplerOp did not set sampler_";
if (unique_) {
OP_REQUIRES(context, num_sampled_ <= sampler_->range(),
absl::InvalidArgumentError("Sampler's range is too small."));
}
// Output candidates and expected_count.
Tensor* out_sampled_candidates = nullptr;
OP_REQUIRES_OK(context,
context->allocate_output(0, TensorShape({num_sampled_}),
&out_sampled_candidates));
Tensor* out_true_expected_count = nullptr;
OP_REQUIRES_OK(context, context->allocate_output(
1, TensorShape({batch_size, num_true_}),
&out_true_expected_count));
Tensor* out_sampled_expected_count = nullptr;
OP_REQUIRES_OK(context,
context->allocate_output(2, TensorShape({num_sampled_}),
&out_sampled_expected_count));
absl::Span<const int64_t> true_candidate(
true_classes.matrix<int64_t>().data(), batch_size * num_true_);
for (const auto& candidate : true_candidate) {
OP_REQUIRES(context, candidate >= 0 && candidate < sampler_->range(),
absl::InvalidArgumentError(absl::StrCat(
"`true_candidate` out of range [", 0, ", ",
sampler_->range(), "), received ", candidate)));
}
absl::Span<int64_t> sampled_candidate(
out_sampled_candidates->vec<int64_t>().data(), num_sampled_);
absl::Span<float> true_expected_count(
out_true_expected_count->matrix<float>().data(),
batch_size * num_true_);
absl::Span<float> sampled_expected_count(
out_sampled_expected_count->vec<float>().data(), num_sampled_);
// Approximately conservatively estimate the number of samples required.
// In cases where rejection sampling is used we may occasionally use more
// samples than expected, which will result in reused random bits.
const int64_t samples32 = 2048 * num_sampled_;
// Pick sampled candidates.
auto local_gen = generator_.ReserveSamples32(samples32);
random::SimplePhilox random(&local_gen);
sampler_->SampleBatchGetExpectedCount(&random, unique_, sampled_candidate,
sampled_expected_count,
true_candidate, true_expected_count);
if (sampler_->NeedsUpdates()) {
sampler_->Update(true_candidate);
}
}
protected:
void set_sampler(RangeSampler* sampler) { sampler_.reset(sampler); }
private:
int32_t num_true_;
int32_t num_sampled_;
bool unique_;
std::unique_ptr<RangeSampler> sampler_;
GuardedPhiloxRandom generator_;
};
template <class RangeSamplerType>
class SimpleCandidateSamplerOp : public BaseCandidateSamplerOp {
public:
explicit SimpleCandidateSamplerOp(OpKernelConstruction* context)
: BaseCandidateSamplerOp(context) {
int64_t range_max;
OP_REQUIRES_OK(context, context->GetAttr("range_max", &range_max));
set_sampler(new RangeSamplerType(range_max));
}
};
REGISTER_KERNEL_BUILDER(Name("UniformCandidateSampler").Device(DEVICE_CPU),
SimpleCandidateSamplerOp<UniformSampler>);
REGISTER_KERNEL_BUILDER(Name("LogUniformCandidateSampler").Device(DEVICE_CPU),
SimpleCandidateSamplerOp<LogUniformSampler>);
REGISTER_KERNEL_BUILDER(
Name("LearnedUnigramCandidateSampler").Device(DEVICE_CPU),
SimpleCandidateSamplerOp<UnigramSampler>);
REGISTER_KERNEL_BUILDER(
Name("ThreadUnsafeUnigramCandidateSampler").Device(DEVICE_CPU),
SimpleCandidateSamplerOp<ThreadUnsafeUnigramSampler>);
class AllCandidateSamplerOp : public BaseCandidateSamplerOp {
public:
explicit AllCandidateSamplerOp(OpKernelConstruction* context)
: BaseCandidateSamplerOp(context) {
int64_t range_max;
OP_REQUIRES_OK(context, context->GetAttr("num_sampled", &range_max));
set_sampler(new AllSampler(range_max));
}
};
REGISTER_KERNEL_BUILDER(Name("AllCandidateSampler").Device(DEVICE_CPU),
AllCandidateSamplerOp);
class FixedUnigramCandidateSamplerOp : public BaseCandidateSamplerOp {
public:
explicit FixedUnigramCandidateSamplerOp(OpKernelConstruction* context)
: BaseCandidateSamplerOp(context) {
int64_t range_max;
OP_REQUIRES_OK(context, context->GetAttr("range_max", &range_max));
std::string vocab_file;
OP_REQUIRES_OK(context, context->GetAttr("vocab_file", &vocab_file));
std::vector<float> unigrams;
OP_REQUIRES_OK(context, context->GetAttr("unigrams", &unigrams));
OP_REQUIRES(context, !vocab_file.empty() || !unigrams.empty(),
absl::InvalidArgumentError(
"Must provide either vocab_file or unigrams."));
OP_REQUIRES(context, vocab_file.empty() || unigrams.empty(),
absl::InvalidArgumentError(
"Must only provide one of vocab_file and unigrams."));
float distortion;
OP_REQUIRES_OK(context, context->GetAttr("distortion", &distortion));
int64_t num_reserved_ids;
OP_REQUIRES_OK(context,
context->GetAttr("num_reserved_ids", &num_reserved_ids));
int64_t num_shards;
OP_REQUIRES_OK(context, context->GetAttr("num_shards", &num_shards));
int64_t shard;
OP_REQUIRES_OK(context, context->GetAttr("shard", &shard));
FixedUnigramSampler* sampler = new FixedUnigramSampler(
range_max, distortion, num_reserved_ids, num_shards, shard);
if (!vocab_file.empty())
OP_REQUIRES_OK(
context, sampler->SetDistributionSampler(context->env(), vocab_file));
else
OP_REQUIRES_OK(context, sampler->SetDistributionSampler(unigrams));
set_sampler(sampler);
}
};
REGISTER_KERNEL_BUILDER(Name("FixedUnigramCandidateSampler").Device(DEVICE_CPU),
FixedUnigramCandidateSamplerOp);
class ComputeAccidentalHitsOp : public OpKernel {
public:
explicit ComputeAccidentalHitsOp(OpKernelConstruction* context)
: OpKernel(context) {
OP_REQUIRES_OK(context, context->GetAttr("num_true", &num_true_));
}
void Compute(OpKernelContext* context) override {
const Tensor& in_true_candidates = context->input(0);
const TensorShape& in_true_candidates_shape = in_true_candidates.shape();
OP_REQUIRES(context,
TensorShapeUtils::IsMatrix(in_true_candidates_shape) &&
in_true_candidates_shape.dim_size(1) == num_true_,
absl::InvalidArgumentError(
"true_candidates must be a batch_size * num_true matrix"));
const int64_t batch_size = in_true_candidates_shape.dim_size(0);
const Tensor& in_sampled_candidates = context->input(1);
OP_REQUIRES(context,
TensorShapeUtils::IsVector(in_sampled_candidates.shape()),
absl::InvalidArgumentError(
"sampled_candidates must be a vector, which is typically "
"an output from CandidateSampler"));
const int64_t num_sampled = in_sampled_candidates.dim_size(0);
absl::flat_hash_map<int64_t, int> sampled_candidate_to_pos;
sampled_candidate_to_pos.reserve(num_sampled);
for (int64_t i = 0; i < num_sampled; ++i) {
sampled_candidate_to_pos[in_sampled_candidates.vec<int64_t>()(i)] = i;
}
// Produce output in the same format as UnpackSparseFeatures.
std::vector<int> indices;
std::vector<int64_t> ids;
std::vector<float> weights;
for (int64_t i = 0; i < batch_size; ++i) {
for (int64_t j = 0; j < num_true_; ++j) {
const int64_t true_candidate =
in_true_candidates.matrix<int64_t>()(i, j);
const auto look = sampled_candidate_to_pos.find(true_candidate);
if (look != sampled_candidate_to_pos.end()) {
indices.push_back(i);
ids.push_back(look->second);
weights.push_back(-FLT_MAX);
}
}
}
Tensor* out_indices = nullptr;
OP_REQUIRES_OK(
context,
context->allocate_output(
0, TensorShape({static_cast<int>(indices.size())}), &out_indices));
Tensor* out_ids = nullptr;
OP_REQUIRES_OK(
context, context->allocate_output(
1, TensorShape({static_cast<int>(ids.size())}), &out_ids));
Tensor* out_weights = nullptr;
OP_REQUIRES_OK(
context,
context->allocate_output(
2, TensorShape({static_cast<int>(weights.size())}), &out_weights));
for (size_t i = 0; i < indices.size(); ++i) {
out_indices->vec<int32_t>()(i) = indices[i];
out_ids->vec<int64_t>()(i) = ids[i];
out_weights->vec<float>()(i) = weights[i];
}
}
private:
int64_t num_true_;
};
REGISTER_KERNEL_BUILDER(Name("ComputeAccidentalHits").Device(DEVICE_CPU),
ComputeAccidentalHitsOp);
} // namespace tensorflow
tensorflow/core/kernels/cast_op.cc
+353
View File
@@ -0,0 +1,353 @@
/* Copyright 2015 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.
==============================================================================*/
// See docs in ../ops/math_ops.cc.
#define EIGEN_USE_THREADS
#include "tensorflow/core/kernels/cast_op.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/kernels/cast_op_impl.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/work_sharder.h"
namespace tensorflow {
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
#define CURRY_TYPES2(FN, arg0) \
FN(arg0, bool); \
FN(arg0, uint8); \
FN(arg0, uint16); \
FN(arg0, uint32); \
FN(arg0, uint64); \
FN(arg0, int8); \
FN(arg0, int16); \
FN(arg0, int32); \
FN(arg0, int64_t); \
FN(arg0, Eigen::half); \
FN(arg0, bfloat16); \
FN(arg0, float); \
FN(arg0, double); \
FN(arg0, std::complex<float>); \
FN(arg0, std::complex<double>)
CastOpBase::CastOpBase(OpKernelConstruction* ctx) : OpKernel(ctx) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("SrcT", &external_src_dtype_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("DstT", &external_dst_dtype_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("Truncate", &use_truncation_));
// Quantized data types use the same underlying format as their non quantized
// version so we use the non quantized implementation for casting.
if (external_dst_dtype_ == DT_QUINT8) {
dst_dtype_ = DT_UINT8;
} else if (external_dst_dtype_ == DT_QINT8) {
dst_dtype_ = DT_INT8;
} else if (external_dst_dtype_ == DT_QINT32) {
dst_dtype_ = DT_INT32;
} else if (external_dst_dtype_ == DT_QINT16) {
dst_dtype_ = DT_INT16;
} else if (external_dst_dtype_ == DT_QUINT16) {
dst_dtype_ = DT_UINT16;
} else {
dst_dtype_ = external_dst_dtype_;
}
if (external_src_dtype_ == DT_QUINT8) {
src_dtype_ = DT_UINT8;
} else if (external_src_dtype_ == DT_QINT8) {
src_dtype_ = DT_INT8;
} else if (external_src_dtype_ == DT_QINT32) {
src_dtype_ = DT_INT32;
} else if (external_src_dtype_ == DT_QINT16) {
src_dtype_ = DT_INT16;
} else if (external_src_dtype_ == DT_QUINT16) {
src_dtype_ = DT_UINT16;
} else {
src_dtype_ = external_src_dtype_;
}
}
void CastOpBase::Compute(OpKernelContext* ctx) {
const Tensor& inp = ctx->input(0);
if (work_ == nullptr) {
ctx->set_output(0, inp);
} else if (external_src_dtype_ != src_dtype_ ||
external_dst_dtype_ != dst_dtype_) {
Tensor in;
// If the type is a quantized type we need to do a bitcast since the
// src_dtype_ is different from external_src_type_.
OP_REQUIRES_OK(ctx, in.BitcastFrom(inp, src_dtype_, inp.shape()));
Tensor* out = nullptr;
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, in.shape(), &out));
out->set_dtype(dst_dtype_);
work_(ctx, in, out, use_truncation_);
out->set_dtype(external_dst_dtype_);
} else {
Tensor* out = nullptr;
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, inp.shape(), &out));
work_(ctx, inp, out, use_truncation_);
}
}
absl::Status CastOpBase::Unimplemented() {
return absl::UnimplementedError(
absl::StrCat("Cast ", DataTypeString(external_src_dtype_), " to ",
DataTypeString(external_dst_dtype_), " is not supported"));
}
CpuCastOp::CpuCastOp(OpKernelConstruction* ctx) : CastOpBase(ctx) {
OP_REQUIRES_OK(ctx, Prepare());
}
absl::Status CpuCastOp::Prepare() {
if (external_src_dtype_ == external_dst_dtype_) {
work_ = nullptr; // Identity
return absl::OkStatus();
}
if (src_dtype_ == DT_BOOL) {
work_ = GetCpuCastFromBool(dst_dtype_);
} else if (src_dtype_ == DT_UINT8) {
work_ = GetCpuCastFromUint8(dst_dtype_);
} else if (src_dtype_ == DT_UINT16) {
work_ = GetCpuCastFromUint16(dst_dtype_);
} else if (src_dtype_ == DT_UINT32) {
work_ = GetCpuCastFromUint32(dst_dtype_);
} else if (src_dtype_ == DT_UINT64) {
work_ = GetCpuCastFromUint64(dst_dtype_);
} else if (src_dtype_ == DT_INT8) {
work_ = GetCpuCastFromInt8(dst_dtype_);
} else if (src_dtype_ == DT_INT16) {
work_ = GetCpuCastFromInt16(dst_dtype_);
} else if (src_dtype_ == DT_INT32) {
work_ = GetCpuCastFromInt32(dst_dtype_);
} else if (src_dtype_ == DT_INT64) {
work_ = GetCpuCastFromInt64(dst_dtype_);
} else if (src_dtype_ == DT_HALF) {
work_ = GetCpuCastFromHalf(dst_dtype_);
} else if (src_dtype_ == DT_FLOAT) {
work_ = GetCpuCastFromFloat(dst_dtype_);
} else if (src_dtype_ == DT_DOUBLE) {
work_ = GetCpuCastFromDouble(dst_dtype_);
} else if (src_dtype_ == DT_COMPLEX64) {
work_ = GetCpuCastFromComplex64(dst_dtype_);
} else if (src_dtype_ == DT_COMPLEX128) {
work_ = GetCpuCastFromComplex128(dst_dtype_);
} else if (src_dtype_ == DT_BFLOAT16) {
work_ = GetCpuCastFromBfloat(dst_dtype_);
} else if (src_dtype_ == DT_FLOAT8_E5M2) {
work_ = GetCpuCastFromFloat8e5m2(dst_dtype_);
} else if (src_dtype_ == DT_FLOAT8_E4M3FN) {
work_ = GetCpuCastFromFloat8e4m3fn(dst_dtype_);
} else if (src_dtype_ == DT_INT4) {
work_ = GetCpuCastFromInt4(dst_dtype_);
} else if (src_dtype_ == DT_UINT4) {
work_ = GetCpuCastFromUint4(dst_dtype_);
}
// TODO(sesse): If CPU casting to or from Eigen::half ever becomes a
// bottleneck, we could probably implement specialized support for
// vectorized versions (not the least based on F16C for Haswell
// or newer).
return work_ == nullptr ? Unimplemented() : absl::OkStatus();
}
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
class GpuCastOp : public CastOpBase {
public:
explicit GpuCastOp(OpKernelConstruction* ctx) : CastOpBase(ctx) {
OP_REQUIRES_OK(ctx, Prepare());
}
private:
absl::Status Prepare() {
if (external_src_dtype_ == external_dst_dtype_) {
work_ = nullptr; // Identity
return absl::OkStatus();
}
if (src_dtype_ == DT_BOOL) {
work_ = GetGpuCastFromBool(dst_dtype_);
} else if (src_dtype_ == DT_UINT8) {
work_ = GetGpuCastFromUint8(dst_dtype_);
} else if (src_dtype_ == DT_UINT16) {
work_ = GetGpuCastFromUint16(dst_dtype_);
} else if (src_dtype_ == DT_UINT32) {
work_ = GetGpuCastFromUint32(dst_dtype_);
} else if (src_dtype_ == DT_UINT64) {
work_ = GetGpuCastFromUint64(dst_dtype_);
} else if (src_dtype_ == DT_INT8) {
work_ = GetGpuCastFromInt8(dst_dtype_);
} else if (src_dtype_ == DT_INT16) {
work_ = GetGpuCastFromInt16(dst_dtype_);
} else if (src_dtype_ == DT_INT32) {
work_ = GetGpuCastFromInt32(dst_dtype_);
} else if (src_dtype_ == DT_INT64) {
work_ = GetGpuCastFromInt64(dst_dtype_);
} else if (src_dtype_ == DT_HALF) {
work_ = GetGpuCastFromHalf(dst_dtype_);
} else if (src_dtype_ == DT_FLOAT) {
work_ = GetGpuCastFromFloat(dst_dtype_);
} else if (src_dtype_ == DT_DOUBLE) {
work_ = GetGpuCastFromDouble(dst_dtype_);
} else if (src_dtype_ == DT_COMPLEX64) {
work_ = GetGpuCastFromComplex64(dst_dtype_);
} else if (src_dtype_ == DT_COMPLEX128) {
work_ = GetGpuCastFromComplex128(dst_dtype_);
} else if (src_dtype_ == DT_BFLOAT16) {
work_ = GetGpuCastFromBfloat(dst_dtype_);
} else if (src_dtype_ == DT_FLOAT8_E5M2) {
work_ = GetGpuCastFromFloat8e5m2(dst_dtype_);
} else if (src_dtype_ == DT_FLOAT8_E4M3FN) {
work_ = GetGpuCastFromFloat8e4m3fn(dst_dtype_);
} else if (src_dtype_ == DT_INT4) {
work_ = GetGpuCastFromInt4(dst_dtype_);
} else if (src_dtype_ == DT_UINT4) {
work_ = GetGpuCastFromUint4(dst_dtype_);
}
return work_ == nullptr ? Unimplemented() : absl::OkStatus();
}
};
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#undef CAST_CASE
REGISTER_KERNEL_BUILDER(Name("Cast").Device(DEVICE_CPU), CpuCastOp);
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
#define REGISTER_CAST_GPU(srctype, dsttype) \
REGISTER_KERNEL_BUILDER(Name("Cast") \
.TypeConstraint<srctype>("SrcT") \
.TypeConstraint<dsttype>("DstT") \
.Device(DEVICE_GPU), \
GpuCastOp)
#if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
CURRY_TYPES2(REGISTER_CAST_GPU, bool);
CURRY_TYPES2(REGISTER_CAST_GPU, int8);
CURRY_TYPES2(REGISTER_CAST_GPU, int16);
CURRY_TYPES2(REGISTER_CAST_GPU, int32);
CURRY_TYPES2(REGISTER_CAST_GPU, int64);
CURRY_TYPES2(REGISTER_CAST_GPU, uint8);
CURRY_TYPES2(REGISTER_CAST_GPU, uint16);
CURRY_TYPES2(REGISTER_CAST_GPU, uint32);
CURRY_TYPES2(REGISTER_CAST_GPU, uint64);
CURRY_TYPES2(REGISTER_CAST_GPU, Eigen::half);
CURRY_TYPES2(REGISTER_CAST_GPU, float);
CURRY_TYPES2(REGISTER_CAST_GPU, double);
CURRY_TYPES2(REGISTER_CAST_GPU, std::complex<float>);
CURRY_TYPES2(REGISTER_CAST_GPU, std::complex<double>);
#else
REGISTER_CAST_GPU(bool, bfloat16);
REGISTER_CAST_GPU(int8_t, bfloat16);
REGISTER_CAST_GPU(int16_t, bfloat16);
REGISTER_CAST_GPU(int32_t, bfloat16);
REGISTER_CAST_GPU(int64_t, bfloat16);
REGISTER_CAST_GPU(uint8_t, bfloat16);
REGISTER_CAST_GPU(uint16_t, bfloat16);
REGISTER_CAST_GPU(uint32_t, bfloat16);
REGISTER_CAST_GPU(uint64_t, bfloat16);
REGISTER_CAST_GPU(Eigen::half, bfloat16);
REGISTER_CAST_GPU(float, bfloat16);
REGISTER_CAST_GPU(double, bfloat16);
REGISTER_CAST_GPU(std::complex<float>, bfloat16);
REGISTER_CAST_GPU(std::complex<double>, bfloat16);
#endif
CURRY_TYPES2(REGISTER_CAST_GPU, bfloat16);
REGISTER_CAST_GPU(float, float8_e5m2);
REGISTER_CAST_GPU(float, float8_e4m3fn);
REGISTER_CAST_GPU(bfloat16, float8_e5m2);
REGISTER_CAST_GPU(bfloat16, float8_e4m3fn);
REGISTER_CAST_GPU(Eigen::half, float8_e5m2);
REGISTER_CAST_GPU(Eigen::half, float8_e4m3fn);
REGISTER_CAST_GPU(float8_e5m2, float);
REGISTER_CAST_GPU(float8_e5m2, bfloat16);
REGISTER_CAST_GPU(float8_e5m2, Eigen::half);
REGISTER_CAST_GPU(float8_e5m2, float8_e5m2);
REGISTER_CAST_GPU(float8_e5m2, float8_e4m3fn);
REGISTER_CAST_GPU(float8_e4m3fn, float);
REGISTER_CAST_GPU(float8_e4m3fn, bfloat16);
REGISTER_CAST_GPU(float8_e4m3fn, Eigen::half);
REGISTER_CAST_GPU(float8_e4m3fn, float8_e5m2);
REGISTER_CAST_GPU(float8_e4m3fn, float8_e4m3fn);
REGISTER_CAST_GPU(int4, int4);
REGISTER_CAST_GPU(int4, int8_t);
REGISTER_CAST_GPU(int4, int16_t);
REGISTER_CAST_GPU(int4, int32_t);
REGISTER_CAST_GPU(int4, int64_t);
REGISTER_CAST_GPU(int4, uint4);
REGISTER_CAST_GPU(int4, uint8_t);
REGISTER_CAST_GPU(int4, uint16_t);
REGISTER_CAST_GPU(int4, uint32_t);
REGISTER_CAST_GPU(int4, uint64_t);
REGISTER_CAST_GPU(int8_t, int4);
REGISTER_CAST_GPU(int16_t, int4);
REGISTER_CAST_GPU(int32_t, int4);
REGISTER_CAST_GPU(int64_t, int4);
REGISTER_CAST_GPU(uint4, int4);
REGISTER_CAST_GPU(uint8_t, int4);
REGISTER_CAST_GPU(uint16_t, int4);
REGISTER_CAST_GPU(uint32_t, int4);
REGISTER_CAST_GPU(uint64_t, int4);
REGISTER_CAST_GPU(uint4, int8_t);
REGISTER_CAST_GPU(uint4, int16_t);
REGISTER_CAST_GPU(uint4, int32_t);
REGISTER_CAST_GPU(uint4, int64_t);
REGISTER_CAST_GPU(uint4, uint4);
REGISTER_CAST_GPU(uint4, uint8_t);
REGISTER_CAST_GPU(uint4, uint16_t);
REGISTER_CAST_GPU(uint4, uint32_t);
REGISTER_CAST_GPU(uint4, uint64_t);
REGISTER_CAST_GPU(int8_t, uint4);
REGISTER_CAST_GPU(int16_t, uint4);
REGISTER_CAST_GPU(int32_t, uint4);
REGISTER_CAST_GPU(int64_t, uint4);
REGISTER_CAST_GPU(uint8_t, uint4);
REGISTER_CAST_GPU(uint16_t, uint4);
REGISTER_CAST_GPU(uint32_t, uint4);
REGISTER_CAST_GPU(uint64_t, uint4);
#undef REGISTER_CAST_GPU
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#undef CURRY_TYPES2
// HostCast differs from Cast in that its input and output are in host memory.
REGISTER_KERNEL_BUILDER(Name("_HostCast").Device(DEVICE_CPU), CpuCastOp);
REGISTER_KERNEL_BUILDER(
Name("_HostCast").Device(DEVICE_DEFAULT).HostMemory("x").HostMemory("y"),
CpuCastOp);
} // end namespace tensorflow
tensorflow/core/kernels/cast_op.h
+351
View File
@@ -0,0 +1,351 @@
/* Copyright 2015 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_CAST_OP_H_
#define TENSORFLOW_CORE_KERNELS_CAST_OP_H_
#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive
#include "tensorflow/core/framework/bfloat16.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/platform/byte_order.h"
#include "tensorflow/core/platform/types.h"
// Note that the GPU cast functor templates need to be instantiated unlike the
// CPU ones, and hence their specializations are different than that for CPUs.
#ifdef SPECIALIZE_FOR_GPUS
#define SPECIALIZE_CAST(DEVICE, OUT_TYPE, IN_TYPE) \
template <typename Device> \
struct CastFunctor<Device, OUT_TYPE, IN_TYPE> { \
void operator()(const Device& d, \
typename TTypes<OUT_TYPE>::Flat out_tensor, \
typename TTypes<IN_TYPE>::ConstFlat in_tensor, \
bool truncate = false) { \
if (truncate) { \
out_tensor.device(d) = \
in_tensor.unaryExpr(LSBZeroSetter<IN_TYPE, OUT_TYPE>()) \
.template cast<OUT_TYPE>(); \
} else { \
out_tensor.device(d) = in_tensor.template cast<OUT_TYPE>(); \
} \
} \
}; \
template struct CastFunctor<DEVICE, OUT_TYPE, IN_TYPE>;
#else
#define SPECIALIZE_CAST(DEVICE, OUT_TYPE, IN_TYPE) \
template <> \
struct CastFunctor<DEVICE, OUT_TYPE, IN_TYPE> { \
void operator()(const DEVICE& d, \
typename TTypes<OUT_TYPE>::Flat out_tensor, \
typename TTypes<IN_TYPE>::ConstFlat in_tensor, \
bool truncate = false) { \
if (truncate) { \
out_tensor.device(d) = \
in_tensor.unaryExpr(LSBZeroSetter<IN_TYPE, OUT_TYPE>()) \
.template cast<OUT_TYPE>(); \
} else { \
out_tensor.device(d) = in_tensor.template cast<OUT_TYPE>(); \
} \
} \
};
#endif
#define CAST_FUNCTORS(devname) \
SPECIALIZE_CAST(devname, float, double) \
SPECIALIZE_CAST(devname, float, std::complex<double>) \
SPECIALIZE_CAST(devname, std::complex<float>, std::complex<double>) \
SPECIALIZE_CAST(devname, std::complex<float>, double) \
SPECIALIZE_CAST(devname, Eigen::half, double) \
SPECIALIZE_CAST(devname, Eigen::half, float) \
SPECIALIZE_CAST(devname, Eigen::half, std::complex<double>) \
SPECIALIZE_CAST(devname, Eigen::half, std::complex<float>) \
SPECIALIZE_CAST(devname, bfloat16, float) \
SPECIALIZE_CAST(devname, float8_e5m2, double) \
SPECIALIZE_CAST(devname, float8_e5m2, float) \
SPECIALIZE_CAST(devname, float8_e5m2, bfloat16) \
SPECIALIZE_CAST(devname, float8_e5m2, Eigen::half) \
SPECIALIZE_CAST(devname, float8_e5m2, float8_e4m3fn) \
SPECIALIZE_CAST(devname, float8_e4m3fn, double) \
SPECIALIZE_CAST(devname, float8_e4m3fn, float) \
SPECIALIZE_CAST(devname, float8_e4m3fn, bfloat16) \
SPECIALIZE_CAST(devname, float8_e4m3fn, Eigen::half) \
template <typename OUT_TYPE, typename IN_TYPE> \
struct CastFunctor<devname, OUT_TYPE, IN_TYPE> { \
void operator()(const devname& d, \
typename TTypes<OUT_TYPE>::Flat out_tensor, \
typename TTypes<IN_TYPE>::ConstFlat in_tensor, \
bool truncate = false) { \
out_tensor.device(d) = in_tensor.template cast<OUT_TYPE>(); \
} \
};
#if defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
// If MLIR kernels are enabled, we don't need the specialized cast from float to
// double or from Eigen::half to double. We still need the specialized cast from
// Eigen::half to float, because it is used in depthwise_conv_grad_op.cc. We
// still need the specialized cast from float to double because it is used in
// resize_bilinear_op.cc.
#define CAST_FUNCTORS_SUBSET(devname) \
SPECIALIZE_CAST(devname, float, double) \
SPECIALIZE_CAST(devname, Eigen::half, float) \
SPECIALIZE_CAST(devname, bfloat16, float) \
SPECIALIZE_CAST(devname, float8_e5m2, double) \
SPECIALIZE_CAST(devname, float8_e5m2, float) \
SPECIALIZE_CAST(devname, float8_e5m2, bfloat16) \
SPECIALIZE_CAST(devname, float8_e5m2, Eigen::half) \
SPECIALIZE_CAST(devname, float8_e5m2, float8_e4m3fn) \
SPECIALIZE_CAST(devname, float8_e4m3fn, double) \
SPECIALIZE_CAST(devname, float8_e4m3fn, float) \
SPECIALIZE_CAST(devname, float8_e4m3fn, bfloat16) \
SPECIALIZE_CAST(devname, float8_e4m3fn, Eigen::half) \
template <typename OUT_TYPE, typename IN_TYPE> \
struct CastFunctor<devname, OUT_TYPE, IN_TYPE> { \
void operator()(const devname& d, \
typename TTypes<OUT_TYPE>::Flat out_tensor, \
typename TTypes<IN_TYPE>::ConstFlat in_tensor, \
bool truncate = false) { \
out_tensor.device(d) = in_tensor.template cast<OUT_TYPE>(); \
} \
};
#endif
namespace tensorflow {
typedef std::function<void(OpKernelContext*, const Tensor&, Tensor*,
bool trunc)>
CastFunctorType;
// Common base class of Cast kernels
class CastOpBase : public OpKernel {
public:
explicit CastOpBase(OpKernelConstruction* ctx);
void Compute(OpKernelContext* ctx) override;
protected:
DataType src_dtype_;
DataType dst_dtype_;
DataType external_src_dtype_;
DataType external_dst_dtype_;
bool use_truncation_;
CastFunctorType work_ = nullptr;
absl::Status Unimplemented();
CastOpBase(const CastOpBase&) = delete;
void operator=(const CastOpBase&) = delete;
};
// CPU implementation of Cast
class CpuCastOp : public CastOpBase {
public:
explicit CpuCastOp(OpKernelConstruction* ctx);
private:
absl::Status Prepare();
};
namespace functor {
template <typename I>
constexpr int MantissaWidth() {
return std::numeric_limits<I>::digits;
}
template <>
constexpr int MantissaWidth<Eigen::half>() {
// Remember, there's 1 hidden bit
return 10 + 1;
}
template <>
constexpr int MantissaWidth<bfloat16>() {
// Remember, there's 1 hidden bit
return 7 + 1;
}
template <typename Device, typename Tout, typename Tin>
void Cast(const Device& d, typename TTypes<Tout>::Flat o,
typename TTypes<Tin>::ConstFlat i) {
o.device(d) = i.template cast<Tout>();
}
template <typename Device, typename Tout, typename Tin>
struct CastFunctor {
void operator()(const Device& d, typename TTypes<Tout>::Flat o,
typename TTypes<Tin>::ConstFlat i, bool truncate = false);
};
template <typename I>
typename std::enable_if<sizeof(I) == 8, void>::type EIGEN_DEVICE_FUNC
EIGEN_STRONG_INLINE static LSBZeroSetterHelper(I& t, int n) {
// Only zero the bits for non-NaNs.
// For NaNs, let the non-truncation version handle it.
if (!Eigen::numext::isnan(t)) {
uint64_t* p = reinterpret_cast<uint64_t*>(&t);
*p &= (0xFFFFFFFFFFFFFFFF << n);
}
}
template <typename I>
typename std::enable_if<sizeof(I) == 4, void>::type EIGEN_DEVICE_FUNC
EIGEN_STRONG_INLINE static LSBZeroSetterHelper(I& t, int n) {
// Only zero the bits for non-NaNs.
// For NaNs, let the non-truncation version handle it.
if (!Eigen::numext::isnan(t)) {
uint32_t* p = reinterpret_cast<uint32_t*>(&t);
*p &= (0xFFFFFFFF << n);
}
}
template <typename I>
typename std::enable_if<sizeof(I) == 2, void>::type EIGEN_DEVICE_FUNC
EIGEN_STRONG_INLINE static LSBZeroSetterHelper(I& t, int n) {
// Only zero the bits for non-NaNs.
// For NaNs, let the non-truncation version handle it.
if (!Eigen::numext::isnan(t)) {
uint16_t* p = reinterpret_cast<uint16_t*>(&t);
*p &= (0xFFFF << n);
}
}
template <typename I>
typename std::enable_if<sizeof(I) == 1, void>::type EIGEN_DEVICE_FUNC
EIGEN_STRONG_INLINE static LSBZeroSetterHelper(I& t, int n) {
// Only zero the bits for non-NaNs.
// For NaNs, let the non-truncation version handle it.
if (!Eigen::numext::isnan(t)) {
uint8_t* p = reinterpret_cast<uint8_t*>(&t);
*p &= (0xFF << n);
}
}
// Set n least significant bits to 0
template <typename I, typename O>
struct LSBZeroSetter {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE I operator()(const I& a) const {
constexpr int bits = MantissaWidth<I>() - MantissaWidth<O>();
static_assert(
bits > 0,
"The output type must have fewer mantissa bits than the input type\n");
I t = a;
LSBZeroSetterHelper(t, bits);
return t;
}
};
template <typename I, typename O>
struct LSBZeroSetter<std::complex<I>, std::complex<O>> {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<I> operator()(
const std::complex<I>& a) const {
constexpr int bits = MantissaWidth<I>() - MantissaWidth<O>();
static_assert(
bits > 0,
"The output type must have fewer mantissa bits than the input type\n");
I re = Eigen::numext::real(a);
I img = Eigen::numext::imag(a);
LSBZeroSetterHelper(re, bits);
LSBZeroSetterHelper(img, bits);
std::complex<I> toReturn(re, img);
return toReturn;
}
};
template <typename I, typename O>
struct LSBZeroSetter<std::complex<I>, O> {
// Sets the 16 LSBits of the float to 0
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<I> operator()(
const std::complex<I>& a) const {
constexpr int bits = MantissaWidth<I>() - MantissaWidth<O>();
static_assert(
bits > 0,
"The output type must have fewer mantissa bits than the input type\n");
I re = Eigen::numext::real(a);
I img = Eigen::numext::imag(a);
LSBZeroSetterHelper(re, bits);
LSBZeroSetterHelper(img, bits);
std::complex<I> toReturn(re, img);
return toReturn;
}
};
} // end namespace functor
} // end namespace tensorflow
namespace Eigen {
namespace internal {
// Eigen can't convert to/from complex numbers, because it is limited to cases
// that can be static_casted. But numpy is able to cast to/from complex, which
// we want to replicate. So we add specializations for complex here.
template <typename From, typename To>
struct scalar_cast_op<std::complex<From>, To> {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE To
operator()(const std::complex<From>& a) const {
// Replicate numpy behavior of returning just the real part
return static_cast<To>(a.real());
}
};
template <typename From>
struct scalar_cast_op<std::complex<From>, bool> {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool operator()(
const std::complex<From>& a) const {
return static_cast<bool>(a.real());
}
};
template <typename From, typename To>
struct scalar_cast_op<From, std::complex<To>> {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<To> operator()(
const From& a) const {
// Replicate numpy behavior of setting the imaginary part to 0
return std::complex<To>(static_cast<To>(a), To(0));
}
};
template <typename From, typename To>
struct scalar_cast_op<std::complex<From>, std::complex<To>> {
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<To> operator()(
const std::complex<From>& a) const {
return std::complex<To>(static_cast<To>(a.real()),
static_cast<To>(a.imag()));
}
};
template <typename From, typename To>
struct functor_traits_complex_impl {
enum { Cost = NumTraits<To>::AddCost, PacketAccess = false };
};
template <typename From>
struct functor_traits<scalar_cast_op<std::complex<From>, bool>>
: functor_traits_complex_impl<std::complex<From>, bool> {};
template <typename From, typename To>
struct functor_traits<scalar_cast_op<std::complex<From>, To>>
: functor_traits_complex_impl<std::complex<From>, To> {};
template <typename From, typename To>
struct functor_traits<scalar_cast_op<From, std::complex<To>>>
: functor_traits_complex_impl<From, std::complex<To>> {};
// Needed to avoid ambiguous partial specialization
template <typename From, typename To>
struct functor_traits<scalar_cast_op<std::complex<From>, std::complex<To>>>
: functor_traits_complex_impl<std::complex<From>, std::complex<To>> {};
} // namespace internal
} // namespace Eigen
#endif // TENSORFLOW_CORE_KERNELS_CAST_OP_H_
tensorflow/core/kernels/cast_op_gpu.cu.cc
+191
View File
@@ -0,0 +1,191 @@
/* Copyright 2015 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.
==============================================================================*/
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
#define EIGEN_USE_GPU
#include "tensorflow/core/framework/bfloat16.h"
#define SPECIALIZE_FOR_GPUS
#include "tensorflow/core/kernels/cast_op.h"
#undef SPECIALIZE_FOR_GPUS
namespace tensorflow {
namespace functor {
typedef Eigen::GpuDevice GPUDevice;
#if defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
CAST_FUNCTORS_SUBSET(GPUDevice);
#else
CAST_FUNCTORS(GPUDevice);
#endif
#define DEFINE(O, I) template struct CastFunctor<GPUDevice, O, I>
#define DEFINE_ALL_FROM(in_type) \
DEFINE(in_type, bool); \
DEFINE(in_type, uint8); \
DEFINE(in_type, uint16); \
DEFINE(in_type, uint32); \
DEFINE(in_type, uint64); \
DEFINE(in_type, int8); \
DEFINE(in_type, int16); \
DEFINE(in_type, int32); \
DEFINE(in_type, int64); \
DEFINE(in_type, Eigen::half); \
DEFINE(in_type, float); \
DEFINE(in_type, double); \
DEFINE(in_type, std::complex<float>); \
DEFINE(in_type, std::complex<double>)
// Required functors not previously specialized for truncation.
DEFINE(double, float8_e5m2);
DEFINE(float, float8_e5m2);
DEFINE(bfloat16, float8_e5m2);
DEFINE(Eigen::half, float8_e5m2);
DEFINE(float8_e5m2, float8_e5m2);
DEFINE(float8_e4m3fn, float8_e5m2);
DEFINE(double, float8_e4m3fn);
DEFINE(float, float8_e4m3fn);
DEFINE(bfloat16, float8_e4m3fn);
DEFINE(Eigen::half, float8_e4m3fn);
DEFINE(float8_e4m3fn, float8_e4m3fn);
#if defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
// The cast from float to double is still needed for resize_bilinear_op.cc
DEFINE(double, float);
#else
DEFINE_ALL_FROM(bool);
DEFINE_ALL_FROM(uint8);
DEFINE_ALL_FROM(uint16);
DEFINE_ALL_FROM(uint32);
DEFINE_ALL_FROM(uint64);
DEFINE_ALL_FROM(int8);
DEFINE_ALL_FROM(int16);
DEFINE_ALL_FROM(int32);
DEFINE_ALL_FROM(int64);
DEFINE_ALL_FROM(double);
DEFINE_ALL_FROM(std::complex<double>);
#endif
#define DEFINE_ALL_TO_FLOAT(out_type) \
DEFINE(out_type, bool); \
DEFINE(out_type, uint8); \
DEFINE(out_type, uint16); \
DEFINE(out_type, uint32); \
DEFINE(out_type, uint64); \
DEFINE(out_type, int8); \
DEFINE(out_type, int16); \
DEFINE(out_type, int32); \
DEFINE(out_type, int64); \
DEFINE(out_type, Eigen::half); \
DEFINE(out_type, bfloat16); \
DEFINE(out_type, float); \
DEFINE(out_type, std::complex<float>)
#define DEFINE_ALL_TO_HALF(out_type) \
DEFINE(out_type, bool); \
DEFINE(out_type, uint8); \
DEFINE(out_type, uint16); \
DEFINE(out_type, uint32); \
DEFINE(out_type, uint64); \
DEFINE(out_type, int8); \
DEFINE(out_type, int16); \
DEFINE(out_type, int32); \
DEFINE(out_type, int64); \
DEFINE(out_type, Eigen::half); \
DEFINE(out_type, bfloat16)
DEFINE_ALL_TO_HALF(bfloat16);
DEFINE(bool, bfloat16);
DEFINE(uint8_t, bfloat16);
DEFINE(uint16_t, bfloat16);
DEFINE(uint32_t, bfloat16);
DEFINE(uint64_t, bfloat16);
DEFINE(int8_t, bfloat16);
DEFINE(int16_t, bfloat16);
DEFINE(int32_t, bfloat16);
DEFINE(int64_t, bfloat16);
DEFINE(std::complex<double>, bfloat16);
DEFINE(double, bfloat16);
DEFINE(bfloat16, std::complex<float>);
DEFINE(bfloat16, std::complex<double>);
DEFINE(bfloat16, double);
#if defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
// The cast from Eigen::half is still needed for depthwise_conv_grad_op.cc.
DEFINE(float, Eigen::half);
// The cast from float to float is still needed for resize_bilinear_op.cc.
DEFINE(float, float);
// The casts from complex to the complex element type is still needed for
// self_adjoint_eig_v2_op_gpu.cc
DEFINE(std::complex<float>, float);
DEFINE(std::complex<double>, double);
DEFINE(float, bfloat16);
DEFINE(Eigen::half, bfloat16);
DEFINE(std::complex<float>, bfloat16);
#else
DEFINE_ALL_TO_HALF(Eigen::half);
DEFINE_ALL_TO_FLOAT(float);
DEFINE_ALL_TO_FLOAT(std::complex<float>);
#endif
#define DEFINE_TO_INT(from_type) \
DEFINE(from_type, int8); \
DEFINE(from_type, int16); \
DEFINE(from_type, int32); \
DEFINE(from_type, int64_t); \
DEFINE(from_type, uint8); \
DEFINE(from_type, uint16); \
DEFINE(from_type, uint32); \
DEFINE(from_type, uint64)
#define DEFINE_FROM_INT(out_type) \
DEFINE(int8, out_type); \
DEFINE(int16, out_type); \
DEFINE(int32, out_type); \
DEFINE(int64_t, out_type); \
DEFINE(uint8, out_type); \
DEFINE(uint16, out_type); \
DEFINE(uint32, out_type); \
DEFINE(uint64, out_type)
DEFINE_TO_INT(int4);
DEFINE_TO_INT(uint4);
DEFINE_FROM_INT(int4);
DEFINE_FROM_INT(uint4);
DEFINE(int4, int4);
DEFINE(int4, uint4);
DEFINE(uint4, int4);
DEFINE(uint4, uint4);
#undef DEFINE_TO_INT
#undef DEFINE_FROM_INT
#undef DEFINE_ALL_TO_FLOAT
#undef DEFINE_ALL_TO_HALF
#undef DEFINE_ALL_FROM
#undef DEFINE
} // end namespace functor
} // end namespace tensorflow
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
tensorflow/core/kernels/cast_op_impl.h
+189
View File
@@ -0,0 +1,189 @@
/* Copyright 2016 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_KERNELS_CAST_OP_IMPL_H_
#define TENSORFLOW_CORE_KERNELS_CAST_OP_IMPL_H_
#include <cstdint>
#include <limits>
#define EIGEN_USE_THREADS
#include "absl/status/status.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/kernels/cast_op.h"
#include "tensorflow/core/platform/errors.h"
namespace tensorflow {
namespace functor {
template <class F, class I>
struct OutOfRange {
bool operator()(const F f) const {
return f < std::numeric_limits<I>::min() ||
f > std::numeric_limits<I>::max();
}
};
#define VALIDATE_CAST(I, F) \
template <> \
struct CastFunctor<Eigen::ThreadPoolDevice, I, F> { \
void operator()(const Eigen::ThreadPoolDevice& d, \
typename TTypes<I>::Flat out_tensor, \
typename TTypes<F>::ConstFlat in_tensor, \
bool truncate = false) const { \
Eigen::Tensor<bool, 0, Eigen::RowMajor> out_of_range = \
in_tensor.unaryExpr(OutOfRange<F, I>{}).any(); \
if (out_of_range()) { \
LOG(ERROR) \
<< "IMPORTANT! The input tensor to Cast contains values out of " \
"range for the target type. This is undefined behavior and " \
"likely a bug in your model. A crash immediately after this " \
"under ubsan is expected."; \
} \
out_tensor.device(d) = in_tensor.template cast<I>(); \
} \
};
// Add additional logging for out of range inputs when running in debug mode.
#ifndef NDEBUG
VALIDATE_CAST(int32_t, float);
VALIDATE_CAST(int64_t, float);
VALIDATE_CAST(int32_t, double);
VALIDATE_CAST(int64_t, double);
#endif
CAST_FUNCTORS(Eigen::ThreadPoolDevice);
} // namespace functor
#define CURRY_TYPES3(FN, arg0, arg1) \
FN(arg0, arg1, bool); \
FN(arg0, arg1, uint8_t); \
FN(arg0, arg1, uint16_t); \
FN(arg0, arg1, uint32_t); \
FN(arg0, arg1, uint64_t); \
FN(arg0, arg1, int8_t); \
FN(arg0, arg1, int16_t); \
FN(arg0, arg1, int32_t); \
FN(arg0, arg1, int64_t); \
FN(arg0, arg1, float); \
FN(arg0, arg1, double); \
FN(arg0, arg1, std::complex<float>); \
FN(arg0, arg1, std::complex<double>) \
FN(arg0, arg1, Eigen::half); \
FN(arg0, arg1, bfloat16);
#define CAST_CASE(DEVICE, IN, OUT) \
if (DataTypeToEnum<OUT>::value == dst_dtype) { \
return [](OpKernelContext* ctx, const Tensor& inp, Tensor* out, \
bool truncate) { \
functor::CastFunctor<DEVICE, OUT, IN> func; \
func(ctx->eigen_device<DEVICE>(), out->flat<OUT>(), inp.flat<IN>(), \
truncate); \
}; \
}
// The functions below are implemented in the cast_op_impl_*.cc files.
CastFunctorType GetCpuCastFromBool(DataType dst_dtype);
CastFunctorType GetCpuCastFromUint8(DataType dst_dtype);
CastFunctorType GetCpuCastFromUint16(DataType dst_dtype);
CastFunctorType GetCpuCastFromInt8(DataType dst_dtype);
CastFunctorType GetCpuCastFromUint32(DataType dst_dtype);
CastFunctorType GetCpuCastFromUint64(DataType dst_dtype);
CastFunctorType GetCpuCastFromInt8(DataType dst_dtype);
CastFunctorType GetCpuCastFromInt16(DataType dst_dtype);
CastFunctorType GetCpuCastFromInt32(DataType dst_dtype);
CastFunctorType GetCpuCastFromInt64(DataType dst_dtype);
CastFunctorType GetCpuCastFromHalf(DataType dst_dtype);
CastFunctorType GetCpuCastFromFloat(DataType dst_dtype);
CastFunctorType GetCpuCastFromDouble(DataType dst_dtype);
CastFunctorType GetCpuCastFromComplex64(DataType dst_dtype);
CastFunctorType GetCpuCastFromComplex128(DataType dst_dtype);
CastFunctorType GetCpuCastFromBfloat(DataType dst_dtype);
CastFunctorType GetCpuCastFromFloat8e5m2(DataType dst_dtype);
CastFunctorType GetCpuCastFromFloat8e4m3fn(DataType dst_dtype);
CastFunctorType GetCpuCastFromInt4(DataType dst_dtype);
CastFunctorType GetCpuCastFromUint4(DataType dst_dtype);
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
// Same, for GPU.
CastFunctorType GetGpuCastFromBool(DataType dst_dtype);
CastFunctorType GetGpuCastFromUint8(DataType dst_dtype);
CastFunctorType GetGpuCastFromUint16(DataType dst_dtype);
CastFunctorType GetGpuCastFromInt8(DataType dst_dtype);
CastFunctorType GetGpuCastFromUint32(DataType dst_dtype);
CastFunctorType GetGpuCastFromUint64(DataType dst_dtype);
CastFunctorType GetGpuCastFromInt16(DataType dst_dtype);
CastFunctorType GetGpuCastFromInt32(DataType dst_dtype);
CastFunctorType GetGpuCastFromInt64(DataType dst_dtype);
CastFunctorType GetGpuCastFromHalf(DataType dst_dtype);
CastFunctorType GetGpuCastFromFloat(DataType dst_dtype);
CastFunctorType GetGpuCastFromDouble(DataType dst_dtype);
CastFunctorType GetGpuCastFromComplex64(DataType dst_dtype);
CastFunctorType GetGpuCastFromComplex128(DataType dst_dtype);
CastFunctorType GetGpuCastFromBfloat(DataType dst_dtype);
CastFunctorType GetGpuCastFromFloat8e5m2(DataType dst_dtype);
CastFunctorType GetGpuCastFromFloat8e4m3fn(DataType dst_dtype);
CastFunctorType GetGpuCastFromInt4(DataType dst_dtype);
CastFunctorType GetGpuCastFromUint4(DataType dst_dtype);
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_CAST_OP_IMPL_H_
tensorflow/core/kernels/cast_op_impl_bfloat.cc
+42
View File
@@ -0,0 +1,42 @@
/* Copyright 2016 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/cast_op_impl.h"
#include "tensorflow/core/util/work_sharder.h"
namespace tensorflow {
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
CastFunctorType GetCpuCastFromBfloat(DataType dst_dtype) {
CURRY_TYPES3(CAST_CASE, CPUDevice, bfloat16);
CAST_CASE(CPUDevice, bfloat16, float8_e5m2);
CAST_CASE(CPUDevice, bfloat16, float8_e4m3fn);
return nullptr;
}
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
CastFunctorType GetGpuCastFromBfloat(DataType dst_dtype) {
CURRY_TYPES3(CAST_CASE, GPUDevice, bfloat16);
CAST_CASE(GPUDevice, bfloat16, float8_e5m2);
CAST_CASE(GPUDevice, bfloat16, float8_e4m3fn);
return nullptr;
}
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow
tensorflow/core/kernels/cast_op_impl_bool.cc
+41
View File
@@ -0,0 +1,41 @@
/* Copyright 2016 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/cast_op_impl.h"
namespace tensorflow {
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
CastFunctorType GetCpuCastFromBool(DataType dst_dtype) {
CURRY_TYPES3(CAST_CASE, CPUDevice, bool);
return nullptr;
}
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
CastFunctorType GetGpuCastFromBool(DataType dst_dtype) {
#if defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
CAST_CASE(GPUDevice, bool, bfloat16);
#else
CURRY_TYPES3(CAST_CASE, GPUDevice, bool);
#endif
return nullptr;
}
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow
tensorflow/core/kernels/cast_op_impl_complex128.cc
+40
View File
@@ -0,0 +1,40 @@
/* Copyright 2016 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/cast_op_impl.h"
namespace tensorflow {
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
CastFunctorType GetCpuCastFromComplex128(DataType dst_dtype) {
CURRY_TYPES3(CAST_CASE, CPUDevice, std::complex<double>);
return nullptr;
}
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
CastFunctorType GetGpuCastFromComplex128(DataType dst_dtype) {
#if defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
CAST_CASE(GPUDevice, std::complex<double>, bfloat16);
#else
CURRY_TYPES3(CAST_CASE, GPUDevice, std::complex<double>);
#endif
return nullptr;
}
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow
tensorflow/core/kernels/cast_op_impl_complex64.cc
+40
View File
@@ -0,0 +1,40 @@
/* Copyright 2016 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/cast_op_impl.h"
namespace tensorflow {
typedef Eigen::ThreadPoolDevice CPUDevice;
typedef Eigen::GpuDevice GPUDevice;
CastFunctorType GetCpuCastFromComplex64(DataType dst_dtype) {
CURRY_TYPES3(CAST_CASE, CPUDevice, std::complex<float>);
return nullptr;
}
#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
(defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
CastFunctorType GetGpuCastFromComplex64(DataType dst_dtype) {
#if defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
CAST_CASE(GPUDevice, std::complex<float>, bfloat16);
#else
CURRY_TYPES3(CAST_CASE, GPUDevice, std::complex<float>);
#endif
return nullptr;
}
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // namespace tensorflow

Some files were not shown because too many files have changed in this diff Show More