kYXDepthBatch = 0, // Same as dist_belief::DF_DEPTH_MAJOR.

kYXBatchDepth, // Same as dist_belief::DF_BATCH_MAJOR.

kBatchYXDepth, // Same as run_brain output, and tensorflow's layout.

kBatchDepthYX, // cuDNN's NCHW layout, data laid out as image, feature,

// maps, rows, columns.

X = 0,

Y = 1,

Z = 2,

k8Bit = 1,

k16Bit = 2,

k32Bit = 4,

kFloat = 0,

kDouble = 1,

kHalf = 2,

kRnnRelu = 0,

kRnnTanh = 1,

kRnnLstm = 2,

kRnnGru = 3,

public:

struct ParamsRegion {

int64 offset;

int64 size;

};

typedef std::vector<ParamsRegion> ParamsRegions;

virtual ~RnnDescriptor() {}

virtual int64 ParamsSizeInBytes() const { return -1; }

virtual ParamsRegions ParamsWeightRegions() const { return ParamsRegions(); }

virtual ParamsRegions ParamsBiasRegions() const { return ParamsRegions(); }

public:

// Creates a "blank" batch descriptor, which should be initialized via the

// named argument helpers.

BatchDescriptor();

explicit BatchDescriptor(int ndims);

// Clones values from 'other' for initialization.

void CloneFrom(const BatchDescriptor& other);

string ToString() const;

string ToShortString() const;

// Accessors.

int64 count() const { return count_; }

int64 feature_map_count() const { return feature_map_count_; }

int64 height() const { return GetDim(spatial_size_, DimIndex::Y); }

int64 width() const { return GetDim(spatial_size_, DimIndex::X); }

int64 spatial_dim(DimIndex dim) const { return GetDim(spatial_size_, dim); }

int ndims() const { return ndims_; }

float value_max() const { return value_max_; }

float value_min() const { return value_min_; }

DataLayout layout() const { return layout_; }

QuantizedActivationMode quantized_activation_mode() const {

return quantized_activation_mode_;

}

// Full dimensions of the underlying data, ordered according to a specific

// layout.

std::vector<int64> full_dims(const DataLayout& layout) const;

// Full strides of the underlying data, ordered according to a specific

// layout.

std::vector<int64> full_strides(const DataLayout& layout) const;

// Named-argument helpers for avoiding user error during construction.

BatchDescriptor& set_count(int64 value) {

count_ = value;

return *this;

}

BatchDescriptor& set_feature_map_count(int64 value) {

feature_map_count_ = value;

return *this;

}

BatchDescriptor& set_height(int64 value) {

SetDim(&spatial_size_, DimIndex::Y, value);

return *this;

}

BatchDescriptor& set_width(int64 value) {

SetDim(&spatial_size_, DimIndex::X, value);

return *this;

}

BatchDescriptor& set_spatial_dim(DimIndex dim, int64 value) {

SetDim(&spatial_size_, dim, value);

return *this;

}

BatchDescriptor& set_value_max(float value) {

value_max_ = value;

return *this;

}

BatchDescriptor& set_value_min(float value) {

value_min_ = value;

return *this;

}

BatchDescriptor& set_layout(DataLayout layout) {

layout_ = layout;

return *this;

}

BatchDescriptor& set_quantized_activation_mode(

QuantizedActivationMode quantized_activation_mode) {

quantized_activation_mode_ = quantized_activation_mode;

return *this;

}

// Return the number of nodes in a single feature map.

int64 NodesPerFeatureMap() const;

// Return the number of nodes across all feature maps. Note that this is not

// affected by the batch count.

int64 NodesAcrossFeatureMaps() const;

// Returns the number of elements (e.g. RGB pixel values) required to hold a

// given batch descriptor, given a no-padding assumption. Note that this is

// affected by the batch count.

int64 ElementCount() const;

// Return the number of weights required to fully connect a layer with

// dimensions given by the 'input' descriptor with a layer with dimensions

// given by the 'output' descriptor.

static int64 FullyConnectedWeightCount(const BatchDescriptor& input,

const BatchDescriptor& output);

// Return the number of biases required to fully connect to an output layer

// with dimensions given the 'output' descriptor.

static int64 FullyConnectedBiasCount(const BatchDescriptor& output);

// Return a BatchDescriptor for the output of a depth concatenation

// with the given input descriptors. The inputs should have the same

// dimensions, except possibly for feature_map_count(), though this

// function does not verify that.

static BatchDescriptor DepthConcatenateOutputDescriptor(

port::ArraySlice<dnn::BatchDescriptor> inputs);

private:

int64 count_;

int64 feature_map_count_;

// Stored as: ..., y, x.

std::vector<int64> spatial_size_;

float value_max_;

float value_min_;

DataLayout layout_;

int ndims_;

QuantizedActivationMode quantized_activation_mode_;

kOutputInputYX = 0, // cuDNN's default filter layout, laid out as:

// (major) output feature maps >> input feature maps >>

// rows >> columns (minor).

kInputYXOutput, // Same as dist_belief's default filter layout.

kYXInputOutput, // Same as tensorflow's default filter layout.

public:

// By default construction, all dimensions are set to zero, so they should all

// be populated by the user via the named-argument helpers below. (See class

// comment for details.)

FilterDescriptor();

explicit FilterDescriptor(int ndims);

~FilterDescriptor();

// Named-argument helpers for avoiding user error during construction.

FilterDescriptor& set_output_feature_map_count(int64 value) {

output_feature_map_count_ = value;

return *this;

}

FilterDescriptor& set_input_feature_map_count(int64 value) {

input_feature_map_count_ = value;

return *this;

}

FilterDescriptor& set_input_filter_height(int64 value) {

SetDim(&input_filter_dims_, DimIndex::Y, value);

return *this;

}

FilterDescriptor& set_input_filter_width(int64 value) {

SetDim(&input_filter_dims_, DimIndex::X, value);

return *this;

}

FilterDescriptor& set_layout(FilterLayout layout) {

layout_ = layout;

return *this;

}

FilterDescriptor& set_spatial_dim(DimIndex dim, int64 value) {

SetDim(&input_filter_dims_, dim, value);

return *this;

}

int ndims() const { return ndims_; }

void CloneFrom(const FilterDescriptor& other);

string ToString() const;

string ToShortString() const;

// Returns the number of weights required as parameters for a convolution

// using this filter descriptor.

int64 ComputeWeightCount() const;

// Returns the number of biases required as parameters for a convolution

// using this filter descriptor.

int64 bias_count() const { return output_feature_map_count_; }

int64 output_feature_map_count() const { return output_feature_map_count_; }

int64 input_feature_map_count() const { return input_feature_map_count_; }

int64 input_filter_height() const {

return GetDim(input_filter_dims_, DimIndex::Y);

}

int64 input_filter_width() const {

return GetDim(input_filter_dims_, DimIndex::X);

}

int64 input_filter_dim(DimIndex dim) const {

return GetDim(input_filter_dims_, dim);

}

FilterLayout layout() const { return layout_; }

std::vector<int64> input_filter_dims() const { return input_filter_dims_; }

private:

int64 output_feature_map_count_;

int64 input_feature_map_count_;

// Stored as: ..., y, x.

std::vector<int64> input_filter_dims_;

int ndims_;

FilterLayout layout_;

public:

// By default construction, there is no zero-padding and the filter stride is

// 1x1 (centering the filter on every cell in the input layer's

// width-by-height area).

ConvolutionDescriptor();

explicit ConvolutionDescriptor(int ndims);

~ConvolutionDescriptor();

string ToString() const;

string ToShortString() const;

ConvolutionDescriptor& set_zero_padding_height(int64 value) {

SetDim(&zero_padding_, DimIndex::Y, value);

return *this;

}

ConvolutionDescriptor& set_zero_padding_width(int64 value) {

SetDim(&zero_padding_, DimIndex::X, value);

return *this;

}

ConvolutionDescriptor& set_zero_padding(DimIndex dim, int64 value) {

SetDim(&zero_padding_, dim, value);

return *this;

}

ConvolutionDescriptor& set_vertical_filter_stride(int64 value) {

SetDim(&filter_strides_, DimIndex::Y, value);

return *this;

}

ConvolutionDescriptor& set_horizontal_filter_stride(int64 value) {

SetDim(&filter_strides_, DimIndex::X, value);

return *this;

}

ConvolutionDescriptor& set_filter_stride(DimIndex dim, int64 value) {

SetDim(&filter_strides_, dim, value);

return *this;

}

int64 zero_padding_height() const {

return GetDim(zero_padding_, DimIndex::Y);

}

int64 zero_padding_width() const {

return GetDim(zero_padding_, DimIndex::X);

}

int64 vertical_filter_stride() const {

return GetDim(filter_strides_, DimIndex::Y);

}

int64 horizontal_filter_stride() const {

return GetDim(filter_strides_, DimIndex::X);

}

int zero_padding(DimIndex dim) const { return GetDim(zero_padding_, dim); }

int filter_stride(DimIndex dim) const { return GetDim(filter_strides_, dim); }

int ndims() const { return ndims_; }

std::vector<int64> strides() const { return filter_strides_; }

std::vector<int64> padding() const { return zero_padding_; }

private:

// Stored as: .. y, x.

std::vector<int64> zero_padding_;

std::vector<int64> filter_strides_;

int ndims_;

// TODO(leary) cudnn provides these fields, but need to characterize what

// their effect is -- they may be boolean rather than integral.

// int64 upscale_input_x;

// int64 upscale_input_y;

public:

PoolingDescriptor();

explicit PoolingDescriptor(int ndims);

PoolingDescriptor& set_pooling_mode(PoolingMode value) {

mode_ = value;

return *this;

}

PoolingDescriptor& set_window_height(int64 value) {

SetDim(&window_, DimIndex::Y, value);

return *this;

}

PoolingDescriptor& set_window_width(int64 value) {

SetDim(&window_, DimIndex::X, value);

return *this;

}

PoolingDescriptor& set_window(DimIndex dim, int64 value) {

SetDim(&window_, dim, value);

return *this;

}

PoolingDescriptor& set_vertical_padding(int64 value) {

SetDim(&padding_, DimIndex::Y, value);

return *this;

}

PoolingDescriptor& set_horizontal_padding(int64 value) {

SetDim(&padding_, DimIndex::X, value);

return *this;

}

PoolingDescriptor& set_padding(DimIndex dim, int64 value) {

SetDim(&padding_, dim, value);

return *this;

}

PoolingDescriptor& set_vertical_stride(int64 value) {

SetDim(&strides_, DimIndex::Y, value);

return *this;

}

PoolingDescriptor& set_horizontal_stride(int64 value) {

SetDim(&strides_, DimIndex::X, value);

return *this;

}

PoolingDescriptor& set_stride(DimIndex dim, int64 value) {

SetDim(&strides_, dim, value);

return *this;

}

int ndims() const { return ndims_; }

void CloneFrom(const PoolingDescriptor& other);

string ToString() const;

string ToShortString() const;

PoolingMode mode() const { return mode_; }

int64 window_height() const { return GetDim(window_, DimIndex::Y); }

int64 window_width() const { return GetDim(window_, DimIndex::X); }

int64 window(DimIndex dim) const { return GetDim(window_, dim); }

int64 vertical_padding() const { return GetDim(padding_, DimIndex::Y); }

int64 horizontal_padding() const { return GetDim(padding_, DimIndex::X); }

int64 padding(DimIndex dim) const { return GetDim(padding_, dim); }

int64 vertical_stride() const { return GetDim(strides_, DimIndex::Y); }

int64 horizontal_stride() const { return GetDim(strides_, DimIndex::X); }

int64 stride(DimIndex dim) const { return GetDim(strides_, dim); }

std::vector<int64> window() const { return window_; }

std::vector<int64> padding() const { return padding_; }

std::vector<int64> strides() const { return strides_; }

private:

PoolingMode mode_;

int ndims_;

// Stored as: ..., y, x.

std::vector<int64> window_;

std::vector<int64> padding_;

std::vector<int64> strides_;

public:

bool is_valid() const { return is_valid_; }

void set_is_valid(bool val) { is_valid_ = val; }

AlgorithmType algorithm() const { return algorithm_; }

void set_algorithm(AlgorithmType val) { algorithm_ = val; }

float elapsed_time_in_ms() const { return elapsed_time_in_ms_; }

void set_elapsed_time_in_ms(float val) { elapsed_time_in_ms_ = val; }

private:

bool is_valid_ = false;

AlgorithmType algorithm_ = kDefaultAlgorithm;

float elapsed_time_in_ms_ = std::numeric_limits<float>::max();

public:

AlgorithmConfig()

: algorithm_(kDefaultAlgorithm),

algorithm_no_scratch_(kDefaultAlgorithm) {}

explicit AlgorithmConfig(AlgorithmType algorithm)

: algorithm_(algorithm), algorithm_no_scratch_(kDefaultAlgorithm) {}

AlgorithmConfig(AlgorithmType algorithm, AlgorithmType algorithm_no_scratch)

: algorithm_(algorithm), algorithm_no_scratch_(algorithm_no_scratch) {}

AlgorithmType algorithm() const { return algorithm_; }

void set_algorithm(AlgorithmType val) { algorithm_ = val; }

AlgorithmType algorithm_no_scratch() const { return algorithm_no_scratch_; }

void set_algorithm_no_scratch(AlgorithmType val) {

algorithm_no_scratch_ = val;

}

private:

AlgorithmType algorithm_;

AlgorithmType algorithm_no_scratch_;

public:

NormalizeDescriptor();

NormalizeDescriptor& set_bias(float bias) {

bias_ = bias;

return *this;

}

NormalizeDescriptor& set_range(int32 range) {

range_ = range;

return *this;

}

NormalizeDescriptor& set_alpha(float alpha) {

alpha_ = alpha;

return *this;

}

NormalizeDescriptor& set_beta(float beta) {

beta_ = beta;

return *this;

}

NormalizeDescriptor& set_wrap_around(bool wrap_around) {

wrap_around_ = wrap_around;

return *this;

}

NormalizeDescriptor& set_segment_size(int32 segment_size) {

segment_size_ = segment_size;

return *this;

}

void CloneFrom(const NormalizeDescriptor& other);

string ToString() const;

string ToShortString() const;

float bias() const { return bias_; }

int32 range() const { return range_; }

float alpha() const { return alpha_; }

float beta() const { return beta_; }

bool wrap_around() const { return wrap_around_; }

int32 segment_size() const { return segment_size_; }

private:

float bias_;

int32 range_;

float alpha_;

float beta_;

bool wrap_around_;

int32 segment_size_;

kSigmoid,

// Rectified linear activation: f(x) = x < 0 ? 0 : x

kRelu,

// Rectified linear activation, where upper maximum is 6.0.

kRelu6,

// Rectified linear activation, where upper maximum specified by

// BatchDescriptor::value_max().

kReluX,

kTanh,

// Like ReluX, but passes all values in the range [-X,X].

kBandPass,

public:

DnnSupport() {}

virtual ~DnnSupport() {}

virtual port::Status Init() = 0;

// Performs a single-precision forward batch normalization operation onto

// the stream.

//

// Arguments:

// stream: borrowed pointer to the stream that the batch normalization

// operation should be enqueued onto.

// x: input data.

// scale: scaling parameters.

// offset: offset parameters.

// estimated_mean: population mean estimated during training.

// Used for inference only; empty for training.

// estimated_variance: population variance estimated during traning,

// used for inference only; empty for training.

// x_desc: dimensions of the input data, which is the same as the dimensions

// of the output.

// scale_offset_desc: dimensions of scale and offset.

// epsilon: a small floating point number added to the variance of x.

// y: output data.

// batch_mean: batch mean, to be used to compute the running mean.

// batch_variance: batch variance, to be used to compute

// the running variance.

// reserve_space_1: saved mean, to be reused in the backward gradient

// computation.

// reserve_space_2: saved variance, to be reused in the backward gradient

// computation.

// is_training: Set to true for training, false for inference.

// var_to_inv_var: a function to convert the variance to inverted variance

// for cuDNN v4 forward inference.

// inv_var_to_var: a function to convert the inverted variance to

// variance for cuDNN v4 forward training, to be used for TensorFlow

// to calculate the running variance.

virtual bool DoBatchNormalizationForward(

Stream* stream, 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>* reserve_space_1,

DeviceMemory<float>* reserve_space_2, bool is_training,

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

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

return false;

}

// Performs a single-precision backward batch normalization gradient

// computation operation onto the stream.

//

// Arguments:

// stream: borrowed pointer to the stream that the batch normalization

// gradient computation operation should be enqueued onto.

// y_backprop: gradient with regard to output y.

// x: input data.

// scale: scaling parameters.

// x_desc: dimensions of the input data, which is the same as the dimensions

// of the output.

// scale_offset_desc: dimensions of scale and offset.

// epsilon: a small floating point number added to the variance of x.

// x_backprop: gradient with respect to input x.

// scale_backprop: gradient with respect to scale.

// offset_backprop: gradient with respect to offset.

virtual bool DoBatchNormalizationBackward(

Stream* stream, 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) {

return false;

}

// Enqueues a single-precision convolution operation onto the stream.

//

// Arguments (all borrowed):

// stream: borrowed pointer to the stream that the 'convolve' operation

// should be enqueued onto.

// input_descriptor: dimensions of the input layer.

// input_data: un-owned device memory region which contains the

// convolution input.

// filter_descriptor: dimensions of the convolution filter.

// weights: coefficients for the convolution filter, these are multiplied

// against values in the input that the filter convolves over.

// convolution_descriptor: stride of the convolution filter.

// output_descriptor: dimensions of the output layer.

// output_data: un-owned device memory region in which to place the

// convolution result.

// scratch_allocator: un-owned, may-be-null object that may allocate scratch

// space in order to speed up the convolution operation.

// algorithm: an integer to specify which algorithm should be used for the

// operation. kDefaultAlgorithm means the system will pick an algorithm

// by default. The coding of the algorithm is be interpretted by the

// underlying implementation.

// output_profile_result: the output profile result for this call. The

// profiling is only enabled when this is not nullptr.

//

// input_descriptor, filter_descriptor, convolution_descriptor and

// output_descriptor together specify exactly how the convolution is aligned

// with the input data:

//

// * (input dimensions - filter size + 1) / filter stride == output dimensions

// corresponds to dist_belief padding = VALID, i.e. the input is not padded.

// * input dimensions / filter stride == output dimensions

// corresponds to dist_belief padding = SAME, i.e. input and output are the

// same size - this requires padding the input.

// * (input dimensions + filter size - 1) / filter stride == output dimensions

// corresponds to dist_belief padding = FULL, i.e. the output is sized so

// that if the inverse of the filter is applied to the output in VALID mode

// the result is the same size as the input - this requires even more

// padding of the input.

virtual bool DoConvolve(

Stream* stream, 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_data, ScratchAllocator* scratch_allocator,

const dnn::AlgorithmConfig& algorithm_config,

ProfileResult* output_profile_result) = 0;

// Return a list of algorithms supported by the forward convolution pass.

virtual bool GetConvolveAlgorithms(

std::vector<AlgorithmType>* out_algorithms);

// Enqueues a double-precision convolution operation onto the stream.

// See DoConvolve above for argument details.

virtual bool DoConvolve(

Stream* stream, const dnn::BatchDescriptor& batch_descriptor,

const DeviceMemory<double>& input_data,

const dnn::FilterDescriptor& filter_descriptor,

const DeviceMemory<double>& filter_data,

const dnn::ConvolutionDescriptor& convolution_descriptor,

const dnn::BatchDescriptor& output_descriptor,

DeviceMemory<double>* output_data) = 0;

// Enqueues a half-precision convolution operation onto the stream.

// See DoConvolve above for argument details.

virtual bool DoConvolve(

Stream* stream, const dnn::BatchDescriptor& batch_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_data,

ScratchAllocator* scratch_allocator,

const dnn::AlgorithmConfig& algorithm_config,

ProfileResult* output_profile_result) = 0;

// Variation of the above with the weight matrix split into two matrices.

// first_weights: Coefficients of the first matrix.

// second_weights: Coefficients of the second matrix.

// depth_multiplier: specifies the columns of the first matrix and rows

// of the second one - first_weights columns = depth_multiplier,

// second_weights rows = depth_multiplier *

// filter_descriptor.input_feature_map_count().

// see go/separable for documentation on separable convolutions.

virtual bool DoSeparableConvolve(

Stream* stream, const BatchDescriptor& input_descriptor,

const DeviceMemory<float>& input_data,

const FilterDescriptor& filter_descriptor, int depth_multiplier,

const DeviceMemory<float>& first_weights,

const DeviceMemory<float>& second_weights,

const ConvolutionDescriptor& convolution_descriptor,

const BatchDescriptor& output_descriptor,

DeviceMemory<float>* output_data) = 0;

// Enqueues a single-precision backward convolution (for data) operation onto

// the stream.

//

// Arguments:

// stream: borrowed pointer to the stream that the 'convolve' operation

// should be enqueued onto.

// filter_descriptor: dimensions of the convolution filter.

// filter_data: coefficients for the convolution filter.

// output_descriptor: dimensions of the output gradients, which is the same

// as the dimensions of the output.

// backward_output_data: un-owned device memory region which contains the

// backprop of the output.

// convolution_descriptor: stride of the convolution filter.

// input_descriptor: dimensions of the input layer.

// backward_input_data: un-owned device memory region in which to place the

// backprop of the input.

// scratch_allocator: un-owned, may-be-null object that may allocate scratch

// space in order to speed up the convolution operation.

virtual bool DoConvolveBackwardData(

Stream* stream, const FilterDescriptor& filter_descriptor,

const DeviceMemory<float>& filter_data,

const BatchDescriptor& output_descriptor,

DeviceMemory<float> backward_output_data,

const ConvolutionDescriptor& convolution_descriptor,

const BatchDescriptor& input_descriptor,

DeviceMemory<float>* backward_input_data,

ScratchAllocator* scratch_allocator,

const dnn::AlgorithmConfig& algorithm_config,

ProfileResult* output_profile_result) = 0;

// Return a list of algorithms supported by the backward convolution pass for

// data.

virtual bool GetConvolveBackwardDataAlgorithms(

std::vector<AlgorithmType>* out_algorithms);

virtual bool DoConvolveBackwardData(

Stream* stream, const FilterDescriptor& filter_descriptor,

const DeviceMemory<Eigen::half>& filter_data,

const BatchDescriptor& output_descriptor,

DeviceMemory<Eigen::half> backward_output_data,

const ConvolutionDescriptor& convolution_descriptor,

const BatchDescriptor& input_descriptor,

DeviceMemory<Eigen::half>* backward_input_data,

ScratchAllocator* scratch_allocator,