diff options
author | A. Unique TensorFlower <gardener@tensorflow.org> | 2018-07-25 08:23:57 -0700 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-07-25 08:27:36 -0700 |
commit | b3771feab49e2122164737a860341727d08c2d8c (patch) | |
tree | 5fb440041db26ef96eb14e7491cb67fe06e7c3d4 /tensorflow/core/kernels/cast_op_impl_bfloat.cc | |
parent | be3d22844025e42e177a21479f3ae73bc5351c1f (diff) |
This change started with an intention of adding an attribute to cast ops to decide
whether bfloat16 casts should use truncation or rounding.
This is a preparatory change before we switch the default float ==> bfloat16 cast
to use rounding instead of truncation. The attribute added can then be specified
on casts that rely on the truncation, e.g., the TensorFlow send/receive operations.
It later emerged that the choice of doing truncation is useful more generally.
Therefore, this change allows the new attribute to be used by all relevant casts
to use truncation instead of rounding.
PiperOrigin-RevId: 205996367
Diffstat (limited to 'tensorflow/core/kernels/cast_op_impl_bfloat.cc')
-rw-r--r-- | tensorflow/core/kernels/cast_op_impl_bfloat.cc | 11 |
1 files changed, 5 insertions, 6 deletions
diff --git a/tensorflow/core/kernels/cast_op_impl_bfloat.cc b/tensorflow/core/kernels/cast_op_impl_bfloat.cc index bfa7ba0d47..96aae15608 100644 --- a/tensorflow/core/kernels/cast_op_impl_bfloat.cc +++ b/tensorflow/core/kernels/cast_op_impl_bfloat.cc @@ -22,20 +22,19 @@ namespace tensorflow { typedef Eigen::ThreadPoolDevice CPUDevice; typedef Eigen::GpuDevice GPUDevice; -std::function<void(OpKernelContext*, const Tensor&, Tensor*)> -GetCpuCastFromBfloat(DataType dst_dtype) { +CastFunctorType GetCpuCastFromBfloat(DataType dst_dtype) { CURRY_TYPES3(CAST_CASE, CPUDevice, bfloat16); return nullptr; } #if GOOGLE_CUDA -std::function<void(OpKernelContext*, const Tensor&, Tensor*)> -GetGpuCastFromBfloat(DataType dst_dtype) { +CastFunctorType GetGpuCastFromBfloat(DataType dst_dtype) { if (dst_dtype == DT_FLOAT) { - return [](OpKernelContext* ctx, const Tensor& inp, Tensor* out) { + return [](OpKernelContext* ctx, const Tensor& inp, Tensor* out, + bool truncate) { functor::CastFunctor<GPUDevice, float, bfloat16> func; func(ctx->eigen_device<GPUDevice>(), out->flat<float>(), - inp.flat<bfloat16>()); + inp.flat<bfloat16>(), truncate); }; } return nullptr; |