/* Copyright 2017 The OpenXLA 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.
==============================================================================*/

#ifndef XLA_SERVICE_GPU_GPU_COMPILER_H_
#define XLA_SERVICE_GPU_GPU_COMPILER_H_

#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <vector>

#include "absl/base/nullability.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "llvm/IR/Module.h"
#include "mlir/IR/MLIRContext.h"
#include "xla/autotune_results.pb.h"
#include "xla/backends/autotuner/codegen_backend.h"
#include "xla/hlo/ir/hlo_module.h"
#include "xla/hlo/pass/hlo_pass_pipeline.h"
#include "xla/hlo/transforms/simplifiers/algebraic_simplifier.h"
#include "xla/pjrt/distributed/key_value_store_interface.h"
#include "xla/service/compiled_module.h"
#include "xla/service/compiler.h"
#include "xla/service/executable.h"
#include "xla/service/gpu/alias_info.h"
#include "xla/service/gpu/autotuning/autotuner_util.h"
#include "xla/service/gpu/compile_module_to_llvm_ir.h"
#include "xla/service/gpu/ir_emission_utils.h"
#include "xla/service/hlo.pb.h"
#include "xla/service/hlo_cost_analysis.h"
#include "xla/service/hlo_module_config.h"
#include "xla/service/llvm_compiler.h"
#include "xla/stream_executor/device_description.h"
#include "xla/stream_executor/device_description.pb.h"
#include "xla/stream_executor/dnn.h"
#include "xla/stream_executor/kernel_stats.h"
#include "xla/stream_executor/platform.h"
#include "xla/stream_executor/semantic_version.h"
#include "xla/stream_executor/stream_executor.h"
#include "xla/tsl/platform/threadpool.h"
#include "xla/util.h"
#include "xla/xla.pb.h"

namespace xla {
namespace gpu {

// The GPU compiler generates efficient GPU executables.
class GpuCompiler : public LLVMCompiler {
 public:
  GpuCompiler(se::Platform::Id platform_id, const char* target_triple,
              const char* data_layout);

  using LLVMCompiler::Compile;

  // An attached device is passed in via stream_exec. We get GPU configuration
  // from the attached device OR from the `options` struct (in which case the
  // attached device is ignored during the compilation).
  // If you call this directly, follow it with RunBackend rather than Compile.
  absl::StatusOr<std::unique_ptr<HloModule>> RunHloPasses(
      std::unique_ptr<HloModule> module, se::StreamExecutor* stream_exec,
      const CompileOptions& options) override;

  absl::StatusOr<std::unique_ptr<Executable>> RunBackend(
      std::unique_ptr<HloModule> module, se::StreamExecutor* stream_exec,
      const CompileOptions& options) override;

  absl::StatusOr<std::vector<std::unique_ptr<CompiledModule>>>
  CompileAheadOfTime(std::unique_ptr<HloModule> hlo_module,
                     AotCompilationOptions const& options) override;

  se::Platform::Id PlatformId() const override { return platform_id_; }

  HloCostAnalysis::ShapeSizeFunction ShapeSizeBytesFunction() const override;

  // Returns a (deserialized) AotCompilationResult from a serialized
  // AotCompilationResult.
  absl::StatusOr<std::unique_ptr<CompiledModule>> LoadAotCompilationResult(
      const std::string& serialized_aot_result) override;

  absl::StatusOr<std::unique_ptr<CompiledModule>> Export(
      Executable* executable) override;

  absl::Status RunPostSchedulingPipelines(
      HloModule* module, int64_t scheduler_mem_limit,
      const se::DeviceDescription& gpu_device_info,
      const GpuAliasInfo* alias_info);

  std::string target_triple() const { return target_triple_; }
  std::string data_layout() const { return data_layout_; }

  const char* GetDataLayout() const { return data_layout_; }

  const char* GetTargetTriple() const { return target_triple_; }

  int64_t GetPointerSize() const { return pointer_size_; }

  static absl::StatusOr<GpuTargetConfig> GetTargetConfig(
      const Compiler::CompileOptions& options, const DebugOptions& debug_opts,
      se::StreamExecutor* executor);

  mlir::MLIRContext* mlir_context() { return &mlir_context_; }

  virtual std::unique_ptr<GpuAliasInfo> GetAliasInfo(
      const se::DeviceDescription& device_description) const {
    return std::make_unique<GpuAliasInfo>(device_description);
  }

  virtual absl::StatusOr<bool> CanUseLinkModules(
      const HloModuleConfig& config,
      const stream_executor::DeviceDescription& device_description) {
    return false;
  }

  enum class AlgebraicSimplifierMode {
    kLayoutInsensitive,
    kPostFusionSimplification,
    kLayoutNormalization,
    kPostLayoutAssignment,
    kAfterSimplifyFPConversions,
    kGpuConvoluationCanonicalization,
  };

  static AlgebraicSimplifierOptions GetAlgebraicSimplifierOptions(
      AlgebraicSimplifierMode mode, const DebugOptions& debug_options,
      bool is_rocm);

  absl::StatusOr<std::unique_ptr<Executable>> LoadExecutableFromAotResult(
      const CompiledModule& aot_result,
      const se::DeviceDescription& device_description) override;

  static std::unique_ptr<HloPassPipeline> GetCublasRewriterPipeline(
      const stream_executor::DeviceDescription& device_description,
      bool enable_cublaslt = false);

  static std::unique_ptr<HloPassPipeline> GetCustomKernelRewriterPipeline(
      const stream_executor::DeviceDescription& device_description);

  // Returns the LLVM command line options that we use for compilation.
  // THey need to be set globally whenever we call into LLVM.
  virtual std::vector<std::string> GetLLVMCommandLineOptions(
      const DebugOptions& debug_options) const = 0;

 protected:
  struct BackendCompileResult {
    std::string asm_text;
    std::vector<uint8_t> binary;
    BinaryMap dnn_compiled_graphs;
    ModuleStats module_stats;
  };

  // During compilation with device, stream_exec != null and autotune_results
  // == null. During deviceless AOT compilation, stream_exec == null and
  // autotune_results != null.
  // thread_pool is used to speed up compilation during autotuning.
  virtual absl::Status OptimizeHloPostLayoutAssignment(
      HloModule* hlo_module, se::StreamExecutor* stream_exec,
      const CompileOptions& options, const GpuTargetConfig& gpu_target_config,
      const GpuAliasInfo* alias_info, tsl::thread::ThreadPool* thread_pool);

  // Add autotuning passes for convolution and gemm.
  // target_config must outlive the pipeline.
  virtual absl::Status AddConvAndGemmAutotuningPass(
      HloPassPipeline* pipeline, HloModule* hlo_module,
      const se::GpuComputeCapability& gpu_version,
      const CompileOptions& options, AutotuneConfig& autotune_config,
      tsl::thread::ThreadPool* thread_pool, se::StreamExecutor* stream_exec,
      const Compiler::GpuTargetConfig* target_config,
      const MultiProcessKeyValueStore& key_value_store,
      const se::SemanticVersion& toolkit_version, const AliasInfo* alias_info,
      const DebugOptions& debug_options, mlir::MLIRContext* mlir_context,
      HloCostAnalysis::ShapeSizeFunction shape_size_fn);

  virtual absl::StatusOr<std::vector<std::unique_ptr<CodegenBackend>>>
  GetCodegenBackends(se::StreamExecutor* stream_exec,
                     const Compiler::GpuTargetConfig* target_config,
                     const AliasInfo* alias_info,
                     const DebugOptions& debug_options,
                     mlir::MLIRContext* mlir_context) {
    return std::vector<std::unique_ptr<CodegenBackend>>();
  }

  // target_config must outlive the pipeline.
  virtual absl::Status AddFusionAutotuningPass(
      HloPassPipeline* pipeline, HloModule* hlo_module,
      const CompileOptions& options, tsl::thread::ThreadPool* thread_pool,
      stream_executor::StreamExecutor* stream_executor,
      const GpuTargetConfig* target_config,
      HloCostAnalysis::ShapeSizeFunction shape_size_fn,
      const MultiProcessKeyValueStore& key_value_store) {
    return absl::OkStatus();
  }

  // Runs cuDNN fusion and custom call compiler passes.
  virtual absl::Status RunCudnnCompilerPasses(HloModule* module,
                                              se::StreamExecutor* stream_exec,
                                              BinaryMap* dnn_compiled_graphs) {
    return absl::OkStatus();
  }

 private:
  struct CompileResultWithMetadata {
    BackendCompileResult backend_result;
    CompileModuleResults compile_module_results;
  };

  // Schedule and compile the module.
  absl::StatusOr<CompileResultWithMetadata> CompileToBackendResult(
      HloModule* module, llvm::LLVMContext* llvm_context,
      const CompileOptions& options,
      const se::DeviceDescription& gpu_device_info);

  absl::StatusOr<BackendCompileResult> CompileAndLink(
      const HloModuleConfig& module_config,
      CompileModuleResults& compile_module_results,
      const stream_executor::DeviceDescription& device_description,
      const CompileOptions& options, const HloModule* debug_module);

  absl::StatusOr<BackendCompileResult> CompileSingleModule(
      const HloModuleConfig& module_config,
      const stream_executor::DeviceDescription& device_description,
      const HloModule* debug_module, llvm::Module* llvm_module,
      bool relocatable, const CompileOptions& options,
      std::optional<int> shard_number);

  absl::Status LoadAutotuneResultsFromFile(const DebugOptions& debug_options);
  absl::Status SerializeAutotuneResultsToFile(
      const DebugOptions& debug_options);

  absl::Status RunPreSchedulingPasses(
      HloModule* module, const se::DeviceDescription& gpu_device_info,
      const GpuAliasInfo* alias_info);
  absl::Status RunCollectiveScheduleLinearizerPasses(
      HloModule* hlo_module, se::StreamExecutor* stream_exec);

  // During compilation with device, stream_exec != null and autotune_results
  // == null. During deviceless AOT compilation, stream_exec == null and
  // autotune_results != null.
  absl::Status OptimizeHloModule(HloModule* hlo_module,
                                 se::StreamExecutor* stream_exec,
                                 const CompileOptions& options,
                                 const GpuTargetConfig& gpu_target_config,
                                 const GpuAliasInfo* alias_info);

  virtual absl::Status OptimizeHloConvolutionCanonicalization(
      HloModule* hlo_module, const se::GpuComputeCapability& gpu_version,
      se::dnn::VersionInfo dnn_version,
      const se::SemanticVersion& toolkit_version) = 0;

  // TODO(timshen): Replace `debug_module` with some portable debug information
  // that accommodates both HLO and MLIR.
  virtual absl::StatusOr<BackendCompileResult> CompileTargetBinary(
      const HloModuleConfig& module_config, llvm::Module* llvm_module,
      const stream_executor::DeviceDescription& device_description,
      bool relocatable, const HloModule* debug_module,
      const CompileOptions& options, std::optional<int> shard_number) = 0;

  // Inserts and optimizes mandatory copies. Necessary for correctness.
  absl::Status RunPreSchedulingCopyInsertion(
      HloModule& hlo_module, const se::DeviceDescription& device_description,
      const GpuAliasInfo* alias_info);

  virtual absl::StatusOr<std::vector<uint8_t>> LinkModules(
      const stream_executor::DeviceDescription& device_description,
      std::vector<std::vector<uint8_t>> modules,
      const DebugOptions& debug_options) {
    return Unimplemented("LinkModules is not implemented.");
  }

  // Runs HLO passes on the given module. If the module has a schedule, it is
  // assumed that the module is already optimized and no passes are run.
  absl::StatusOr<std::unique_ptr<HloModule>> RunHloPassesIfNeeded(
      std::unique_ptr<HloModule> hlo_module,
      se::StreamExecutor* absl_nullable executor,
      const CompileOptions& compile_options);

  // New AOT compilation which compiles up the the Thunk generation stage.
  absl::StatusOr<std::vector<std::unique_ptr<CompiledModule>>>
  NewCompileAheadOfTime(std::unique_ptr<HloModule> hlo_module,
                        se::StreamExecutor* executor,
                        const CompileOptions& compile_options);
  // Legacy AOT compilation.
  absl::StatusOr<std::vector<std::unique_ptr<CompiledModule>>>
  LegacyCompileAheadOfTime(std::unique_ptr<HloModule> hlo_module,
                           const AotCompilationOptions& options);

  se::Platform::Id platform_id_;

  // The triple that represents our target.
  const char* target_triple_;

  // The data layout of the emitted module.
  const char* data_layout_;

  // The size in bytes of a pointer. Used by ShapeSizeBytesFunction.
  const int64_t pointer_size_;

  GpuCompiler(const GpuCompiler&) = delete;
  GpuCompiler& operator=(const GpuCompiler&) = delete;

  // A MLIR context that can be used by pre-codegen passes. For codegen, we will
  // need to have a context with more dialects registered.
  mlir::MLIRContext mlir_context_;
};

}  // namespace gpu
}  // namespace xla

#endif  // XLA_SERVICE_GPU_GPU_COMPILER_H_