/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.

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

    http://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 "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.h"

#include <algorithm>
#include <cstdint>
#include <optional>
#include <queue>
#include <stack>
#include <string>

#include "absl/container/flat_hash_set.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/ErrorHandling.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
#include "mlir/IR/Diagnostics.h"  // from @llvm-project
#include "mlir/IR/Location.h"  // from @llvm-project
#include "mlir/IR/MLIRContext.h"  // from @llvm-project
#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/IR/SymbolTable.h"  // from @llvm-project
#include "mlir/IR/TypeRange.h"  // from @llvm-project
#include "mlir/IR/Value.h"  // from @llvm-project
#include "mlir/IR/ValueRange.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/lite/quantization/quantization_utils.h"
#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/stablehlo_type_utils.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/quantization_unit_loc.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/xla_call_module_attrs.h"
#include "tensorflow/core/ir/types/dialect.h"
#include "tensorflow/core/platform/mutex.h"

namespace mlir::quant {

// Default version number for native serialization.
constexpr int64_t kDefaultVersion = 9;
// Default platform for XlaCallModuleOp.
constexpr StringRef kPlatformCpu = "CPU";
// Name of `tf.XlaCallModule`'s dictionary attribute for keeping the
// deserialized stablehlo module's attributes.
constexpr llvm::StringRef kStablehloModuleAttrsAttrName =
    "_stablehlo_module_attrs";
// Attribute required for running shape refinement pass enabled in XlaCallModule
// version 8 and above.
constexpr llvm::StringRef kUsesShapePolymorphismAttr =
    "jax.uses_shape_polymorphism";

// Checks if the op is inside a lifted function.
bool IsInLiftedFunc(Operation& op) {
  return op.getParentOfType<func::FuncOp>()->hasAttr(kFusedFunctionAttr);
}

// Inserts the function to the symbol table of the module thread-safely.
StringAttr InsertToSymbolTable(Operation& module, Operation& function,
                               const std::string& func_name) {
  static tensorflow::mutex* mtx = new tensorflow::mutex();
  tensorflow::mutex_lock lock(*mtx);

  SymbolTable symbol_table(&module);
  std::string unique_name = func_name;
  int32_t uniquing_counter = 0;
  while (symbol_table.lookup(unique_name) != nullptr) {
    ++uniquing_counter;
    unique_name = func_name + "_" + std::to_string(uniquing_counter);
  }
  function.setAttr("sym_name",
                   StringAttr::get(module.getContext(), unique_name));
  return symbol_table.insert(&function);
}

// Creates the TF::PartitionedCallOp with the given arguments and output types.
// This function call op is for invoking the TF subgraphs.
ValueRange createTFPartitionedCallOp(OpBuilder builder, Location location,
                                     StringRef func_name,
                                     TypeRange output_types, ValueRange args) {
  TF::PartitionedCallOp call_op = builder.create<TF::PartitionedCallOp>(
      location, output_types, args,
      FlatSymbolRefAttr::get(builder.getStringAttr(func_name)),
      /*config=*/"", /*config_proto=*/"", /*executor_type=*/"");

  // Set the attribute to annotate this function call op as a quantizable spot.
  call_op->setAttr(
      kQuantTraitAttrName,
      builder.getStringAttr(llvm::StringRef(
          std::string(QuantTraitValues[QuantizationTrait::FullyQuantizable]))));

  return call_op.getOutput();
}

// Creates the TF::XlaCallModuleOp with the given arguments and output types.
// This function call op is for invoking the StableHLO subgraphs.
ValueRange createTFXlaCallModuleOp(OpBuilder builder, Location location,
                                   StringRef func_name, TypeRange output_types,
                                   ValueRange args) {
  auto ctx = builder.getContext();
  // Collect the shapes of the output to fill up the Sout attribute.
  SmallVector<Attribute> shape_attrs;
  for (const Type result_type : output_types) {
    shape_attrs.push_back(
        tf_type::ShapeAttr::get(ctx, result_type.cast<ShapedType>()));
  }
  auto empty_array_attr = ArrayAttr::get(ctx, {});
  auto platforms = ArrayAttr::get(ctx, {StringAttr::get(ctx, kPlatformCpu)});

  TF::XlaCallModuleOp call_op = builder.create<TF::XlaCallModuleOp>(
      location,
      /*output=*/output_types,
      /*args=*/args,
      /*version=*/kDefaultVersion, /*module=*/"",
      /*Sout=*/ArrayAttr::get(ctx, shape_attrs),
      /*dim_args_spec=*/empty_array_attr,
      /*platforms=*/platforms,
      /*function_list=*/empty_array_attr,
      /*has_token_input_output=*/false,
      /*disabled_checks=*/empty_array_attr);

  // Set the function name. This will be controlled by the
  // XlaCallModuleSerialization related passes directly, which means that the
  // function name can be changed by those passes.
  call_op->setAttr(TF::kStablehloEntryFunctionAttrName,
                   FlatSymbolRefAttr::get(builder.getStringAttr(func_name)));
  // Store the custom attribute to restore the function name when loading it
  // back in the post calibration stage. As mentioned above, the above entry
  // function attribute is not reliable.
  call_op->setAttr(kOriginalStablehloEntryFunctionAttrName,
                   builder.getStringAttr(func_name));

  // Set the attribute to annotate this function call op as a quantizable spot.
  call_op->setAttr(
      kQuantTraitAttrName,
      builder.getStringAttr(llvm::StringRef(
          std::string(QuantTraitValues[QuantizationTrait::FullyQuantizable]))));

  // Set jax.uses_shape_polymorphism=true to enable shape refinement at runtime.
  // This is needed for native serialization version >= 8.
  call_op->setAttr(kStablehloModuleAttrsAttrName,
                   builder.getDictionaryAttr(builder.getNamedAttr(
                       kUsesShapePolymorphismAttr, builder.getBoolAttr(true))));

  return call_op.getOutput();
}

// Creates the function call op based on the given call_op_type argument.
ValueRange createFunctionCallOp(OpBuilder builder, Location location,
                                FunctionCallOpType call_op_type,
                                StringRef func_name, TypeRange output_types,
                                ValueRange args) {
  switch (call_op_type) {
    case FunctionCallOpType::TFXlaCallModuleOp:
      return createTFXlaCallModuleOp(builder, location, func_name, output_types,
                                     args);
    case FunctionCallOpType::TFPartitionedCallOp:
      return createTFPartitionedCallOp(builder, location, func_name,
                                       output_types, args);
    default:
      llvm_unreachable("unhandled call op type");
  }
}

// Finds ops in the paths from arguments to results. The ops is listed in an
// order that the former ops shouldn't have any dependencies on the later ones.
llvm::SmallVector<Operation*> FindOpsFromArgumentsToResults(
    const llvm::SmallVector<Value>& arguments,
    const llvm::SmallVector<Value>& results) {
  std::queue<Value> value_queue;
  for (Value result : results) {
    value_queue.push(result);
  }
  absl::flat_hash_set<mlir::detail::ValueImpl*> argument_set;
  for (Value argument : arguments) {
    argument_set.insert(argument.getImpl());
  }

  // Searching for ops from results to arguments. Duplicate ops in the op stack
  // are intentional in order to make sure the op on the top of the stack
  // doesn't depends on any ops below it.
  std::stack<Operation*> op_stack;
  while (!value_queue.empty()) {
    Value current_value = value_queue.front();
    value_queue.pop();

    Operation* defining_node = current_value.getDefiningOp();
    if (defining_node == nullptr) continue;
    op_stack.push(defining_node);
    for (const auto& arg : defining_node->getOperands()) {
      if (!argument_set.contains(arg.getImpl())) {
        value_queue.push(arg);
      }
    }
  }

  // Remove duplicate ops from the op stack.
  llvm::SmallVector<Operation*> sorted_ops;
  absl::flat_hash_set<Operation*> unique_ops;
  while (!op_stack.empty()) {
    Operation* current_op = op_stack.top();
    op_stack.pop();
    if (unique_ops.contains(current_op)) continue;
    sorted_ops.push_back(current_op);
    unique_ops.insert(current_op);
  }
  return sorted_ops;
}

// Finds the name of each attribute in `attributes` and set the attr_map
// attribute which maps an attribute identifier to its attribute name. The
// identifier is the order of that attribute in `attributes`. This map
// is then used to set attributes in the quantized functions in the
// QuantizeCompositeFunctionsPass.
// For example, for tf.MatMul with `attributes` = {{"transpose_a", false},
// {"transpose_b", false}}, the generated attr_map is
// "0:transpose_a,1:transpose_b", where 0 and 1 are the respective attribute
// identifiers.
// This function returns success if all attributes could be found.
LogicalResult SetAttributeMap(
    MLIRContext& context, const llvm::SmallVector<NamedAttribute>& attributes,
    const llvm::SmallVector<Operation*>& ops) {
  // A map to find which operation an attribute belongs to.
  // The key for this map uses the entire NamedAttribute object, i.e. the
  // {attribute_name, attribute_value} pair.
  llvm::SmallDenseMap<NamedAttribute, Operation*> attr_to_op_map;
  for (Operation* op : ops) {
    for (const NamedAttribute named_attr : op->getAttrs()) {
      attr_to_op_map.insert({named_attr, op});
    }
  }

  for (int idx : llvm::seq<int>(0, attributes.size())) {
    const NamedAttribute& attribute = attributes[idx];
    // Skip the following steps if the attribute value is `NullAttribute`.
    if (const auto string_attr =
            attribute.getValue().dyn_cast_or_null<StringAttr>();
        string_attr != nullptr &&
        string_attr.getValue().equals(kNullAttributeValue)) {
      continue;
    }

    if (std::find_if(
            attr_to_op_map.begin(), attr_to_op_map.end(), [&](auto attr_op) {
              return std::get<0>(attr_op).getName() == attribute.getName();
            }) == attr_to_op_map.end()) {
      mlir::emitError(UnknownLoc::get(&context),
                      "Could not find attribute: " + attribute.getName().str());
      return failure();
    }

    Operation* owner_op;
    for (const auto& [attr, val] : attr_to_op_map) {
      if (attr.getName() == attribute.getName()) owner_op = val;
    }
    if (stablehlo::IsStablehloOp(owner_op)) {
      owner_op->setAttr(StringRef(attribute.getName()), attribute.getValue());
    } else {
      owner_op = attr_to_op_map[attribute];

      std::string new_attr_map_str{};
      if (owner_op->hasAttr(kAttrMapAttribute)) {
        new_attr_map_str =
            owner_op->getAttrOfType<StringAttr>(kAttrMapAttribute).str();
        absl::StrAppend(&new_attr_map_str, ",");
      }

      // Append "<identifier>:<attribute_name>". Ex) "0:transpose_a".
      const std::string identifier = std::to_string(idx);
      const mlir::StringAttr attribute_name = attribute.getName();
      absl::StrAppend(&new_attr_map_str, identifier, ":", attribute_name.str());
      owner_op->setAttr(kAttrMapAttribute,
                        StringAttr::get(&context, new_attr_map_str));
    }
  }
  return success();
}

// Creates a function to wrap the section between arguments and results.
llvm::SmallVector<Value, 4> LiftAsFunctionCall(
    OpBuilder builder, Location location, FunctionCallOpType call_op_type,
    StringRef func_name, const llvm::SmallVector<Value>& arguments,
    const llvm::SmallVector<Value>& results,
    const llvm::SmallVector<NamedAttribute>& attributes) {
  MLIRContext* context = builder.getContext();
  if (results.empty()) {
    mlir::emitError(UnknownLoc::get(context), "No result values specified");
    return {};
  }
  Operation* result_op = results[0].getDefiningOp();
  auto module = result_op->getParentOfType<ModuleOp>();

  // Create a private function and copy all ops between arguments and results.
  auto current_func = result_op->getParentOfType<func::FuncOp>();
  auto guard = OpBuilder::InsertionGuard(builder);
  builder.setInsertionPointAfter(current_func);
  TypeRange arg_types{ValueRange{arguments}};
  TypeRange result_types{ValueRange{results}};
  auto func_type = FunctionType::get(context, arg_types, result_types);

  llvm::SmallVector<Location> arg_locs;
  for (const auto& arg : arguments) {
    arg_locs.push_back(arg.getLoc());
  }
  auto wrap_func = builder.create<func::FuncOp>(location, func_name, func_type);
  wrap_func.setVisibility(SymbolTable::Visibility::Private);
  // The callee function for TF::XlaCallModuleOp must have this attribute.
  if (call_op_type == FunctionCallOpType::TFXlaCallModuleOp) {
    wrap_func->setAttr(TF::kFromXlaCallModuleAttrName, builder.getUnitAttr());
  }
  wrap_func->setAttr(kFusedFunctionAttr, builder.getUnitAttr());
  builder.createBlock(&wrap_func.getBody(), wrap_func.begin(), arg_types,
                      arg_locs);

  IRMapping mapping;
  for (int32_t i : llvm::seq<int32_t>(0, arguments.size())) {
    mapping.map(arguments[i], wrap_func.getArgument(i));
  }

  auto cloning_ops = FindOpsFromArgumentsToResults(arguments, results);
  // Set the location of call op to QuantizationUnitLoc if found.
  Location call_op_loc = location;
  for (Operation* op : cloning_ops) {
    std::optional<QuantizationUnitLoc::QuantizationUnit> unit =
        FindQuantizationUnitFromLoc(op->getLoc());
    if (unit.has_value()) {
      call_op_loc = QuantizationUnitLoc(builder.getContext(), unit.value());
    }
  }

  if (failed(SetAttributeMap(*context, attributes, cloning_ops))) {
    current_func.emitError() << "Some attributes couldn't be found.";
  }
  for (Operation* op : cloning_ops) {
    builder.clone(*op, mapping);
  }

  llvm::SmallVector<Value> return_values;
  for (Value result : results) {
    return_values.push_back(mapping.lookupOrNull(result));
  }
  builder.create<mlir::func::ReturnOp>(location, return_values);

  // Create a function call to the newly created function.
  StringAttr new_func_name =
      InsertToSymbolTable(*module, *wrap_func, func_name.str());
  builder.setInsertionPointAfter(result_op);
  ValueRange new_results =
      createFunctionCallOp(builder, call_op_loc, call_op_type,
                           new_func_name.getValue(), result_types, arguments);
  return llvm::SmallVector<Value, 4>(new_results.begin(), new_results.end());
}

llvm::SmallVector<Value, 4> LiftAsFunctionCall(
    OpBuilder builder, Location location, FunctionCallOpType call_op_type,
    StringRef func_name, const llvm::SmallVector<Value>& arguments,
    const llvm::SmallVector<Value>& results) {
  llvm::SmallVector<NamedAttribute> attributes;
  return LiftAsFunctionCall(builder, location, call_op_type, func_name,
                            arguments, results, attributes);
}

llvm::SmallVector<Value> AppendToVector(
    const llvm::SmallVector<Value>& arguments, Value append) {
  llvm::SmallVector<Value> ret(arguments);
  ret.push_back(append);
  return ret;
}

// Check if the given einsum equation is supported by XlaDotV2.
// Conditions:
// 1. Two inputs & one output.
// 2. No ... in the equation.
// 3. Batch dimensions should be the same, or only the left equation should have
//    the batch dimension. This condition is from the XlaDotV2 specification. It
//    could process the following equation by setting the attributes properly:
//    abc,cd->abd.
// 4. The output should be in the form: [batch dims][lhs dims][rhs dims]
bool IsEinsumSupportedByXlaDotV2(mlir::StringAttr equation_attr) {
  StringRef equation = equation_attr.getValue();

  if (!absl::StrContains(equation, "->") || !absl::StrContains(equation, ",") ||
      absl::StrContains(equation, ".")) {
    return false;
  }

  // Parse equation.
  int idx_arrow = equation.find("->");
  StringRef calc_eq = equation.substr(0, idx_arrow);
  StringRef out_eq = equation.substr(idx_arrow + 2);

  int idx_comma = calc_eq.find(',');
  StringRef lhs_eq = calc_eq.substr(0, idx_comma);
  StringRef rhs_eq = calc_eq.substr(idx_comma + 1);

  if (absl::StrContains(rhs_eq, ",")) return false;

  int lhs_out_idx_start = out_eq.size();
  int lhs_out_idx_end = -1;
  int rhs_out_idx_start = out_eq.size();
  int rhs_out_idx_end = -1;
  int lhs_batch_dim_size = 0;
  int rhs_batch_dim_size = 0;
  for (const char c : lhs_eq) {
    if (absl::StrContains(out_eq, c) && absl::StrContains(rhs_eq, c)) {
      lhs_batch_dim_size++;
    } else if (absl::StrContains(out_eq, c)) {
      const int out_idx = out_eq.find(c);
      if (out_idx < lhs_out_idx_end) {
        // Left-hand equation is reversed in the output.
        return false;
      }
      lhs_out_idx_start = std::min(lhs_out_idx_start, out_idx);
      lhs_out_idx_end = std::max(lhs_out_idx_end, out_idx);
    }
  }

  for (const char c : rhs_eq) {
    if (absl::StrContains(out_eq, c) && absl::StrContains(lhs_eq, c)) {
      rhs_batch_dim_size++;
    } else if (absl::StrContains(out_eq, c)) {
      int out_idx = out_eq.find(c);
      if (out_idx < rhs_out_idx_end) {
        return false;
      }
      if (out_idx < rhs_out_idx_start) rhs_out_idx_start = out_idx;
      if (out_idx > rhs_out_idx_end) rhs_out_idx_end = out_idx;
    }
  }

  if (lhs_batch_dim_size != rhs_batch_dim_size && lhs_batch_dim_size != 0 &&
      rhs_batch_dim_size != 0) {
    // Batch dimension does not match.
    return false;
  }

  // All the lhs equations should come first.
  if (lhs_out_idx_end > rhs_out_idx_start) return false;

  // All the lhs out dim and rhs out dim should be larger than the batch dims,
  // and they should not be mixed.
  int batch_dim_size = std::max(rhs_batch_dim_size, lhs_batch_dim_size);
  return lhs_out_idx_start >= batch_dim_size &&
         rhs_out_idx_start >= batch_dim_size;
}

}  // namespace mlir::quant
