limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/mlir/tensorflow/transforms/mlprogram.h"

#include <string>
#include <utility>

#include "llvm/ADT/STLFunctionalExtras.h"
#include "llvm/ADT/Twine.h"
#include "mlir/Transforms/Passes.h"  // from @llvm-project

  // CHECK-DAG: ml_program.global public @vars.Variable_1([]) : tensor<!tf_type.string>
  func.func @lowers_string_ops(%arg0: tensor<128xi32>, %arg1: tensor<128xi32>, %arg2: tensor<128x1xi32>, %arg3: tensor<128x90xi32>, %arg4: tensor<128x90xi32>, %arg5: tensor<128x90xi32>, %arg6: tensor<128x90x64xf32>, %arg7: tensor<128x90x64xf32>) -> tensor<!tf_type.string> {

Matthias Kramm <******@****.***> 1662565451 -0700

#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/transforms/mlprogram.h"

namespace tensorflow {

void RegisterMlProgramPasses() {
  mlir::registerPassPipeline(
      "tf-lower-to-mlprogram-and-hlo", "Lower TF to ml_program + mhlo",
      [](mlir::OpPassManager& pm, llvm::StringRef options,
         llvm::function_ref<mlir::LogicalResult(const llvm::Twine&)>

#include "mlir/Dialect/Func/Extensions/AllExtensions.h"  // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/Dialect/MLProgram/IR/MLProgram.h"  // from @llvm-project
#include "mlir/Dialect/MLProgram/IR/MLProgramAttributes.h"  // from @llvm-project
#include "mlir/IR/Dialect.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"

#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallSet.h"
#include "mlir/Dialect/Affine/Utils.h"  // from @llvm-project
#include "mlir/Dialect/MLProgram/IR/MLProgram.h"  // from @llvm-project
#include "mlir/Dialect/MLProgram/IR/MLProgramAttributes.h"  // from @llvm-project
#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/IR/Dialect.h"  // from @llvm-project

#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"  // from @llvm-project
#include "mlir/Analysis/DataFlowFramework.h"  // from @llvm-project
#include "mlir/Dialect/MLProgram/IR/MLProgram.h"  // from @llvm-project
#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project

}

def LowerGlobalsToMlProgramPass: Pass<"tf-saved-model-lower-globals-to-mlprogram", "mlir::ModuleOp"> {
  let summary = "Remove (and remap) function arguments that map to global tensors.";
  let description = [{
    This pass will lower AssignVariableOp and ReadVariableOp to their ml_program
    equivalents, for those resources which stem from a "bound input" function argument

// RUN: tf-opt %s --allow-unregistered-dialect --tf-saved-model-lower-globals-to-mlprogram --split-input-file --verify-diagnostics

module attributes {tf_saved_model.semantics} {
  "tf_saved_model.global_tensor"() { is_mutable, sym_name = "v", type = tensor<?xf32>, value = dense<1.> : tensor<1xf32> } : () -> ()
  "tf_saved_model.global_tensor"() { is_mutable, sym_name = "v1", type = tensor<?xf32>, value = dense<1.> : tensor<1xf32> } : () -> ()

// RUN: tf-opt --allow-unregistered-dialect --split-input-file --tf-saved-model-lower-variable-ops-to-mlprogram %s | FileCheck %s

// CHECK-LABEL: module
module attributes {tf_saved_model.semantics} {
  // CHECK: ml_program.global {{.*}} @vars.v
  // CHECK: ml_program.global {{.*}} @vars.bar
  "tf_saved_model.global_tensor"() { is_mutable, sym_name = "v", type = tensor<10xf32>, value = dense<[0.,1.,2.,3.,4.,5.,6.,7.,8.,9.]> : tensor<10xf32> } : () -> ()