/* 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/stream_executor/stream.h"

#include "tensorflow/stream_executor/platform/port.h"

#include "tensorflow/stream_executor/blas.h"

#include "tensorflow/stream_executor/host_buffer.h"

#include "tensorflow/stream_executor/lib/stacktrace.h"

#include "tensorflow/stream_executor/lib/strcat.h"

#include "tensorflow/stream_executor/platform.h"

#include "tensorflow/stream_executor/platform/logging.h"

#include "tensorflow/stream_executor/rng.h"

#include "tensorflow/stream_executor/stream_executor_internal.h"

#include "tensorflow/stream_executor/stream_executor_pimpl.h"

namespace perftools {

namespace gputools {

namespace {

// Code to turn parameters to functions on stream into strings that

// will be VLOG'ed. We need overloads, instead of

// e.g. BatchDescriptorToVlogString(), as the code that calls these

// functions does not know what the type of the parameter is.

string ToVlogString(const dnn::BatchDescriptor &descriptor) {

return descriptor.ToShortString();

}

string ToVlogString(const dnn::FilterDescriptor &descriptor) {

return descriptor.ToShortString();

}

string ToVlogString(const dnn::ConvolutionDescriptor &descriptor) {

return descriptor.ToShortString();

}

string ToVlogString(const dnn::PoolingDescriptor &descriptor) {

return descriptor.ToShortString();

}

string ToVlogString(const dnn::NormalizeDescriptor &descriptor) {

return descriptor.ToShortString();

}

string ToVlogString(dnn::ActivationMode mode) {

return dnn::ActivationModeString(mode);

}

string ToVlogString(dnn::ElementwiseOperation op) {

return dnn::ElementwiseOperationString(op);

}

string ToVlogString(dnn::QuantizedActivationMode mode) {

return dnn::QuantizedActivationModeString(mode);

}

string ToVlogString(blas::Transpose t) { return blas::TransposeString(t); }

string ToVlogString(blas::UpperLower ul) { return blas::UpperLowerString(ul); }

string ToVlogString(blas::Diagonal d) { return blas::DiagonalString(d); }

string ToVlogString(blas::Side s) { return blas::SideString(s); }

string ToVlogString(const void *ptr) {

if (ptr == nullptr) {

return "null";

}

// StrCat does not convert pointers to text.

std::ostringstream out;

out << ptr;

return out.str();

}

string ToVlogString(const HostBuffer &buffer) { return buffer.AsString(); }

template <class T>

string ToVlogString(const std::complex<T> &c) {

// StrCat does not convert std::complex to text.

std::ostringstream out;

out << c;

return out.str();

}

template <class T>

string ToVlogString(const std::function<T> &f) {

return f == nullptr ? "null" : "<non-null function>";

}

string ToVlogString(const DeviceMemoryBase &memory) {

return ToVlogString(memory.opaque());

}

string ToVlogString(const DeviceMemoryBase *memory) {

return ToVlogString(*memory);

}

string ToVlogString(int i) { return port::StrCat(i); }

string ToVlogString(uint32 i) { return port::StrCat(i); }

string ToVlogString(uint64 i) { return port::StrCat(i); }

string ToVlogString(int64 i) { return port::StrCat(i); }

string ToVlogString(float f) { return port::StrCat(f); }

string ToVlogString(double d) { return port::StrCat(d); }

template <class T>

string ToVlogString(port::ArraySlice<T> elements) {

string str = port::StrCat(

ToVlogString(reinterpret_cast<const void *>(elements.data())), "[",

elements.size(), "]{");

const char *separator = "";

size_t max_to_show = std::numeric_limits<size_t>::max();

if (!VLOG_IS_ON(2)) {

max_to_show = 5;

} else if (!VLOG_IS_ON(3)) {

max_to_show = 20;

} else if (!VLOG_IS_ON(11)) {

max_to_show = 1000;

}

for (size_t i = 0; i < elements.size(); ++i) {

if (i == max_to_show) {

str += ", ...";

break;

}

port::StrAppend(&str, separator, ToVlogString(elements[i]));

separator = ", ";

}

str += "}";

return str;

}

template <class T>

string ToVlogString(port::MutableArraySlice<T> elements) {

return ToVlogString(port::ArraySlice<T>(elements));

}

string ToVlogString(dnn::DepthToSpaceLayout depth_to_space_layout) {

switch (depth_to_space_layout) {

case dnn::DepthToSpaceLayout::DepthHeightWidth:

return "DepthToSpaceLayout::DepthHeightWidth";

}

}

// Used together with PARAM to VLOG calls made to the stream. Intended

// to be used like this:

//

// VLOG(1) << CallStr("MyFunction", this, {PARAM(a), PARAM(b)});

//

// where a and b are the parameters to MyFunction.

//

// See VLOG_CALL for a short-hand for this. This way of doing it saves

// a tremendous amount of boilerplate code given how many functions

// there are on Stream and how many parameters they each have.

string CallStr(const char *function_name, Stream *stream,

std::vector<std::pair<const char *, string>> params) {

// Do not call this function unless VLOG is on since just

// constructing all the strings in params is expensive.

CHECK(VLOG_IS_ON(1));

string str = port::StrCat("Called Stream::", function_name, "(");

const char *separator = "";

for (const auto &param : params) {

port::StrAppend(&str, separator, param.first, "=", param.second);

separator = ", ";

}

port::StrAppend(&str, ") stream=", ToVlogString(stream));

if (VLOG_IS_ON(10)) {

port::StrAppend(&str, " ", port::CurrentStackTrace(), "\n");

}

return str;

}

// Use this macro to avoid having to type every parameter twice to log

// it with VLOG and CallStr.

#define PARAM(parameter) \

{ #parameter, ToVlogString(parameter) }

// Use this macro to avoid having to type out the name of each

// function and to save some boilerplate. Intended to be used like this:

//

// VLOG_CALL(PARAM(a), PARAM(b))

//

// This saves a tremendous amount of boilerplate compared to the alternative:

//

// VLOG(1) << "Calling MyFunction(a=" << ToVlogString(a)

// << ", b=" << ToVlogString(b);

//

// Note here that most of the parameter names are not short and that

// most of the functions take many more than 2 parameters.

#define VLOG_CALL(...) VLOG(1) << CallStr(__func__, this, {__VA_ARGS__})

} // namespace

Stream::Stream(StreamExecutor *parent)

: parent_(parent),

implementation_(parent->implementation()->GetStreamImplementation()),

allocated_(false),

ok_(false),

temporary_memory_manager_(this) {

VLOG_CALL(PARAM(parent));

}

Stream::Stream(StreamExecutor *parent,

internal::StreamInterface *implementation)

: parent_(parent),

implementation_(implementation),

allocated_(false),

ok_(false),

temporary_memory_manager_(this) {

VLOG_CALL(PARAM(parent), PARAM(implementation));

}

Stream::~Stream() {

VLOG_CALL();

temporary_memory_manager_.ForceDeallocateAll();

if (allocated_) {

parent_->DeallocateStream(this);

}

}

Stream &Stream::Init() {

VLOG_CALL();

mutex_lock lock{mu_};

CHECK_EQ(false, allocated_)

<< "stream appears to already have been initialized";

CHECK(!ok_) << "stream should be in !ok() state pre-initialization";

if (parent_->AllocateStream(this)) {

// Successful initialization!

allocated_ = true;

ok_ = true;

} else {

LOG(ERROR) << "failed to allocate stream during initialization";

}

return *this;

}

Stream &Stream::InitTimer(Timer *timer) {

VLOG_CALL(PARAM(timer));

if (ok()) {

CheckError(parent_->AllocateTimer(timer));

} else {

LOG(INFO) << "did not allocate timer: " << timer;

}

return *this;

}

Stream &Stream::InitWithTimer(Timer *timer) {

VLOG_CALL(PARAM(timer));

return Init().InitTimer(timer);

}

Stream &Stream::ThenRecordEvent(Event *event) {

VLOG_CALL(PARAM(event));

port::Status status = parent_->RecordEvent(this, event);

if (!status.ok()) {

LOG(ERROR) << "Error recording event in stream: " << status.error_message()

<< "; not marking stream as bad, as the Event object may be "

<< "at fault. Monitor for further errors.";

}

return *this;

}

Stream &Stream::ThenBatchNormalizationForward(

const DeviceMemory<float> &x, const DeviceMemory<float> &scale,

const DeviceMemory<float> &offset,

const DeviceMemory<float> &estimated_mean,

const DeviceMemory<float> &estimated_variance,

const dnn::BatchDescriptor &x_desc,

const dnn::BatchDescriptor &scale_offset_desc, const double epsilon,

DeviceMemory<float> *y, DeviceMemory<float> *batch_mean,

DeviceMemory<float> *batch_var, DeviceMemory<float> *saved_mean,

DeviceMemory<float> *saved_inv_var, bool is_training,

std::function<const DeviceMemory<float> &()> var_to_inv_var,

std::function<void()> inv_var_to_var) {

VLOG_CALL(PARAM(x), PARAM(scale), PARAM(offset), PARAM(x_desc),

PARAM(scale_offset_desc), PARAM(epsilon), PARAM(y));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoBatchNormalizationForward(

this, x, scale, offset, estimated_mean, estimated_variance, x_desc,

scale_offset_desc, epsilon, y, batch_mean, batch_var, saved_mean,

saved_inv_var, is_training, std::move(var_to_inv_var),

std::move(inv_var_to_var)));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenBatchNormalizationBackward(

const DeviceMemory<float> &y_backprop, const DeviceMemory<float> &x,

const DeviceMemory<float> &scale, const DeviceMemory<float> &mean,

const DeviceMemory<float> &variance, 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) {

VLOG_CALL(PARAM(y_backprop), PARAM(x), PARAM(scale), PARAM(x_desc),

PARAM(scale_offset_desc), PARAM(epsilon), PARAM(x_backprop),

PARAM(scale_backprop), PARAM(offset_backprop));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoBatchNormalizationBackward(

this, y_backprop, x, scale, mean, variance, x_desc, scale_offset_desc,

epsilon, x_backprop, scale_backprop, offset_backprop));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenConvolveWithScratch(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<Eigen::half> &input_data,

const dnn::FilterDescriptor &filter_descriptor,

const DeviceMemory<Eigen::half> &filter_data,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<Eigen::half> *output,

ScratchAllocator *scratch_allocator) {

VLOG_CALL(PARAM(input_descriptor), PARAM(input_data),

PARAM(filter_descriptor), PARAM(filter_data),

PARAM(convolution_descriptor), PARAM(output_descriptor),

PARAM(output));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoConvolve(

this, input_descriptor, input_data, filter_descriptor, filter_data,

convolution_descriptor, output_descriptor, output,

/*scratch_allocator=*/scratch_allocator, dnn::AlgorithmConfig(),

nullptr));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenConvolveWithScratch(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<float> &input_data,

const dnn::FilterDescriptor &filter_descriptor,

const DeviceMemory<float> &filter_data,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::BatchDescriptor &output_descriptor, DeviceMemory<float> *output,

ScratchAllocator *scratch_allocator) {

VLOG_CALL(PARAM(input_descriptor), PARAM(input_data),

PARAM(filter_descriptor), PARAM(filter_data),

PARAM(convolution_descriptor), PARAM(output_descriptor),

PARAM(output));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoConvolve(

this, input_descriptor, input_data, filter_descriptor, filter_data,

convolution_descriptor, output_descriptor, output,

/*scratch_allocator=*/scratch_allocator, dnn::AlgorithmConfig(),

nullptr));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenConvolveWithAlgorithm(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<float> &input_data,

const dnn::FilterDescriptor &filter_descriptor,

const DeviceMemory<float> &filter_data,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::BatchDescriptor &output_descriptor, DeviceMemory<float> *output,

ScratchAllocator *scratch_allocator,

const dnn::AlgorithmConfig &algorithm_config,

dnn::ProfileResult *output_profile_result) {

VLOG_CALL(PARAM(input_descriptor), PARAM(input_data),

PARAM(filter_descriptor), PARAM(filter_data),

PARAM(convolution_descriptor), PARAM(output_descriptor),

PARAM(output));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

auto status = dnn->DoConvolve(

this, input_descriptor, input_data, filter_descriptor, filter_data,

convolution_descriptor, output_descriptor, output, scratch_allocator,

algorithm_config, output_profile_result);

if (!status && !output_profile_result) {

SetError();

}

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenConvolveWithAlgorithm(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<Eigen::half> &input_data,

const dnn::FilterDescriptor &filter_descriptor,

const DeviceMemory<Eigen::half> &filter_data,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<Eigen::half> *output, ScratchAllocator *scratch_allocator,

const dnn::AlgorithmConfig &algorithm_config,

dnn::ProfileResult *output_profile_result) {

VLOG_CALL(PARAM(input_descriptor), PARAM(input_data),

PARAM(filter_descriptor), PARAM(filter_data),

PARAM(convolution_descriptor), PARAM(output_descriptor),

PARAM(output));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

auto status = dnn->DoConvolve(

this, input_descriptor, input_data, filter_descriptor, filter_data,

convolution_descriptor, output_descriptor, output, scratch_allocator,

algorithm_config, output_profile_result);

if (!status && !output_profile_result) {

SetError();

}

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenConvolve(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<float> &input_data,

const dnn::FilterDescriptor &filter_descriptor,

const DeviceMemory<float> &filter_data,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<float> *output) {

return ThenConvolveWithScratch(input_descriptor, input_data,

filter_descriptor, filter_data,

convolution_descriptor, output_descriptor,

output, /*scratch_allocator=*/nullptr);

}

Stream &Stream::ThenConvolveQuantized(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<float> &input_data,

const dnn::FilterDescriptor &filter_descriptor,

const DeviceMemory<int8> &filter_coefficients,

const DeviceMemory<float> &coefficient_scales,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<float> *output) {

VLOG_CALL(PARAM(input_descriptor), PARAM(input_data),

PARAM(filter_descriptor), PARAM(filter_coefficients),

PARAM(coefficient_scales), PARAM(convolution_descriptor),

PARAM(output_descriptor), PARAM(output));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoConvolveQuantized(

this, input_descriptor, input_data, filter_descriptor,

filter_coefficients, coefficient_scales, convolution_descriptor,

output_descriptor, output));

} else {

SetError();

LOG(WARNING)

<< "attempting to perform DNN operation using StreamExecutor "

"without DNN support";

}

}

return *this;

}

Stream &Stream::ThenConvolveQuantized(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<float> &input_data,

const dnn::FilterDescriptor &filter_descriptor,

const DeviceMemory<int16> &filter_coefficients,

const DeviceMemory<float> &coefficient_scales,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<float> *output) {

VLOG_CALL(PARAM(input_descriptor), PARAM(input_data),

PARAM(filter_descriptor), PARAM(filter_coefficients),

PARAM(coefficient_scales), PARAM(convolution_descriptor),

PARAM(output_descriptor), PARAM(output));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoConvolveQuantized(

this, input_descriptor, input_data, filter_descriptor,

filter_coefficients, coefficient_scales, convolution_descriptor,

output_descriptor, output));

} else {

SetError();

LOG(WARNING)

<< "attempting to perform DNN operation using StreamExecutor "

"without DNN support";

}

}

return *this;

}

Stream &Stream::ThenSeparableConvolve(

const dnn::BatchDescriptor &batch_descriptor,

const DeviceMemory<float> &input_data,

const dnn::FilterDescriptor &filter_descriptor, int depth_multiplier,

const DeviceMemory<float> &first_weights,

const DeviceMemory<float> &second_weights,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<float> *output) {

VLOG_CALL(

PARAM(batch_descriptor), PARAM(input_data), PARAM(filter_descriptor),

PARAM(depth_multiplier), PARAM(first_weights), PARAM(second_weights),

PARAM(convolution_descriptor), PARAM(output_descriptor), PARAM(output));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoSeparableConvolve(

this, batch_descriptor, input_data, filter_descriptor,

depth_multiplier, first_weights, second_weights,

convolution_descriptor, output_descriptor, output));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenConvolveBackwardDataWithScratch(

const dnn::FilterDescriptor &filter_descriptor,

const DeviceMemory<float> &filter_data,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<float> backward_output_data,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::BatchDescriptor &input_descriptor,

DeviceMemory<float> *backward_input_data,

ScratchAllocator *scratch_allocator) {

VLOG_CALL(PARAM(filter_descriptor), PARAM(filter_data),

PARAM(output_descriptor), PARAM(backward_output_data),

PARAM(convolution_descriptor), PARAM(input_descriptor),

PARAM(backward_input_data));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoConvolveBackwardData(

this, filter_descriptor, filter_data, output_descriptor,

backward_output_data, convolution_descriptor, input_descriptor,

backward_input_data, scratch_allocator, dnn::AlgorithmConfig(),

nullptr));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenConvolveBackwardDataWithAlgorithm(

const dnn::FilterDescriptor &filter_descriptor,

const DeviceMemory<float> &filter_data,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<float> backward_output_data,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::BatchDescriptor &input_descriptor,

DeviceMemory<float> *backward_input_data,

ScratchAllocator *scratch_allocator,

const dnn::AlgorithmConfig &algorithm_config,

dnn::ProfileResult *output_profile_result) {

VLOG_CALL(PARAM(filter_descriptor), PARAM(filter_data),

PARAM(output_descriptor), PARAM(backward_output_data),

PARAM(convolution_descriptor), PARAM(input_descriptor),

PARAM(backward_input_data));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

auto status = dnn->DoConvolveBackwardData(

this, filter_descriptor, filter_data, output_descriptor,

backward_output_data, convolution_descriptor, input_descriptor,

backward_input_data, scratch_allocator, algorithm_config,

output_profile_result);

if (!status && !output_profile_result) {

SetError();

}

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenConvolveBackwardDataWithAlgorithm(

const dnn::FilterDescriptor &filter_descriptor,

const DeviceMemory<Eigen::half> &filter_data,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<Eigen::half> backward_output_data,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::BatchDescriptor &input_descriptor,

DeviceMemory<Eigen::half> *backward_input_data,

ScratchAllocator *scratch_allocator,

const dnn::AlgorithmConfig &algorithm_config,

dnn::ProfileResult *output_profile_result) {

VLOG_CALL(PARAM(filter_descriptor), PARAM(filter_data),

PARAM(output_descriptor), PARAM(backward_output_data),

PARAM(convolution_descriptor), PARAM(input_descriptor),

PARAM(backward_input_data));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

auto status = dnn->DoConvolveBackwardData(

this, filter_descriptor, filter_data, output_descriptor,

backward_output_data, convolution_descriptor, input_descriptor,

backward_input_data, scratch_allocator, algorithm_config,

output_profile_result);

if (!status && !output_profile_result) {

SetError();

}

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenConvolveBackwardDataWithScratch(

const dnn::FilterDescriptor &filter_descriptor,

const DeviceMemory<Eigen::half> &filter_data,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<Eigen::half> backward_output_data,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::BatchDescriptor &input_descriptor,

DeviceMemory<Eigen::half> *backward_input_data,

ScratchAllocator *scratch_allocator) {

VLOG_CALL(PARAM(filter_descriptor), PARAM(filter_data),

PARAM(output_descriptor), PARAM(backward_output_data),

PARAM(convolution_descriptor), PARAM(input_descriptor),

PARAM(backward_input_data));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoConvolveBackwardData(

this, filter_descriptor, filter_data, output_descriptor,

backward_output_data, convolution_descriptor, input_descriptor,

backward_input_data, scratch_allocator, dnn::AlgorithmConfig(),

nullptr));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenConvolveBackwardData(

const dnn::FilterDescriptor &filter_descriptor,

const DeviceMemory<float> &filter_data,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<float> backward_output_data,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::BatchDescriptor &input_descriptor,

DeviceMemory<float> *backward_input_data) {

return ThenConvolveBackwardDataWithScratch(

filter_descriptor, filter_data, output_descriptor, backward_output_data,

convolution_descriptor, input_descriptor, backward_input_data,

/*scratch_allocator=*/nullptr);

}

Stream &Stream::ThenConvolveBackwardFilterWithScratch(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<float> &input_data,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<float> backward_output_data,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::FilterDescriptor &filter_descriptor,

DeviceMemory<float> *backward_filter_data,

ScratchAllocator *scratch_allocator) {

VLOG_CALL(PARAM(input_descriptor), PARAM(input_data),

PARAM(output_descriptor), PARAM(backward_output_data),

PARAM(convolution_descriptor), PARAM(filter_descriptor),

PARAM(backward_filter_data));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoConvolveBackwardFilter(

this, input_descriptor, input_data, output_descriptor,

backward_output_data, convolution_descriptor, filter_descriptor,

backward_filter_data, scratch_allocator, dnn::AlgorithmConfig(),

nullptr));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenConvolveBackwardFilterWithAlgorithm(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<float> &input_data,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<float> backward_output_data,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::FilterDescriptor &filter_descriptor,

DeviceMemory<float> *backward_filter_data,

ScratchAllocator *scratch_allocator,

const dnn::AlgorithmConfig &algorithm_config,

dnn::ProfileResult *output_profile_result) {

VLOG_CALL(PARAM(input_descriptor), PARAM(input_data),

PARAM(output_descriptor), PARAM(backward_output_data),

PARAM(convolution_descriptor), PARAM(filter_descriptor),

PARAM(backward_filter_data));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

auto status = dnn->DoConvolveBackwardFilter(

this, input_descriptor, input_data, output_descriptor,

backward_output_data, convolution_descriptor, filter_descriptor,

backward_filter_data, scratch_allocator, algorithm_config,

output_profile_result);

if (!status && !output_profile_result) {

SetError();

}

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenConvolveBackwardFilterWithScratch(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<Eigen::half> &input_data,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<Eigen::half> backward_output_data,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::FilterDescriptor &filter_descriptor,

DeviceMemory<Eigen::half> *backward_filter_data,

ScratchAllocator *scratch_allocator) {

VLOG_CALL(PARAM(input_descriptor), PARAM(input_data),

PARAM(output_descriptor), PARAM(backward_output_data),

PARAM(convolution_descriptor), PARAM(filter_descriptor),

PARAM(backward_filter_data));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoConvolveBackwardFilter(

this, input_descriptor, input_data, output_descriptor,

backward_output_data, convolution_descriptor, filter_descriptor,

backward_filter_data, scratch_allocator, dnn::AlgorithmConfig(),

nullptr));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenConvolveBackwardFilterWithAlgorithm(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<Eigen::half> &input_data,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<Eigen::half> backward_output_data,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::FilterDescriptor &filter_descriptor,

DeviceMemory<Eigen::half> *backward_filter_data,

ScratchAllocator *scratch_allocator,

const dnn::AlgorithmConfig &algorithm_config,

dnn::ProfileResult *output_profile_result) {

VLOG_CALL(PARAM(input_descriptor), PARAM(input_data),

PARAM(output_descriptor), PARAM(backward_output_data),

PARAM(convolution_descriptor), PARAM(filter_descriptor),

PARAM(backward_filter_data));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

auto status = dnn->DoConvolveBackwardFilter(

this, input_descriptor, input_data, output_descriptor,

backward_output_data, convolution_descriptor, filter_descriptor,

backward_filter_data, scratch_allocator, algorithm_config,

output_profile_result);

if (!status && !output_profile_result) {

SetError();

}

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenConvolveBackwardFilter(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<float> &input_data,

const dnn::BatchDescriptor &output_descriptor,

DeviceMemory<float> backward_output_data,

const dnn::ConvolutionDescriptor &convolution_descriptor,

const dnn::FilterDescriptor &filter_descriptor,

DeviceMemory<float> *backward_filter_data) {

return ThenConvolveBackwardFilterWithScratch(

input_descriptor, input_data, output_descriptor, backward_output_data,

convolution_descriptor, filter_descriptor, backward_filter_data,

/*scratch_allocator=*/nullptr);

}

template <typename T>

Stream &Stream::ThenConvolveBackwardBiasImpl(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<T> &input_data,

const dnn::BatchDescriptor &bias_descriptor,

DeviceMemory<T> *backward_bias_data) {

VLOG_CALL(PARAM(input_descriptor), PARAM(input_data), PARAM(bias_descriptor),

PARAM(backward_bias_data));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoConvolveBackwardBias(this, input_descriptor, input_data,

bias_descriptor,

backward_bias_data));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenConvolveBackwardBias(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<double> &input_data,

const dnn::BatchDescriptor &bias_descriptor,

DeviceMemory<double> *backward_bias_data) {

return ThenConvolveBackwardBiasImpl(input_descriptor, input_data,

bias_descriptor, backward_bias_data);

}

Stream &Stream::ThenConvolveBackwardBias(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<float> &input_data,

const dnn::BatchDescriptor &bias_descriptor,

DeviceMemory<float> *backward_bias_data) {

return ThenConvolveBackwardBiasImpl(input_descriptor, input_data,

bias_descriptor, backward_bias_data);

}

Stream &Stream::ThenConvolveBackwardBias(

const dnn::BatchDescriptor &input_descriptor,

const DeviceMemory<Eigen::half> &input_data,

const dnn::BatchDescriptor &bias_descriptor,

DeviceMemory<Eigen::half> *backward_bias_data) {

return ThenConvolveBackwardBiasImpl(input_descriptor, input_data,

bias_descriptor, backward_bias_data);

}

Stream &Stream::ThenMatMul(const DeviceMemory<float> &input_data,

const DeviceMemory<float> &weights,

const dnn::BatchDescriptor &input_dimensions,

const dnn::BatchDescriptor &output_dimensions,

DeviceMemory<float> *output_data) {

VLOG_CALL(PARAM(input_data), PARAM(weights), PARAM(input_dimensions),

PARAM(output_dimensions), PARAM(output_data));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoMatMul(this, input_data, weights, input_dimensions,

output_dimensions, output_data));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenMatMulQuantized(

const DeviceMemory<float> &input_data, const DeviceMemory<int8> &weights,

const DeviceMemory<float> &weight_scales,

const dnn::BatchDescriptor &input_dimensions,

const dnn::BatchDescriptor &output_dimensions,

DeviceMemory<float> *output_data) {

VLOG_CALL(PARAM(input_data), PARAM(weights), PARAM(weight_scales),

PARAM(input_dimensions), PARAM(output_dimensions),

PARAM(output_data));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoMatMulQuantized(this, input_data, weights,

weight_scales, input_dimensions,

output_dimensions, output_data));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenMatMulQuantized(

const DeviceMemory<float> &input_data, const DeviceMemory<int16> &weights,

const DeviceMemory<float> &weight_scales,

const dnn::BatchDescriptor &input_dimensions,

const dnn::BatchDescriptor &output_dimensions,

DeviceMemory<float> *output_data) {

VLOG_CALL(PARAM(input_data), PARAM(weights), PARAM(weight_scales),

PARAM(input_dimensions), PARAM(output_dimensions),

PARAM(output_data));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoMatMulQuantized(this, input_data, weights,

weight_scales, input_dimensions,

output_dimensions, output_data));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenBiasAdd(const DeviceMemory<float> &input_data,

const DeviceMemory<float> &biases,

const dnn::BatchDescriptor &dimensions,

DeviceMemory<float> *output_data) {

VLOG_CALL(PARAM(input_data), PARAM(biases), PARAM(dimensions),

PARAM(output_data));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(

dnn->DoBiasAdd(this, input_data, biases, dimensions, output_data));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenPoolForward(

const dnn::PoolingDescriptor &pooling_dimensions,

const dnn::BatchDescriptor &input_dimensions,

const DeviceMemory<float> &input_data,

const dnn::BatchDescriptor &output_dimensions,

DeviceMemory<float> *output_data) {

VLOG_CALL(PARAM(pooling_dimensions), PARAM(input_dimensions),

PARAM(input_data), PARAM(output_dimensions), PARAM(output_data));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoPoolForward(this, pooling_dimensions, input_dimensions,

input_data, output_dimensions,

output_data));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}

Stream &Stream::ThenPoolForward(

const dnn::PoolingDescriptor &pooling_dimensions,

const dnn::BatchDescriptor &input_dimensions,

const DeviceMemory<Eigen::half> &input_data,

const dnn::BatchDescriptor &output_dimensions,

DeviceMemory<Eigen::half> *output_data) {

VLOG_CALL(PARAM(pooling_dimensions), PARAM(input_dimensions),

PARAM(input_data), PARAM(output_dimensions), PARAM(output_data));

if (ok()) {

if (dnn::DnnSupport *dnn = parent_->AsDnn()) {

CheckError(dnn->DoPoolForward(this, pooling_dimensions, input_dimensions,

input_data, output_dimensions,

output_data));

} else {

SetErrorAndLogNoDnnSupport();

}

}

return *this;

}