// Copyright 2024 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

#ifndef IREE_COMPILER_CODEGEN_DIALECT_GPU_TARGETUTILS_CONFIGUTILS_H_
#define IREE_COMPILER_CODEGEN_DIALECT_GPU_TARGETUTILS_CONFIGUTILS_H_

#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
#include "mlir/IR/Operation.h"
#include "mlir/Interfaces/FunctionInterfaces.h"

namespace mlir::iree_compiler::IREE::GPU {

/// Helper for setting up a data tiled multi-MMA inner_tiled config based on the
/// specified target.
LogicalResult setDataTiledMmaInnerTiledLoweringConfig(
    IREE::GPU::TargetAttr target, mlir::FunctionOpInterface entryPoint,
    Operation *op, IREE::Codegen::UKernelDescriptorAttr ukernelConfig,
    std::optional<uint64_t> prefetchNumStages);

/// Helper for setting up a convolution config using direct convolution based on
/// the specified target.
LogicalResult setDirectConvolutionLoweringConfig(
    IREE::GPU::TargetAttr target, mlir::FunctionOpInterface entryPoint,
    Operation *op, std::optional<uint64_t> prefetchNumStages);

/// Helper for setting up a convolution config using IGEMM based on the
/// specified target.
/// TODO: Currently this only succeeds if the target supports an mma
/// kind. Add support for a fallback direct lowering path.
LogicalResult setIGEMMConvolutionLoweringConfig(
    IREE::GPU::TargetAttr target, mlir::FunctionOpInterface entryPoint,
    Operation *op, bool useDirectLoad, bool padConv,
    std::optional<uint64_t> prefetchNumStages);

/// Helper for setting up a matmul config based on the specified target.
/// TODO: Currently this only succeeds if the target supports an mma
/// kind. Add support for a fallback direct lowering path.
LogicalResult
setMatmulLoweringConfig(IREE::GPU::TargetAttr target,
                        mlir::FunctionOpInterface entryPoint, Operation *op,
                        bool useDirectLoad,
                        std::optional<uint64_t> prefetchNumStages);

/// Helper for setting up a default tile and fuse config for targeting
/// simple thread distribution. Currently restricted to linalg ops.
LogicalResult setTileAndFuseLoweringConfig(IREE::GPU::TargetAttr target,
                                           mlir::FunctionOpInterface entryPoint,
                                           Operation *op);

// Helper for setting tile sizes for scatter.
LogicalResult setScatterLoweringConfig(IREE::GPU::TargetAttr target,
                                       mlir::FunctionOpInterface entryPoint,
                                       Operation *op);

LogicalResult setSortConfig(IREE::GPU::TargetAttr target,
                            mlir::FunctionOpInterface entryPoint,
                            Operation *op);

/// Helper for setting tile sizes for arg_compare. Targets the TileAndFuse
/// pipeline so the parallel iteration space is distributed across threads via
/// `scf.forall` + `GPUDistributeForall`, while the reduction dimension stays
/// untiled (each thread runs the deterministic scan on its own slice).
LogicalResult setArgCompareConfig(IREE::GPU::TargetAttr target,
                                  mlir::FunctionOpInterface entryPoint,
                                  Operation *op);

/// Helper for setting up a memory bound reduction configuration, focusing
/// on getting peak global memory bandwidth. Supports linalg::LinalgOp and
/// LinalgExt::ArgCompareOp.
LogicalResult setReductionConfig(IREE::GPU::TargetAttr target,
                                 mlir::FunctionOpInterface entryPoint,
                                 Operation *op);

//===----------------------------------------------------------------------===//
// Pass Pipeline Options
//===----------------------------------------------------------------------===//

using IREE::GPU::ReorderWorkgroupsStrategy;

struct GPUPipelineOptions {
  bool enableReduceSharedMemoryBankConflicts = true;
  // Number of software pipelining stages for shared memory prefetching.
  // 0 or 1 disables prefetching, 2+ enables it with that many stages.
  unsigned prefetchNumStages = 0;
  bool useIgemmConvolution = false;
  bool enableUkernels = false;
  std::optional<ReorderWorkgroupsStrategy> reorderStrategy;
};

llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
                              const GPUPipelineOptions &options);

GPUPipelineOptions
getPipelineOptions(FunctionOpInterface funcOp,
                   IREE::Codegen::TranslationInfoAttr translationInfo);

} // namespace mlir::iree_compiler::IREE::GPU

#endif // IREE_COMPILER_CODEGEN_DIALECT_GPU_TARGETUTILS_CONFIGUTILS_H_
