// LINT: LEGACY_NAMES
syntax = "proto3";

package stream_executor.dnn;

option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/stream_executor";

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

// 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.
  //
  // Note: In cudnn, kBatchDepthYX4 and kBatchDepthYX32 are the same layout
  // (namely, NCHW_VECT_C).  It differentiates between these two by using a
  // different data type (int8x4 vs int8x32).  In StreamExecutor we use
  // different layouts for these, because we don't usually pass an explicit data
  // type to StreamExecutor functions.
  kYXDepthBatch = 0;
  kYXBatchDepth = 1;
  kBatchYXDepth = 2;    // cuDNN's NHWC layout
  kBatchDepthYX = 3;    // cuDNN's NCHW layout
  kBatchDepthYX4 = 4;   // cuDNN's NCHW_VECT_C with 4-elem vectors (e.g. int8x4)
  kBatchDepthYX32 = 5;  // cuDNN's NCHW_VECT_C with 32-elem vects (e.g. int8x32)
}

// 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 with 4-elem vectors
  kOutputInputYX32 = 5;  // cuDNN's NCHW_VECT_C layout with 32-elem vectors
  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;
}

enum ConvolutionKind {
  INVALID = 0;
  FORWARD = 1;
  BACKWARD_FILTER = 2;
  BACKWARD_DATA = 3;
  FORWARD_BIAS_ACTIVATION = 4;
}

// 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;
  // cuDNN v8 uses a string to uniquely represent the backend plan.
  string exec_plan_id = 3;
}

// 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;
  // Tensorflow node name, same as in NodeDef, for debugging purposes.
  string name = 7;
}
