// LINT: LEGACY_NAMES
syntax = "proto3";

package stream_executor.dnn;

// Specifies the data type used by an operation.
enum DataType {
  kFloat = 0;
  kDouble = 1;
  kHalf = 2;
  kInt8 = 3;
  kInt32 = 4;
}

// Describes how a convolution input or output layer's data is formatted.
enum DataLayout {
  // Naming convention:
  // Y <-> row or height
  // X <-> column or width
  // Batch <-> batch, or N
  // Depth <-> feature, or channel
  // TODO(timshen): turn them into cuDNN names, e.g. kNCHW.
  kYXDepthBatch = 0;
  kYXBatchDepth = 1;
  kBatchYXDepth = 2;   // cuDNN's NHWC layout
  kBatchDepthYX = 3;   // cuDNN's NCHW layout
  kBatchDepthYX4 = 4;  // cuDNN's NCHW_VECT_C layout
}

// Describes how a convolution filter is laid out in the memory.
enum FilterLayout {
  // Naming convention:
  // Y <-> row or height
  // X <-> column or width
  // Output <-> output feature, or N
  // Input <-> input feature, or N
  // TODO(timshen): turn them into cuDNN names, e.g. kNCHW.
  kOutputInputYX = 0;   // cuDNN's NCHW layout
  kOutputYXInput = 1;   // cuDNN's NHWC layout
  kOutputInputYX4 = 2;  // cuDNN's NCHW_VECT_C layout
  kInputYXOutput = 3;
  kYXInputOutput = 4;
}

// Describes a kind of non-linearity (threshold-like mathematical function).
enum ActivationMode {
  kNone = 0;
  kSigmoid = 1;
  // Rectified linear activation: f(x) = x < 0 ? 0 : x
  kRelu = 2;
  // Rectified linear activation; where upper maximum is 6.0.
  kRelu6 = 3;
  // Rectified linear activation; where upper maximum specified by
  // BatchDescriptor::value_max().
  kReluX = 4;
  kTanh = 5;
  // Like ReluX; but passes all values in the range [-X,X].
  kBandPass = 6;
}

// Describe the math definition for the conv op. The popular behavior is
// actually called cross-correlation in math, despite the operation is often
// referred as convolution. See cuDNN cudnnConvolutionMode_t.
enum ConvolutionMode {
  CROSS_CORRELATION = 0;
  CONVOLUTION = 1;
}

// Generic tensor representation.
message TensorDescriptorProto {
  repeated int64 dimensions = 1;
  DataType data_type = 2;
  oneof layout_oneof {
    DataLayout data_layout = 3;
    FilterLayout filter_layout = 4;
  }
}

// Generic algorithm representation.
message AlgorithmProto {
  enum MathType {
    DEFAULT_MATH = 0;
    // The GPU may operate 4x4 matrix FMA.
    // See cuDNN's documentation for CUDNN_TENSOR_OP_MATH.
    TENSOR_OP_MATH = 1;
  }
  int64 algo_id = 1;
  MathType math_type = 2;
}

// Convolution-specific parameters.
message ConvolutionDescriptorProto {
  repeated int64 paddings = 1;
  repeated int64 strides = 2;
  repeated int64 dilations = 3;
  // The "accumulator" type. For example, use F32 as an accumulator for F16
  // convolutions.
  // See cuDNN's cudnnConvolutionMode_t.
  DataType compute_mode = 4;
  // See cuDNN's group count.
  int32 group_count = 5;
  ConvolutionMode convolution_mode = 6;
}