/* Copyright 2022 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.
==============================================================================*/

syntax = "proto3";

package orbax;

message ConverterOptionsV2 {
  // TPU conversion options.
  repeated TpuFunction tpu_functions = 1;

  // Specifications for function to be executed with XLA compilation and their
  // device assignments. Note that the converter will return an error when the
  // specified functions have pre-existing device assignment.
  //
  // TODO(b/278562978): Deprecate tpu_functions above to use this field for TPU
  // conversion interface once the support for the function spec is complete.
  repeated XlaFunctionSpec xla_functions = 4;

  // Set to true to enable saftety checks that ensure every op that is placed
  // on the TPU without enable_soft_device_place=True is supported by TPUs.
  // If an op is not supported by TPUs it will error during the execution of the
  // model. This safety check lets the user know early on that there is an
  // issue.
  //
  // This flag is NOT enabled by default for two reasons. 1) In order to perform
  // this check users must ensure all their ops are linked into the Converter
  // binary. The Converter uses these dependencies to verify what ops are
  // supported by TPUs. If a user doesn't link them, the Converter cannot know
  // they exist so the safety check will return a false error. 2) The TF2XLA
  // bridge may perform automatic head and tail outside compilation and move
  // some incompatible ops on the CPU automatically at runtime. It isn't
  // possible to know what ops it will move until the model is actually
  // executed so setting this flag to True could cause false errors.
  bool enable_device_placement_safety_checks = 2;

  // If enabled, the TPU replicate attribute cluster name will be shared
  // across all TPU clusters. By default, each cluster name is unique.
  //
  // Enabling this flag is not recommended for standard TensorFlow models.
  // It was added for Pathways models.
  bool enable_shared_replicate_cluster_name = 3;

  // The state of an optimization.
  enum State {
    // When state is set to default, the optimization will perform its
    // default behavior. For some optimizations this is disabled and for others
    // it is enabled. To check a specific optimization, read the optimization's
    // description.
    DEFAULT = 0;
    // Enabled.
    ENABLED = 1;
    // Disabled.
    DISABLED = 2;
  }

  // Batch options to apply to the graph.
  //
  // Usage:
  // 0. Leave it empty to skip batching.
  // 1. Add one entry with the batch_component unset. The batch option will be
  //    applied to all tpu subgraphs.
  // 2. Add one or more entries, within each specify the batch_component (it can
  //    be a signature or a function). The batch options will be applied to
  //    those components correspondingly.
  repeated BatchOptionsV2 batch_options = 100;

  // Global flag to disable all optimizations that are enabled by default.
  // When enabled, all optimizations that run by default are disabled. If a
  // default optimization is explicitly enabled, this flag will have no affect
  // on that optimization.
  //
  // This flag defaults to false.
  bool disable_default_optimizations = 202;

  // If enabled, apply an optimization that reshapes the tensors going into
  // and out of the TPU. This reshape operation improves performance by reducing
  // the transfer time to and from the TPU.
  //
  // This optimization is incompatible with `input_shape_opt` which is disabled.
  // by default. If `input_shape_opt` is enabled, this option should be
  // disabled.
  //
  // This optimization defaults to enabled for TPU conversions. It is not
  // supported for GPU.
  State io_shape_optimization = 200;

  // If enabled, apply an optimization that updates float variables and float
  // ops on the TPU to bfloat16. This optimization improves performance and
  // throughtput by reducing HBM usage and taking advantage of TPU support for
  // bfloat16.
  //
  // This optimization may cause a loss of accuracy for some models. If an
  // unacceptable loss of accuracy is detected, disable this optimization.
  //
  // This optimization defaults to enabled.
  //
  // Note: This optimization defaults to disabled for models with XlaCallModule,
  // i.e., models converted with the JAX native serialization. If explicitly
  // enabled, it will apply the StableHLO bfloat16 optimization, which may
  // cause output regression. Though it is no longer actively supported, the
  // the APEX team (apex-all@) may offer assistance on a best-effort basis.
  State bfloat16_optimization = 201;

  BFloat16OptimizationOptions bfloat16_optimization_options = 203;

  // If enabled, apply an optimization for XLA I/O. It currently only supports
  // XLA input packing where XLA inputs below the threshold get packed into a
  // single tensor for each type.
  //
  // Inputs of the same data type are reshaped into a single dimension and
  // concatenated outside the XLA function. Then, inside the XLA function, the
  // inputs are split and restored to their original shapes.
  //
  // It currently supports packing tensors with statically known shapes and
  // supports packing inputs to the XLA GPU functions only. It will result in an
  // error when enabled without an XLA GPU function.
  //

  State xla_io_optimization = 207;

  XlaIoOptimizationOptions xla_io_optimization_options = 208;

  // The settings for XLA sharding. If set, XLA sharding is enabled.
  XlaShardingOptions xla_sharding_options = 204;

  // Options to perform device placement for non-XLA functions and ops.
  //
  // Currently, the non-xla device assignment is used on XLA GPU model only.
  NonXlaDevicePlacementOptions non_xla_device_placement_options = 206;

  ExperimentalOptions experimental_options = 209;

  // Features not supported natively by the Inference Converter. These features
  // are supported by individual owners outside the Inference Converter team and
  // may not work for all models. They may be experimental or broken at times.
  // Please consult the individual owners to assess whether they are appropriate
  // for your use case.
}

message XlaFunctionSpec {
  // Type of the device on which the function below is placed. A valid device
  // type must be specified. DEVICE_TYPE_UNSPECIFIED will result in an error.
  enum DeviceType {
    DEVICE_TYPE_UNSPECIFIED = 0;
    GPU = 1;
    TPU = 2;
  }
  DeviceType device_type = 1;

  // Specifies whether the XLA GPU function is intended to be executed on a
  // single or multiple devices.
  //
  // As of June 2025, use single device for regular TFRT, and multiple device
  // for TFRT IFRT unified runtime.
  //
  // DEVICE_EXECUTION_TYPE_UNSPECIFIED will not error and behaves the same as
  // DEFAULT_DEVICE_EXECUTION_TYPE.
  enum DeviceExecutionType {
    DEVICE_EXECUTION_TYPE_UNSPECIFIED = 0;
    DEFAULT_DEVICE_EXECUTION_TYPE = 1;  // Default type is SINGLE_DEVICE.
    SINGLE_DEVICE = 2;
    MULTIPLE_DEVICE = 3;
  }
  DeviceExecutionType device_execution_type = 9;

  // Name of the device on which the specified function should be placed. This
  // field is used only for the GPU conversion and ignored for TPU conversion.
  // For GPU conversion, it defaults to "/device:GPU:0" when empty.
  string device_name = 2;

  oneof function_spec {
    // When jit_compile_functions is set to True, all jit compiled functions
    // are placed on the device.
    //
    // To use this option, decorate the relevant function(s) with
    // @tf.function(jit_compile=True), before exporting. Then set this flag to
    // True. The converter will find all functions that were tagged with
    // jit_compile=True and place them on the device.
    //
    // When using this option, all other settings for the XlaFunctionSpec such
    // as device_name will apply to all functions tagged with jit_compile=True.
    //
    // This option will place all jit_compile=True functions on the device.
    // If only some jit_compile=True functions should be placed on the device,
    // use function_alias or concrete_function_name.
    bool jit_compile_functions = 3;

    // The name of the function alias associated with the function that
    // should be placed on the device. Function aliases are created during model
    // export via the tf.saved_model.SaveOptions.
    //
    // This is a recommended way to specify which function should be placed
    // on the device.
    string function_alias = 4;

    // The name of the concrete function that should be placed on the device.
    // This is the name of the function as it found in the GraphDef and the
    // FunctionDefLibrary.
    //
    // This is NOT the recommended way to specify which function should be
    // placed on the device because concrete function names change every time a
    // model is exported.
    string concrete_function_name = 5;

    // The name of the signature to be placed on the device. The user must make
    // sure there is no incompatible op for the device under the entire
    // signature. Or soft device placement can be enabled. But this will incur
    // performance loss.
    string signature_name = 6;
  }

  // Set to true to enable outside compilation for this function. This is
  // supported only for the function being placed on TPU. If the TPU function
  // has TPU-incompatible ops, outside compilation can automatically move the
  // ops to execute on the CPU at the runtime. For more details, see
  //
  bool enable_soft_device_placement = 7;
}

message TpuFunction {
  // The function(s) that should be placed on the TPU. Only provide a given
  // function once. Duplicates will result in errors. For example, if
  // you provide a specific function via function_alias do not also provide the
  // same function via concrete_function_name or jit_compile_functions.
  oneof name {
    // The name of the function alias associated with the function that
    // should be placed on the TPU. Function aliases are created during model
    // export via the tf.saved_model.SaveOptions.
    //
    // This is a recommended way to specify which function should be placed
    // on the TPU.
    string function_alias = 1;

    // The name of the concrete function that should be placed on the TPU. This
    // is the name of the function as it found in the GraphDef and the
    // FunctionDefLibrary.
    //
    // This is NOT the recommended way to specify which function should be
    // placed on the TPU because concrete function names change every time a
    // model is exported.
    string concrete_function_name = 3;

    // The name of the signature to be placed on the TPU. The user must make
    // sure there is no TPU-incompatible op under the entire signature. Or soft
    // device placement can be enabled. But this will incur performance loss.
    string signature_name = 5;

    // When jit_compile_functions is set to True, all jit compiled functions
    // are placed on the TPU.
    //
    // To use this option, decorate the relevant function(s) with
    // @tf.function(jit_compile=True), before exporting. Then set this flag to
    // True. The converter will find all functions that were tagged with
    // jit_compile=True and place them on the TPU.
    //
    // When using this option, all other settings for the TpuFunction such as
    // enable_soft_device_placement will apply to all functions tagged with
    // jit_compile=True.
    //
    // This option will place all jit_compile=True functions on the TPU.
    // If only some jit_compile=True functions should be placed on the TPU,
    // use function_alias or concrete_function_name.
    bool jit_compile_functions = 4;
  }

  // Set to true to enable outside compilation for this TPU function. If the TPU
  // function has TPU-incompatible ops, outside compilation can automatically
  // move the ops to execute on the CPU at the runtime. For more details, see
  //
  bool enable_soft_device_placement = 2;
}

message BatchOptionsV2 {
  // The component to be batched.
  // 1. Unset if it's for all tpu subgraphs.
  // 2. Set function_alias or concrete_function_name if it's for a function.
  // 3. Set signature_name if it's for a signature.
  oneof batch_component {
    // The function alias associated with the function. Function alias is
    // created during model export via the tf.saved_model.SaveOptions, and is
    // the recommended way to specify functions.
    //
    // Note: Avoid using function_alias when using the converter to update the
    // batch parameters of an existing batch function that was inserted by the
    // inference converter. The inference converter breaks the function alias
    // map and your function alias might no longer be valid (See b/278894058).
    string function_alias = 11;

    // The concreate name of the function. This is the name of the function as
    // it found in the GraphDef and the FunctionDefLibrary. This is NOT the
    // recommended way to specify functions, because concrete function names
    // change every time a model is exported.
    string concrete_function_name = 12;

    // The name of the signature.
    string signature_name = 13;
  }

  // Number of scheduling threads for processing batches of work. Determines
  // the number of batches processed in parallel. This should be roughly in line
  // with the number of TPU cores available.
  int32 num_batch_threads = 1;

  // The maximum allowed batch size. Can be larger than allowed_batch_sizes to
  // utilize large batch splitting.
  int32 max_batch_size = 2;

  // Maximum number of microseconds to wait before outputting an incomplete
  // batch.
  int32 batch_timeout_micros = 3;

  // Optional list of allowed batch sizes. If left empty,
  // does nothing. Otherwise, supplies a list of batch sizes, causing the op
  // to pad batches up to one of those sizes. The entries must increase
  // monotonically, and the final entry must be lower or equal to
  // max_batch_size.
  repeated int32 allowed_batch_sizes = 4;

  // Maximum number of batches enqueued for processing before requests are
  // failed fast.
  int32 max_enqueued_batches = 5;

  // If set, disables large batch splitting which is an efficiency improvement
  // on batching to reduce padding inefficiency.
  // Please see
  bool disable_large_batch_splitting = 6;

  // Enum representing the policy.
  enum MixedPriorityBatchingPolicy {
    // Defaults to LOW_PRIORITY_PADDING_WITH_MAX_BATCH_SIZE which is the
    // BatchFunction op's default.
    UNSPECIFIED_POLICY = 0;
    // Pad the high priority batch with low priority requests up to the max high
    // priority batch size. Low priority only batches can still be processed as
    // in the PRIORITY_ISOLATION policy below.
    LOW_PRIORITY_PADDING_WITH_MAX_BATCH_SIZE = 1;
    // Pad the high priority batch with low priority requests up to the next
    // allowed high priority batch size. Low priority only batches can still be
    // processed as in the PRIORITY_ISOLATION policy below.
    LOW_PRIORITY_PADDING_WITH_NEXT_ALLOWED_BATCH_SIZE = 2;
    // High priority requests and low priority requests are never processed in
    // the same batch. Low priority batches are scheduled iff the batches either
    // max out or time out and there is no high priority request waiting in to
    // be processed. This applies to the policies above as well.
    PRIORITY_ISOLATION = 3;
  }

  // Optional set of batch options for low priority requests, which will be used
  // in priority queue batching.
  message LowPriorityBatchOptions {
    // The maximum allowed batch size.
    int32 max_batch_size = 1;

    // Maximum number of microseconds to wait before outputting an incomplete
    // batch.
    int32 batch_timeout_micros = 2;

    // Optional list of allowed batch sizes. If left empty,
    // does nothing. Otherwise, supplies a list of batch sizes, causing the op
    // to pad batches up to one of those sizes. The entries must increase
    // monotonically, and the final entry must equal max_batch_size.
    repeated int32 allowed_batch_sizes = 3;

    // Maximum number of batches enqueued for processing before requests are
    // failed fast.
    int32 max_enqueued_batches = 4;

    // Policy the determine the mixed priority batching behavior.
    MixedPriorityBatchingPolicy mixed_priority_batching_policy = 5;
  }
  LowPriorityBatchOptions low_priority_batch_options = 8;

  reserved 7;
}

message BFloat16OptimizationOptions {
  // Indicates where the BFloat16 optimization should be applied.
  enum Scope {
    // The scope currently defaults to any computation on XLA devices.
    DEFAULT = 0;
    // Apply the bfloat16 optimization to TPU computation.
    TPU = 1;
    // Apply the bfloat16 optimization to the entire model including CPU
    // computations. The signature types will be updated to bfloat16 unless
    // `preserve_signature` is set.
    ALL = 2;
    // Apply the bfloat16 optimization to XLA GPU computation.
    GPU = 3;
  }

  // This field indicates where the bfloat16 optimization should be applied.
  //
  // The scope defaults to TPU.
  Scope scope = 1;

  // If set, the normal safety checks are skipped. For example, if the model
  // already contains bfloat16 ops, the bfloat16 optimization will error because
  // pre-existing bfloat16 ops can cause issues with the optimization. By
  // setting this flag, the bfloat16 optimization will skip the check.
  //
  // This is an advanced feature and not recommended for almost all models.
  //
  // This flag is off by default.
  bool skip_safety_checks = 2;

  // Ops that should not be converted to bfloat16.
  // Inputs into these ops will be cast to float32, and outputs from these ops
  // will be cast back to bfloat16.
  repeated string filterlist = 3;

  // If set to true, preserves the original types of the signature. If false
  // (default), will convert the signature float (float32) types to bfloat16.
  // Only effective when scope == ALL. Other scopes do not modify the signature
  // by default, and will raise an error if this flag is set.
  bool preserve_signature = 4;

  // Experimental bfloat16 optimization features, which are not guaranteed to
  // work for all models.
  message Experimental {
    // If true, checkpoint float values will be converted to bfloat16.
    // It is only supported if `scope` == ALL.
    bool convert_checkpoint = 1;
  }
  Experimental experimental = 5;

  // If set to true, bfloat16 type cast happens outside function around call
  // node. If false (default) the cast happens inside function after _Arg and
  // before _Retval. Only valid when the batching scope is set to only batch
  // the TPU computation.
  bool enable_cast_outside = 110;
}

message XlaIoOptimizationOptions {
  // Threshold used to select the tensors to pack among the eligible tensors.
  // Any input tensors that are above this threshold size will not be packed.
  // Packing small tensors can provide a performance improvement but packing
  // larger tensors can hurt performance.
  //
  // If not set or set to zero, it defaults to 1,000,000 bytes (1 MB).
  uint64 input_packing_threshold = 1;
}

message PolicyBasedXlaShardingSpec {
  enum XlaShardingPolicy {
    // Returns an error.
    POLICY_UNSPECIFIED = 0;

    // Shards all tensors using `shard_dim` specified below.
    FIXED_DIM = 1;

    // Shards all tensors using dim 0
    FIRST_DIM = 2;

    // Shards all tensors using the last dimension.
    LAST_DIM = 3;

    // Shards tensors alternately using the first and last dimension (i.e. we
    // keep track of a count of all tensors we have seen till now, and shard
    // all the even counts by the first dimension and the odd counts by the
    // last dimension).
    ALTERNATE_FIRST_AND_LAST_DIM = 4;

    // Shards the largest dimension
    LARGEST_DIM = 5;
  }

  // The sharding policy to apply. It should always have a value other than
  // POLICY_UNSPECIFIED. Otherwise, it will result in an error.
  XlaShardingPolicy policy = 6;

  // The dimension on which sharding will be applied. This only applies to the
  // FIXED_DIM policy and it has to be present. If present with other policies,
  // returns an error. Negative values indicate counting from the last
  // dimension.
  optional int32 shard_dim = 1;

  // Minimum tensor rank for policy-based sharding.
  int32 min_tensor_rank = 3;

  // Minimum size of the dimension identified for policy-based sharding.
  int32 min_dimension_size = 4;
}

message XlaShardingOptions {
  // num_cores_per_replica for TPUReplicateMetadata.
  //
  // This is the number of cores you wish to split your model into using XLA
  // SPMD. It always needs to have a positive value.
  int32 num_cores_per_replica = 1;

  // (optional) device_assignment for TPUReplicateMetadata.
  //
  // This is in a flattened [x, y, z, core] format (i.e. core 1 of the chip
  // located in 2,3,0 will be stored as [2,3,0,1]).
  //
  // If this is not specified, then the device assignments will utilize the same
  // topology as specified in the `topology` attribute.
  repeated int32 device_assignment = 2;

  // A serialized string of `tensorflow.tpu.TopologyProto` object, used for
  // `topology` attribute in TPUReplicateMetadata.
  //
  // You must specify the `mesh_shape` and `device_coordinates` attributes in
  // the topology object.
  //
  // This option is required for num_cores_per_replica > 1 cases due to
  // ambiguity of num_cores_per_replica, e.g., pf_1x2x1 with megacore and df_1x1
  // both have num_cores_per_replica = 2, but topology is (1,2,1,1) for pf and
  // (1,1,1,2) for df.
  // - For pf_1x2x1, mesh shape and device_coordinates looks like:
  //   mesh_shape = [1,2,1,1]
  //   device_coordinates=flatten([0,0,0,0], [0,1,0,0])
  // - For df_1x1, mesh shape and device_coordinates looks like:
  //   mesh_shape = [1,1,1,2]
  //   device_coordinates=flatten([0,0,0,0], [0,0,0,1])
  // - For df_2x2, mesh shape and device_coordinates looks like:
  //   mesh_shape = [2,2,1,2]
  //   device_coordinates=flatten(
  //    [0,0,0,0],[0,0,0,1],[0,1,0,0],[0,1,0,1]
  //    [1,0,0,0],[1,0,0,1],[1,1,0,0],[1,1,0,1])
  bytes topology = 3;

  message SplitSpec {
    // Dimension to split (0-indexed).
    uint32 dimension = 1;

    // Number of devices across which the dimension is split.
    uint32 num_devices = 2;
  }

  // Sharding specs to apply to variables.
  message XlaShardingSpec {
    // Regular expression to match the variable name.
    string var_name_regex = 1;

    // Split spec to apply to the variable.
    SplitSpec split_spec = 2;
  }
  repeated XlaShardingSpec sharding_specs = 5;

  message XlaArgShardingSpec {
    // Index of the argument to apply the sharding to.
    uint32 arg_index = 1;

    // Split spec to apply to the argument.
    SplitSpec split_spec = 2;
  }

  message XlaFunctionShardingSpec {
    oneof function_spec {
      // The name of the function alias associated with the function that
      // should be placed on the device. Function aliases are created during
      // model export via the tf.saved_model.SaveOptions.
      string function_alias = 1;

      // The name of the concrete function that should be placed on the device.
      // This is the name of the function as it is found in the GraphDef and the
      // FunctionDefLibrary.
      string concrete_function_name = 2;
    }

    // Sharding specs to apply to the arguments of the function.
    repeated XlaArgShardingSpec arg_sharding_specs = 3;
  }
  // Argument sharding specs for XLA functions, if set for a function that's
  // already converted to an XLA function, it must have been converted with the
  // inference converter V2.
  repeated XlaFunctionShardingSpec function_sharding_specs = 6;

  // Experimental sharding features, which is not guaranteed to work for all
  // models and have no performance guarantee.
  message Experimental {
    // If present, sharding annotations will be automatically added for eligible
    // variables to be sharded across num_cores_per_replica cores based on the
    // provided policy. All pre-existing shardings will be removed.
    PolicyBasedXlaShardingSpec policy_based_sharding_spec = 1;
  }
  Experimental experimental = 4;

  // If true, use Shardy (go/shardy) partitioner. If false, use GSPMD.
  bool use_shardy_partitioner = 7;
}

message NonXlaDevicePlacementOptions {
  // Clear the device assignment for the AssignVariableOps in the restore
  // function gets cleared. See b/281550152 for the context. It is enabled by
  // default. It cannot be enabled at the same time
  // place_unplaced_non_xla_ops_on_cpu is enabled.
  ConverterOptionsV2.State clear_restore_function_device_assignment = 1;

  // Place unplaced ops that are not XLA ops on CPU, except for {VarHandleOp,
  // AssignVariableOp, ReadVariableOp}. It could reduce inference latency when
  // the model has a lot of interleaving between GPU-eligible ops and CPU-only
  // ops. e.g. Bert model. It is disabled by default. It cannot be enabled at
  // the same time clear_restore_function_device_assignment is enabled.
  ConverterOptionsV2.State place_unplaced_non_xla_ops_on_cpu = 2;

  // Place variable related ops {VarHandleOp, AssignVariableOp, ReadVariableOp}
  // on CPU. This only takes effect if place_unplaced_non_xla_ops_on_cpu is
  // enabled. It is disabled by default.
  ConverterOptionsV2.State place_variables_on_cpu = 3;
}

message ExperimentalOptions {
  // Enabled by default. If enabled, removes unused functions from the model.
  // During conversion, unoptimized functions are replaced by their optimized
  // counterparts. The unoptimized functions, though no longer used, remain in
  // in the model (unless this option is enabled).
  //
  // If it causes issues, please leave a comment under b/342436826, or file a
  // bug at
  ConverterOptionsV2.State graph_pruning = 1;
}
