// Copyright 2023 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//===----------------------------------------------------------------------===//
//
// This file includes the LLVMGPU Passes.
//
//===----------------------------------------------------------------------===//

#ifndef IREE_COMPILER_CODEGEN_LLVMGPU_PASSES_H_
#define IREE_COMPILER_CODEGEN_LLVMGPU_PASSES_H_

#include <optional>

#include "iree/compiler/Codegen/Common/GPU/Passes.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
#include "iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.h"
#include "iree/compiler/Codegen/Utils/CodegenOptions.h"
#include "iree/compiler/Codegen/Utils/CodegenPipelineOptions.h"
#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
#include "mlir/Pass/Pass.h"

namespace mlir::iree_compiler {

using IREE::GPU::GPUPipelineOptions;

//----------------------------------------------------------------------------//
// LLVMGPU Backend Pass Pipelines
//----------------------------------------------------------------------------//

/// Simple lowering only distributute linalg ops on blocks and threads. This
/// will result in scalar operations. Expects pass manager to be a
/// module-level pass manager.
void addGPUSimpleDistributePassPipeline(OpPassManager &funcPassManager);

/// Lowering config driven pipeline that uses greedy tile + fuse to distribute
/// to threads.
void addGPUTileAndFusePassPipeline(OpPassManager &funcPassManager,
                                   const GPUPipelineOptions &pipelineOptions,
                                   bool forROCDL);

/// Transform dialect-based path.
void addGPUTransformDialectPasses(OpPassManager &funcPassManager,
                                  StringRef entryPoint);

/// Lowering calling vectorization patterns. Expects pass manager to be a
/// module-level pass manager.
void addGPUVectorizationPassPipeline(OpPassManager &funcPassManager);

/// Lowering for winograd transform ops. Follows `VectorizationPassPipeline`
/// with different tiling and distribution passes.
void addGPUWinogradVectorizePassPipeline(OpPassManager &funcPassManager);

/// Lowering based on vector distribution patterns.
void addGPUVectorDistributePassPipeline(OpPassManager &funcPassManager,
                                        const GPUPipelineOptions &options,
                                        bool forROCDL);

/// Default pass pipeline on GPU, currently used only for the ukernel path.
void addGPUDefaultPassPipeline(OpPassManager &funcPassManager,
                               const GPUPipelineOptions &options);

/// Pass pipeline to lower IREE HAL executables without tiling and distribution.
void addGPUBaseLoweringPassPipeline(OpPassManager &pm);

/// Populates the common passes needed to preprocess and select the translation
/// strategy.
void buildLLVMGPUCodegenCommonConfigurationPassPipeline(
    OpPassManager &modulePassManager);

/// Populates passes needed to preprocess and select the translation strategy.
void buildLLVMGPUCodegenConfigurationPassPipeline(
    OpPassManager &modulePassManager);

/// Populates passes needed to lower a XLA HLO op to NVVM/ROCDL dialect via
/// the structured ops path. The `modulePassManager` should operate on the
/// module within the IREE::HAL::ExecutableOp.
void buildLLVMGPUCodegenPassPipeline(OpPassManager &modulePassManager,
                                     bool useROCM, bool preserveDebugInfo,
                                     bool includeLLVMLowering = true,
                                     bool useSPIRV = false);

/// Wraps GPUPipelineOptions and forROCDL for passing through
/// PipelineAttrInterface::buildPipeline.
struct GPUCodegenPipelineOptions final
    : CodegenPipelineOptionsBase<GPUCodegenPipelineOptions> {
  GPUCodegenPipelineOptions(const GPUPipelineOptions &options, bool forROCDL)
      : options(options), forROCDL(forROCDL) {}

  GPUPipelineOptions options;
  bool forROCDL = false;
};

/// Verify configuration set for the LLVMGPUVectorDistribute pass pipeline.
LogicalResult verifyLLVMGPUVectorDistributePipeline(
    Operation *op, IREE::GPU::LoweringConfigAttr loweringConfig);

//----------------------------------------------------------------------------//
// LLVMGPU Linking Passes and Pipelines
//----------------------------------------------------------------------------//

/// Populates passes needed to link HAL executables across LLVMGPU targets.
void buildLLVMGPULinkingPassPipeline(
    OpPassManager &modulePassManager,
    std::optional<std::string> target = std::nullopt);

//------------------------------------------------------------------------------
// Wrappers that do not use tablegen options
//------------------------------------------------------------------------------

enum class GPUTensorCoreType {
  WMMA = 0,
  MMA_SYNC = 1,
};

std::unique_ptr<InterfacePass<FunctionOpInterface>>

createLLVMGPUTensorCoreVectorizationPass(GPUTensorCoreType tensorCoreType);
std::unique_ptr<InterfacePass<FunctionOpInterface>>
createLLVMGPUVectorToGPUPass(GPUTensorCoreType tensorCoreType);

std::unique_ptr<InterfacePass<FunctionOpInterface>>
createLLVMGPUTileAndDistributePass(bool distributeToWarp);

//----------------------------------------------------------------------------//
// Register LLVMGPU Passes
//----------------------------------------------------------------------------//

#define GEN_PASS_DECL
#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc" // IWYU pragma: keep

void registerCodegenLLVMGPUPasses();

} // namespace mlir::iree_compiler

#endif // IREE_COMPILER_CODEGEN_LLVMGPU_PASSES_H_
