// license that can be found in the LICENSE file.

package unmarshal

import (
	_ "embed"
	"go/ast"
	"go/types"

	"golang.org/x/tools/go/analysis"
	"golang.org/x/tools/go/analysis/passes/inspect"
	"golang.org/x/tools/go/analysis/passes/internal/analysisutil"
	"golang.org/x/tools/go/ast/inspector"
	"golang.org/x/tools/go/types/typeutil"
	"golang.org/x/tools/internal/typesinternal"
)

//go:embed doc.go
var doc string

namespace mlir::quant::stablehlo {

#define GEN_PASS_DEF_REMOVESHARDINGCUSTOMCALLPASS
#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"

// Include patterns generated from `remove_sharding_custom_call.td`.
#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/remove_sharding_custom_call.inc"

class RemoveShardingCustomCallPass
    : public impl::RemoveShardingCustomCallPassBase<

#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/fill_quantization_options.h"

namespace stablehlo {

# Verifies golang.org/issue/37555.

[short] skip

# 'go test -json' should say a test passes if it says it passes.
go test -json ./pass
stdout '"Action":"pass","Package":"[^"]*","Elapsed":[^"]*}\n\z'
! stdout '"Test":.*\n\z'

# 'go test -json' should say a test passes if it exits 0 and prints nothing.
# TODO(golang.org/issue/29062): this should fail in the future.
go test -json ./exit0main
stdout '"Action":"pass".*\n\z'
! stdout '"Test":.*\n\z'

#include <string>
#include <utility>

#include "llvm/ADT/STLFunctionalExtras.h"
#include "llvm/ADT/Twine.h"
#include "mlir/Transforms/Passes.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/transforms/bridge.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_saved_model_passes.h"

// traditional control flow structure (Switch/Merge/etc.) into functional
// control flow structure (XlaIf/XlaWhile). Following passes must
// handle those FunctionDef correctly.

// POST_REWRITE_FOR_EXEC passes that support auto-clustering to enable XLA:

REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 5,
                      CloneConstantsForBetterClusteringPass);

void CreateReplicatedClusteringPipeline(OpPassManager &pm,
                                        llvm::StringRef module_name) {
  // Since the internal bridge clustering passes are shared among TF1/TF2
  // TF2-only passes should go here. However, this should be very rare and
  // new passes generally should go into the internal
  // AddReplicatedBridgeClusteringPipelinePasses.
  pm.addPass(mlir::TFTPU::CreateTPUValidateInputsPass());

#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/passes.h"
#include "tensorflow/compiler/mlir/lite/transforms/passes.h"  // IWYU pragma: keep
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"  // IWYU pragma: keep
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace odml {
namespace {

// This file is generated from `passes.td` and provides the implementation base
// class.

#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/config.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/post_calibration.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.h"  // IWYU pragma: keep
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"  // IWYU pragma: keep

#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir::quant::stablehlo {

#define GEN_PASS_DEF_XLACALLMODULETOCALLPASS
#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"

namespace {

// Converts XlaCallModuleOps to func.call.