// RUN: odml_to_stablehlo %s -skip-resize -smuggle-disallowed-ops -o - | FileCheck %s
// RUN: odml-to-stablehlo-opt %s --smuggle-disallowed-ops-pass | FileCheck %s --check-prefix=CHECK-OPT

// CHECK-LABEL: @main
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 975 : i32}, tf_saved_model.semantics}  {

    %0 = stablehlo.composite "odml.scaled_dot_product_attention" %arg0, %arg1, %arg2, %arg3, %arg4 {decomposition = @odml.scaled_dot_product_attention.impl} : (tensor<1x100x32x4xf32>, tensor<1x500x4x4xf32>,...

  let dependentDialects = ["TF::TensorFlowDialect"];
  let constructor = "mlir::odml::CreateLegalizeHloToTfPass()";
}

def LegalizeHloToTfLitePass : Pass<"tfl-legalize-hlo", "mlir::ModuleOp"> {
  let summary = "Legalize from MHLO to the TFLite dialect";
  let dependentDialects = ["TFL::TensorFlowLiteDialect"];
  let constructor = "mlir::odml::CreateLegalizeHloToTfLitePass()";
}

// RUN: odml-to-stablehlo-opt %s -stablehlo-custom-call-legalize-composite | FileCheck %s

// CHECK-LABEL: module
module {
  // CHECK-LABEL: @main
  func.func @main(%arg0: tensor<1xf32>, %arg1: tensor<2xf32>) {
    // CHECK: stablehlo.custom_call @foo
    stablehlo.custom_call @foo() : () -> ()
    // CHECK-NOT: stablehlo.custom_call

namespace mlir {
namespace odml {

#define GEN_PASS_DEF_LEGALIZECOMPOSITETOCUSTOMOPPASS
#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/passes.h.inc"

namespace {
bool IsSupportedComposite(::mlir::stablehlo::CompositeOp op) {
  // List of supported composites to represent using CustomOp.
  return llvm::is_contained(
      {"odml.update_kv_cache", "odml.scaled_dot_product_attention"},
      op.getName());

#define TENSORFLOW_COMPILER_MLIR_LITE_STABLEHLO_TRANSFORMS_COMPOSITE_LOWERING_PASS_H_

namespace mlir {
namespace odml {

std::unique_ptr<Pass> CreateCompositeLoweringPass();

}  // namespace odml
}  // namespace mlir

  // directly. Additionally remove the composite to custom once ODML long term
  // solution lands.
  pass_manager.addPass(
      mlir::odml::createLegalizeStablehloCustomCallToCompositePass());
  pass_manager.addNestedPass<mlir::func::FuncOp>(
      mlir::odml::createLegalizeCompositeToCustomOpPass());
}

 \___/|___/|_|  |_|____|  \___\___/_||_\_/\___|_|  \__\___|_|
)";

int main(int argc, char* argv[]) {
  tensorflow::InitMlir y(&argc, &argv);
  llvm::errs() << art << "\n";

  mlir::odml::registerODMLConverterPasses();
  mlir::odml::registerLegalizeStablehloToVhloPass();

  mlir::DialectRegistry registry;
  registry.insert<mlir::func::FuncDialect, mlir::stablehlo::StablehloDialect,

    outputDir = file('foo-out')
    testOutputDir = file('foo-out-test')
}
"""

        //then
        def iml = getFile([:], 'root.iml').text
        assert iml.contains('inherit-compiler-output="false"')
        assert iml.contains('foo-out')
        assert iml.contains('foo-out-test')
    }

    @Test
    @ToBeFixedForConfigurationCache
    void dslSupportsShortFormsForModule() {

namespace mlir {
namespace odml {

// Adds passes which transform TF Ops to StableHLO Ops.
void AddLegalizeTFToStablehloPasses(OpPassManager& pm,
                                    bool skip_quantization_ops,
                                    bool skip_resize,
                                    bool skip_partitioned_calls);

}  // namespace odml
}  // namespace mlir