// Copyright 2019 Google LLC
//
// 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
//
//      https://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.

#include "integrations/tensorflow/compiler/Passes.h"
#include "iree/base/signature_mangle.h"
#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "llvm/ADT/STLExtras.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/Utils.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"

namespace mlir {
namespace iree_compiler {

using ::iree::SipSignatureMangler;

namespace {

LogicalResult setRawSignatureIndex(FuncOp funcOp, SipSignatureMangler &mangler,
                                   int rawIndex, ArrayAttr indexPathAttr) {
  llvm::SmallVector<SipSignatureMangler::Key, 8> indexKeys;
  for (auto &indexAttr : indexPathAttr) {
    if (auto stringAttr = indexAttr.dyn_cast<StringAttr>()) {
      auto stringRef = stringAttr.getValue();
      indexKeys.emplace_back(
          absl::string_view(stringRef.data(), stringRef.size()));
    } else if (auto intAttr = indexAttr.dyn_cast<IntegerAttr>()) {
      indexKeys.emplace_back(intAttr.getInt());
    } else {
      return funcOp.emitError()
             << "Each index path component must be a string or integer";
    }
  }

  if (!mangler.SetRawSignatureIndex(rawIndex, indexKeys)) {
    return funcOp.emitError()
           << "Unable to generate mangled form for index path";
  }

  return success();
}

LogicalResult ImportTfSavedModelGlobalTensorsToIREEFlow(ModuleOp module) {
  OpBuilder global_builder(module.getBodyRegion());
  SymbolTable symbol_table(module);

  DenseMap<StringRef, std::string> sym_name_to_flow_sym_name;
  for (auto global_tensor : module.getOps<tf_saved_model::GlobalTensorOp>()) {
    auto exported_names = tf_saved_model::GetExportedNames(global_tensor);
    std::string flow_sym_name;
    if (exported_names.empty()) {
      flow_sym_name = "__iree_flow_" + global_tensor.sym_name().str();
    } else if (exported_names.size() == 1) {
      flow_sym_name = exported_names[0].str();
    } else {
      return global_tensor.emitError()
             << "Multiple exported names for global tensor not supported yet";
    }
    sym_name_to_flow_sym_name[global_tensor.sym_name()] = flow_sym_name;
    global_builder.create<IREE::Flow::VariableOp>(
        global_tensor.getLoc(), flow_sym_name, global_tensor.is_mutable(),
        global_tensor.type(), global_tensor.value());
  }

  for (auto func : module.getOps<FuncOp>()) {
    SmallVector<unsigned, 4> args_to_erase;
    for (int i = 0, e = func.getNumArguments(); i < e; i++) {
      tf_saved_model::GlobalTensorOp global_tensor =
          tf_saved_model::LookupBoundInput(func, i, symbol_table);
      if (!global_tensor) {
        continue;
      }
      args_to_erase.push_back(i);
      auto flow_sym_ref = global_builder.getSymbolRefAttr(
          sym_name_to_flow_sym_name[global_tensor.sym_name()]);
      Value arg = func.getArgument(i);
      if (global_tensor.is_mutable()) {
        // The value is a tensor<*x!tf.resource> type, which flows into
        // tf.ReadVariableOp/tf.AssignVariableOp.
        // XLA resource functionalization should have canonicalized everything
        // to uses of those two ops in the body of the tf_saved_model exported
        // function.
        for (OpOperand &operand : llvm::make_early_inc_range(arg.getUses())) {
          if (auto read_variable =
                  dyn_cast<TF::ReadVariableOp>(operand.getOwner())) {
            auto load = OpBuilder(read_variable)
                            .create<IREE::Flow::VariableLoadOp>(
                                read_variable.getLoc(),
                                read_variable.value().getType(), flow_sym_ref);
            read_variable.value().replaceAllUsesWith(load.result());
            read_variable.erase();
            continue;
          }
          if (auto assign_variable =
                  dyn_cast<TF::AssignVariableOp>(operand.getOwner())) {
            OpBuilder(assign_variable)
                .create<IREE::Flow::VariableStoreOp>(assign_variable.getLoc(),
                                                     assign_variable.value(),
                                                     flow_sym_ref);
            assign_variable.erase();
            continue;
          }
          return operand.getOwner()->emitError()
                 << "unknown op operating on resource for global tensor";
        }
      } else {
        // The value is already a tensor value type. Just RAUW it with a
        // `flow.variable.load`.
        auto load =
            OpBuilder(func.getBody())
                .create<IREE::Flow::VariableLoadOp>(
                    global_tensor.getLoc(), arg.getType(), flow_sym_ref);
        arg.replaceAllUsesWith(load.result());
      }
    }
    func.eraseArguments(args_to_erase);
  }

  // Erase all the global tensors.
  for (auto global_tensor : llvm::make_early_inc_range(
           module.getOps<tf_saved_model::GlobalTensorOp>())) {
    global_tensor.erase();
  }
  return success();
}

}  // namespace

class TFSavedModelAdoptExportsPass
    : public ModulePass<TFSavedModelAdoptExportsPass> {
 public:
  void runOnModule() override {
    if (failed(run())) {
      signalPassFailure();
    }
  }

  LogicalResult run() {
    mlir::Builder builder(getModule());
    Identifier savedModelIndexPathIdent =
        builder.getIdentifier("tf_saved_model.index_path");
    Identifier ireeReflectionIdent = builder.getIdentifier("iree.reflection");
    Identifier ireeModuleExportIdent =
        builder.getIdentifier("iree.module.export");
    Identifier sipIdent = builder.getIdentifier("sip");
    Identifier abiIdent = builder.getIdentifier("abi");
    Identifier abiVersionIdent = builder.getIdentifier("abiv");

    if (failed(ImportTfSavedModelGlobalTensorsToIREEFlow(getModule()))) {
      return failure();
    }

    // Handle saved model exported functions.
    for (auto func : getModule().getOps<FuncOp>()) {
      // Transfer exported names to IREE.
      auto exported_names = mlir::tf_saved_model::GetExportedNames(func);
      if (exported_names.empty()) continue;

      // TODO(laurenzo): After VM rework, we should just keep the
      // function name as-is and create explicit export ops for each exported
      // function.
      if (exported_names.size() > 1) {
        return func.emitError() << "Multiple exported names not supported yet";
      }
      func.setName(exported_names.front());

      // Function level reflection attributes.
      SipSignatureMangler inputsMangler;
      SipSignatureMangler resultsMangler;
      SmallVector<NamedAttribute, 3> funcReflectAttrs;
      funcReflectAttrs.push_back(
          builder.getNamedAttr(abiIdent, builder.getStringAttr(sipIdent)));
      funcReflectAttrs.push_back(
          builder.getNamedAttr(abiVersionIdent, builder.getI32IntegerAttr(1)));

      // Tag it as an IREE exported function.
      func.setAttr(ireeModuleExportIdent, builder.getUnitAttr());

      // Process per-argument attrs and generate reflection metadata.
      for (int i = 0, e = func.getNumArguments(); i < e; i++) {
        auto indexPathAttr =
            func.getArgAttrOfType<mlir::ArrayAttr>(i, savedModelIndexPathIdent);
        if (!indexPathAttr) {
          return func.emitError()
                 << "Missing argument attribute: " << savedModelIndexPathIdent;
        }
        func.removeArgAttr(i, savedModelIndexPathIdent);

        if (failed(
                setRawSignatureIndex(func, inputsMangler, i, indexPathAttr))) {
          return failure();
        }
      }

      // Process per-result attrs and generate reflection metadata.
      for (int i = 0, e = func.getNumResults(); i < e; i++) {
        auto indexPathAttr = func.getResultAttrOfType<mlir::ArrayAttr>(
            i, savedModelIndexPathIdent);
        if (!indexPathAttr) {
          return func.emitError()
                 << "Missing result attribute: " << savedModelIndexPathIdent;
        }
        func.removeResultAttr(i, savedModelIndexPathIdent);

        if (failed(
                setRawSignatureIndex(func, resultsMangler, i, indexPathAttr))) {
          return failure();
        }
      }

      // Add the function level reflection attribute.
      auto functionSignature = SipSignatureMangler::ToFunctionSignature(
          inputsMangler, resultsMangler);
      if (!functionSignature) {
        return func.emitError() << "Unable to generate sip function signature";
      }
      funcReflectAttrs.push_back(builder.getNamedAttr(
          sipIdent, builder.getStringAttr(functionSignature->encoded())));

      if (!funcReflectAttrs.empty()) {
        func.setAttr(ireeReflectionIdent,
                     builder.getDictionaryAttr(funcReflectAttrs));
      }

      // Remove its designation as a saved model export.
      func.removeAttr("tf_saved_model.exported_names");
    }

    // We should have now removed anything requiring saved model semantics.
    getModule().removeAttr("tf_saved_model.semantics");
    return success();
  }
};

std::unique_ptr<OpPassBase<ModuleOp>> createTFSavedModelAdoptExportsPass() {
  return std::make_unique<TFSavedModelAdoptExportsPass>();
}

static PassRegistration<TFSavedModelAdoptExportsPass> pass(
    "iree-tf-saved-model-adopt-exports", "Adopts TF saved model exports");

}  // namespace iree_compiler
}  // namespace mlir