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

Merge pull request from ROCmSoftwarePlatform:google_upstream_rocm_platform_fix_190517

Browse Source 
PiperOrigin-RevId: 249244017
This commit is contained in:
TensorFlower Gardener 2019-05-21 07:47:06 -07:00
commit 2fad0483f2
2 changed files with 20 additions and 5 deletions
tensorflow/stream_executor/rocm
rocm_dnn.ccrocm_dnn.h

12
tensorflow/stream_executor/rocm/rocm_dnn.cc
View File

@ -3004,6 +3004,8 @@ bool MIOpenSupport::DoBatchNormalizationForward(
DeviceMemory<Eigen::half>* y, DeviceMemory<float>* batch_mean,
DeviceMemory<float>* batch_var, DeviceMemory<float>* saved_mean,
DeviceMemory<float>* saved_inv_var, bool is_training,
ScratchAllocator* reserve_space_allocator,
ScratchAllocator* workspace_allocator,
std::function<const DeviceMemory<float>&()> var_to_inv_var,
std::function<void()> inv_var_to_var) {
return DoBatchNormalizationForwardImpl<Eigen::half, float>(
@ -3023,6 +3025,8 @@ bool MIOpenSupport::DoBatchNormalizationForward(
DeviceMemory<float>* y, DeviceMemory<float>* batch_mean,
DeviceMemory<float>* batch_var, DeviceMemory<float>* saved_mean,
DeviceMemory<float>* saved_inv_var, bool is_training,
ScratchAllocator* reserve_space_allocator,
ScratchAllocator* workspace_allocator,
std::function<const DeviceMemory<float>&()> var_to_inv_var,
std::function<void()> inv_var_to_var) {
return DoBatchNormalizationForwardImpl<float, float>(
@ -3089,7 +3093,9 @@ bool MIOpenSupport::DoBatchNormalizationBackward(
const dnn::BatchDescriptor& x_desc,
const dnn::BatchDescriptor& scale_offset_desc, const double epsilon,
DeviceMemory<Eigen::half>* x_backprop, DeviceMemory<float>* scale_backprop,
DeviceMemory<float>* offset_backprop) {
DeviceMemory<float>* offset_backprop,
DeviceMemory<uint8>* reserve_space_data,
ScratchAllocator* workspace_allocator) {
return DoBatchNormalizationBackwardImpl<Eigen::half, float>(
stream, miopenHalf, miopenFloat, y_backprop, x, scale, mean, inv_var,
x_desc, scale_offset_desc, epsilon, x_backprop, scale_backprop,
@ -3103,7 +3109,9 @@ bool MIOpenSupport::DoBatchNormalizationBackward(
const dnn::BatchDescriptor& x_desc,
const dnn::BatchDescriptor& scale_offset_desc, const double epsilon,
DeviceMemory<float>* x_backprop, DeviceMemory<float>* scale_backprop,
DeviceMemory<float>* offset_backprop) {
DeviceMemory<float>* offset_backprop,
DeviceMemory<uint8>* reserve_space_data,
ScratchAllocator* workspace_allocator) {
return DoBatchNormalizationBackwardImpl<float, float>(
stream, miopenFloat, miopenFloat, y_backprop, x, scale, mean, variance,
x_desc, scale_offset_desc, epsilon, x_backprop, scale_backprop,

13
tensorflow/stream_executor/rocm/rocm_dnn.h
View File

@ -215,6 +215,8 @@ class MIOpenSupport : public dnn::DnnSupport {
DeviceMemory<float>* y, DeviceMemory<float>* batch_mean,
DeviceMemory<float>* batch_var, DeviceMemory<float>* saved_mean,
DeviceMemory<float>* saved_inv_var, bool is_training,
ScratchAllocator* reserve_space_allocator,
ScratchAllocator* workspace_allocator,
std::function<const DeviceMemory<float>&()> var_to_inv_var,
std::function<void()> inv_var_to_var) override;
@ -228,6 +230,8 @@ class MIOpenSupport : public dnn::DnnSupport {
DeviceMemory<Eigen::half>* y, DeviceMemory<float>* batch_mean,
DeviceMemory<float>* batch_var, DeviceMemory<float>* saved_mean,
DeviceMemory<float>* saved_inv_var, bool is_training,
ScratchAllocator* reserve_space_allocator,
ScratchAllocator* workspace_allocator,
std::function<const DeviceMemory<float>&()> var_to_inv_var,
std::function<void()> inv_var_to_var) override;
@ -238,7 +242,9 @@ class MIOpenSupport : public dnn::DnnSupport {
const dnn::BatchDescriptor& x_desc,
const dnn::BatchDescriptor& scale_offset_desc, const double epsilon,
DeviceMemory<float>* x_backprop, DeviceMemory<float>* scale_backprop,
DeviceMemory<float>* offset_backprop) override;
DeviceMemory<float>* offset_backprop,
DeviceMemory<uint8>* reserve_space_data,
ScratchAllocator* workspace_allocator) override;
bool DoBatchNormalizationBackward(
Stream* stream, const DeviceMemory<Eigen::half>& y_backprop,
@ -247,8 +253,9 @@ class MIOpenSupport : public dnn::DnnSupport {
const dnn::BatchDescriptor& x_desc,
const dnn::BatchDescriptor& scale_offset_desc, const double epsilon,
DeviceMemory<Eigen::half>* x_backprop,
DeviceMemory<float>* scale_backprop,
DeviceMemory<float>* offset_backprop) override;
DeviceMemory<float>* scale_backprop, DeviceMemory<float>* offset_backprop,
DeviceMemory<uint8>* reserve_space_data,
ScratchAllocator* workspace_allocator) override;
port::Status DoConvolve(
dnn::ConvolutionKind kind, dnn::DataType element_type, Stream* stream,