/* Copyright 2021 The JAX Authors

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 "jaxlib/mlir.h"

#include <string>
#include <string_view>

#include "absl/log/log.h"
#include "absl/log/vlog_is_on.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir-c/IR.h"
#include "mlir/Bindings/Python/IRCore.h"
#include "mlir/Bytecode/BytecodeWriter.h"
#include "mlir/CAPI/IR.h"
#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/OwningOpRef.h"
#include "mlir/Parser/Parser.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Support/LogicalResult.h"
#include "nanobind/nanobind.h"
#include "nanobind/stl/optional.h"  // IWYU pragma: keep
#include "nanobind/stl/string.h"  // IWYU pragma: keep
#include "nanobind/stl/string_view.h"  // IWYU pragma: keep
#include "shardy/dialect/mpmd/ir/dialect.h"
#include "shardy/dialect/sdy/ir/dialect.h"
#include "stablehlo/dialect/Serialization.h"
#include "xla/hlo/builder/xla_computation.h"
#include "xla/hlo/translate/stablehlo.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"
#include "xla/pjrt/mlir_to_hlo.h"
#include "xla/pjrt/status_casters.h"
#include "xla/python/refine_polymorphic_shapes.h"
#include "xla/python/version.h"
#include "xla/service/hlo.pb.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"

using ::mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::DefaultingPyMlirContext;
using ::mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::PyModule;

namespace nb = nanobind;

namespace jax {
namespace {

std::string PrintModule(mlir::ModuleOp module) {
  std::string s;
  llvm::raw_string_ostream os(s);
  mlir::OpPrintingFlags flags;
  flags.enableDebugInfo();
  module->print(os, flags);
  return s;
}

absl::StatusOr<std::string> SerializeUsingBytecode(mlir::ModuleOp module) {
  std::string bytecode;
  llvm::raw_string_ostream os(bytecode);
  mlir::BytecodeWriterConfig config;
  if (mlir::failed(mlir::writeBytecodeToFile(module, os, config))) {
    return absl::InvalidArgumentError("mlir::writeBytecodeToFile failed");
  }
  return bytecode;
}

void EnablePrintBeforeAndAfter(mlir::PassManager& pm) {
  auto print_before = [](mlir::Pass*, mlir::Operation*) { return true; };
  auto print_after = [](mlir::Pass*, mlir::Operation*) { return true; };
  pm.enableIRPrinting(print_before, print_after);
}

absl::StatusOr<nb::bytes> HloToStableHlo(const nb::bytes& hlo_module_proto) {
  mlir::MLIRContext context;
  if (VLOG_IS_ON(3)) context.disableMultithreading();
  xla::HloModuleProto proto;
  absl::string_view proto_view(hlo_module_proto.c_str(),
                               hlo_module_proto.size());
  if (!proto.ParseFromString(proto_view)) {
    return absl::InvalidArgumentError(
        "Failed to deserialize HloModuleProto");
  }
  ABSL_ASSIGN_OR_RETURN(mlir::OwningOpRef<mlir::ModuleOp> module,
                        ConvertHloToStablehlo(context, &proto));
  ABSL_ASSIGN_OR_RETURN(std::string bytecode, SerializeUsingBytecode(*module));
  return nb::bytes(bytecode.data(), bytecode.size());
}

// Converts an XlaComputation to a StableHLO mlir::Module string.
// Exists for backwards compatibility.
// TODO(phawkins): port remaining users of XlaComputations to use mlir::Modules
// instead and delete this function.
absl::StatusOr<std::string> PyXlaComputationToMlirModule(
    const xla::XlaComputation& computation) {
  mlir::MLIRContext context;
  if (VLOG_IS_ON(3)) context.disableMultithreading();
  ABSL_ASSIGN_OR_RETURN(mlir::OwningOpRef<mlir::ModuleOp> module,
                        ConvertHloToStablehlo(context, &computation.proto()));
  return PrintModule(*module);
}

absl::StatusOr<xla::XlaComputation> PyMlirModuleToXlaComputation(
    std::string_view mlir_module, bool use_tuple_args, bool return_tuple) {
  mlir::MLIRContext context;
  ABSL_ASSIGN_OR_RETURN(mlir::OwningOpRef<mlir::ModuleOp> module,
                        xla::ParseMlirModuleString(mlir_module, context));
  xla::XlaComputation computation;
  ABSL_RETURN_IF_ERROR(MlirToXlaComputation(*module, computation,
                                            use_tuple_args, return_tuple,
                                            /*exec_build_options=*/nullptr));
  return computation;
}

absl::StatusOr<nb::bytes> PyMhloToStablehlo(std::string_view mlir_module) {
  mlir::MLIRContext context;
  if (VLOG_IS_ON(3)) context.disableMultithreading();
  // JAX can be customized in a way that involves operations from custom
  // dialects showing up in JAX IR.
  // `xla::ParseMlirModuleString` won't know about these dialects, but that's
  // fine since we just want to convert MHLO ops to StableHLO ops here and leave
  // everything else unchanged.
  // In order to achieve that, we're allowing unregistered dialects here.
  context.allowUnregisteredDialects(true);
  ABSL_ASSIGN_OR_RETURN(mlir::OwningOpRef<mlir::ModuleOp> module,
                        xla::ParseMlirModuleString(mlir_module, context));
  mlir::PassManager pm(&context);
  if (VLOG_IS_ON(3)) EnablePrintBeforeAndAfter(pm);
  pm.addPass(mlir::mhlo::createHloLegalizeToStablehloPass());
  if (!mlir::succeeded(pm.run(*module))) {
    return tsl::errors::InvalidArgument("MHLO => StableHLO failed");
  }
  // Use bytecode, passing unregistered dialects with properties causes issues
  // when using textual assembly.
  ABSL_ASSIGN_OR_RETURN(std::string bytecode, SerializeUsingBytecode(*module));
  return nb::bytes(bytecode.data(), bytecode.size());
}

absl::StatusOr<nb::bytes> PySerializePortableArtifact(
    std::string_view mlir_module, std::string_view target,
    std::string_view sdy_version, bool use_mixed_serialization) {
  mlir::MLIRContext context;
  context.loadDialect<mlir::mpmd::MpmdDialect>();
  if (VLOG_IS_ON(3)) context.disableMultithreading();
  ABSL_ASSIGN_OR_RETURN(mlir::OwningOpRef<mlir::ModuleOp> module,
                        xla::ParseMlirModuleString(mlir_module, context));

  // Serialize portable artifact
  std::string bytecode;
  {
    nb::gil_scoped_release gil_release;
    ABSL_ASSIGN_OR_RETURN(
        bytecode, xla::SerializeUsingVersionedStablehlo(
                      *module, target, sdy_version,
                      /*inplace=*/true,
                      /*allow_mixed_serialization=*/use_mixed_serialization));
  }
  return nb::bytes(bytecode.data(), bytecode.size());
}

absl::StatusOr<nb::bytes> PySerializePortableArtifact(
    PyModule& module, std::string_view target, std::string_view sdy_version,
    bool use_mixed_serialization) {
  mlir::ModuleOp module_op = unwrap(module.get());
  std::string bytecode;
  {
    nb::gil_scoped_release gil_release;
    ABSL_ASSIGN_OR_RETURN(
        bytecode, xla::SerializeUsingVersionedStablehlo(
                      module_op, target, sdy_version,
                      /*inplace=*/false,
                      /*allow_mixed_serialization=*/use_mixed_serialization));
  }
  return nb::bytes(bytecode.data(), bytecode.size());
}

absl::StatusOr<nb::object> PyDeserializePortableArtifact(
    const nb::bytes& bytecode_str, DefaultingPyMlirContext ctx) {
  MlirContext c_context = ctx->get();
  mlir::MLIRContext* context = unwrap(c_context);
  context->loadDialect<mlir::sdy::SdyDialect, mlir::mpmd::MpmdDialect>();
  std::string_view bytecode(bytecode_str.c_str(), bytecode_str.size());
  mlir::OwningOpRef<mlir::ModuleOp> module;
  {
    nb::gil_scoped_release gil_release;
    module = mlir::stablehlo::deserializePortableArtifact(bytecode, context);
  }
  if (!module)
    return tsl::errors::InvalidArgument("Failed to deserialize StableHLO");
  return PyModule::forModule(wrap(module.release())).releaseObject();
}

}  // namespace

void BuildMlirSubmodule(nb::module_& m) {
  nb::module_ mlir_module = m.def_submodule("mlir", "MLIR/XLA integration");

  mlir_module.def("hlo_to_stablehlo", xla::ValueOrThrowWrapper(HloToStableHlo),
                  nb::arg("computation"));

  mlir_module.def("xla_computation_to_mlir_module",
                  xla::ValueOrThrowWrapper(PyXlaComputationToMlirModule),
                  nb::arg("computation"),
                  nb::sig(
                      // clang-format off
                      "def xla_computation_to_mlir_module("
                      "computation: jaxlib._hlo.XlaComputation) -> str"  // NOLINT
                      // clang-format on
                      ));
  mlir_module.def(
      "mlir_module_to_xla_computation",
      [](const nb::bytes& bytecode, bool use_tuple_args, bool return_tuple) {
        return xla::ValueOrThrow(PyMlirModuleToXlaComputation(
            std::string_view(bytecode.c_str(), bytecode.size()), use_tuple_args,
            return_tuple));
      },
      nb::arg("mlir_module"), nb::arg("use_tuple_args") = false,
      nb::arg("return_tuple") = false,
      nb::sig(
          // clang-format off
      "def mlir_module_to_xla_computation("
      "mlir_module: bytes, "
      "use_tuple_args: bool = ..., "
      "return_tuple: bool = ..."
      ") -> jaxlib._hlo.XlaComputation"  // NOLINT
                            // clang-format on
          ));
  mlir_module.def("mlir_module_to_xla_computation",
                  xla::ValueOrThrowWrapper(PyMlirModuleToXlaComputation),
                  nb::arg("mlir_module"), nb::arg("use_tuple_args") = false,
                  nb::arg("return_tuple") = false,
                  nb::sig(
                      // clang-format off
      "def mlir_module_to_xla_computation("
      "mlir_module: str, "
      "use_tuple_args: bool = ..., "
      "return_tuple: bool = ..."
      ") -> jaxlib._hlo.XlaComputation"  // NOLINT
                                        // clang-format on
                      ));
  mlir_module.def(
      "mhlo_to_stablehlo",
      [](const nb::bytes& bytecode) {
        return xla::ValueOrThrow(PyMhloToStablehlo(
            std::string_view(bytecode.c_str(), bytecode.size())));
      },
      nb::arg("mlir_module"));
  mlir_module.def("mhlo_to_stablehlo",
                  xla::ValueOrThrowWrapper(PyMhloToStablehlo),
                  nb::arg("mlir_module"));
  mlir_module.def(
      "serialize_portable_artifact",
      [](nb::any mlir_module, std::string_view target,
         bool use_mixed_serialization) {
        // TODO(hyeontaek): Take a Sdy version as an argument to
        // `serialize_portable_artifact()`.
        const std::string sdy_version = xla::GetDefaultSdyVersion();
        if (nb::isinstance<PyModule>(mlir_module)) {
          PyModule& module = nb::cast<PyModule&>(mlir_module);
          return xla::ValueOrThrow(PySerializePortableArtifact(
              module, target, sdy_version, use_mixed_serialization));
        }
        if (nb::bytes bytecode;
            nb::try_cast<nb::bytes>(mlir_module, bytecode)) {
          return xla::ValueOrThrow(PySerializePortableArtifact(
              std::string_view(bytecode.c_str(), bytecode.size()), target,
              sdy_version, use_mixed_serialization));
        }
        if (std::string str_module;
            nb::try_cast<std::string>(mlir_module, str_module)) {
          return xla::ValueOrThrow(PySerializePortableArtifact(
              str_module, target, sdy_version, use_mixed_serialization));
        }
        throw nb::type_error("mlir_module must be bytes, str, or MlirModule");
      },
      nb::arg("mlir_module"), nb::arg("target"),
      nb::arg("use_mixed_serialization") = false,
      nb::sig("def serialize_portable_artifact(mlir_module: typing.Any, "
              "target: str, use_mixed_serialization: bool = False) -> bytes"));
  mlir_module.def("deserialize_portable_artifact",
                  xla::ValueOrThrowWrapper(PyDeserializePortableArtifact),
                  nb::arg("mlir_module"), nb::arg("context") = nb::none(),
                  nb::sig("def deserialize_portable_artifact(mlir_module: "
                          "bytes, context: typing.Any = None) -> typing.Any"));
  mlir_module.def(
      "refine_polymorphic_shapes",
      [](nb::bytes bytecode, bool enable_shape_assertions,
         bool validate_static_shapes, bool enable_shardy) -> nb::bytes {
        std::string buffer;
        llvm::raw_string_ostream os(buffer);
        xla::ThrowIfError(xla::RefinePolymorphicShapes(
            std::string_view(bytecode.c_str(), bytecode.size()), os,
            enable_shape_assertions, validate_static_shapes, enable_shardy));
        return nb::bytes(buffer.data(), buffer.size());
      },
      nb::arg("mlir_module"), nb::arg("enable_shape_assertions") = true,
      nb::arg("validate_static_shapes") = true,
      nb::arg("enable_shardy") = false,
      R"(Refines the dynamic shapes for a module.
        The "main" function must have static shapes and all the
        intermediate dynamic shapes depend only on the input static
        shapes. Optionally, also validates that the resulting module has
        only static shapes.
      )");
}

}  // namespace jax