// RUN: odml-to-stablehlo-opt -composite-lowering -verify-diagnostics %s | FileCheck %s

func.func @hardswish(%arg0: tensor<2xf32>) -> (tensor<*xf32>) {
  %0 = mhlo.composite "aten.hardswish.default" %arg0 {decomposition = @XlaCallModule_aten.hardswish.default.impl_0} : (tensor<2xf32>) -> tensor<2xf32>
  %1 = "tf.Identity"(%0) {device = ""} : (tensor<2xf32>) -> tensor<*xf32>
  %2 = "tf.Identity"(%1) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>

(Neg(64|32|16|8) x) => (I64Sub (I64Const [0]) x)
(Neg(64|32)F ...) => (F(64|32)Neg ...)

(Com(64|32|16|8) x) => (I64Xor x (I64Const [-1]))

(Not ...) => (I64Eqz ...)

// Lowering pointer arithmetic
(OffPtr ...) => (I64AddConst ...)

// Lowering extension
// It is unnecessary to extend loads
(SignExt32to64        x:(I64Load32S _ _)) => x
(SignExt16to(64|32)   x:(I64Load16S _ _)) => x
(SignExt8to(64|32|16) x:(I64Load8S  _ _)) => x

  // canonicalized to plain HLO constant if statically shaped. Add the
  // canonicalization pattern to pattern list to enable multi-hop lowering.
  chlo::ConstantLikeOp::getCanonicalizationPatterns(patterns, context);

  return ApplyPatterns(op, patterns, legalize_chlo);
}

// Performs the lowering to XLA dialect.
void LegalizeTF::runOnOperation() {
  auto op = getOperation();
  auto op_name = op->getName().getStringRef().str();

                          tensorflow::ConvertTensor(tensor, &builder));
      builder.setInsertionPoint(
          &func_op.getFunctionBody().getBlocks().front().front());
      // Use mhlo.Constant when it is consumed by the lowering passes since they
      // can't lower tf.Const.
      Value cst;
      if (use_mhlo_const) {
        cst = builder.create<mhlo::ConstantOp>(func_op->getLoc(), attrs);
      } else {

# Target Aware Conversion (TAC)

Different hardwares have different capabilities and restrictions.

TAC is designed to leverage hardwares' capabilities to:

*   Perform device-specific optimizations (such as unsupported ops lowering,
    layout transformations, etc.)
*   Graph partitioning based on the hardware costs modeling.
*   It supports general import/export where you can hook your own

	Types          Types
	lowerBlock     blockRewriter  // block lowering function, first round
	lowerValue     valueRewriter  // value lowering function, first round
	lateLowerBlock blockRewriter  // block lowering function that needs to be run after the first round; only used on some architectures
	lateLowerValue valueRewriter  // value lowering function that needs to be run after the first round; only used on some architectures

  // CHECK: mlrt.await_handle [[handle_1]]
  mlrt.await_handle %handle_0
  mlrt.await_handle %handle_1

  // CHECK: return [[result]]
  return %result : tensor<i32>
}

// -----

// Test lowering tf.If

func.func @then(%x: tensor<i32>, %y: tensor<i32>) -> tensor<i32> {
  return %x: tensor<i32>
}

func.func @else(%x: tensor<i32>, %y: tensor<i32>) -> tensor<i32> {
  return %y: tensor<i32>
}

//   passes for backend-specific ops.
// lower_to_xla_hlo: Temporary parameter to be removed in imminent update. If
//   true, includes legalization and MHLO lowering passes.
// allow_partial_conversion: when this is true, allow operations that can't be
//   legalized.
ABSL_DEPRECATED("Use v2/legalize_tf.h::LegalizeMlirToHlo instead.")
void CreateConvertMlirToXlaHloPipeline(

      /*parallel_iterations=*/10, false,
      /*shape_invariant=*/false);

  // `_lower_using_switch_merge` is the default for While ops created
  // in TensorFlow and allows lowering to V1 control flow for loop
  // parallelization.
  dataset_while->setAttr("_lower_using_switch_merge",
                         builder.getBoolAttr(true));

  return dataset_while;
}

// `variable_snapshots` is a map of {index of the input to the
// compilation_result -> underlying Tensor the variable is/was pointing to (i.e.
// the value of the variable at the time of lowering/compilation)}.
//
// The obtained PjRtBuffers are populated to `args` vector.
// `non_donatable_input_indices` will also be set, which contains the indices of
// the input that should not be donated to output.
//