// Copyright 2021 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_UTILS_GPUUTILS_H_
#define IREE_COMPILER_CODEGEN_UTILS_GPUUTILS_H_

#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUInterfaces.h"
#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Interfaces/FunctionInterfaces.h"

namespace mlir::iree_compiler {

static constexpr int32_t kNumGPUDims = 3;
static constexpr int32_t kWarpSize = 32;

// Universally-safe ceiling for `gpu.shuffle`-based codegen: at this width or
// below, every backend's `gpu.shuffle` lowering emits working code. Wider
// values rely on the backend decomposing them, which today only ROCDL does
// (cf. llvm/llvm-project#136320; NVVM tracking at llvm/llvm-project#197080).
// See `targetSupportsShuffleBitwidth`.
constexpr unsigned kShuffleNativeBits = 32;

//===----------------------------------------------------------------------===//
// GPU processor IDs and sizes
//===----------------------------------------------------------------------===//

llvm::SmallVector<linalg::ProcInfo, 2>
getGPUThreadIdsAndCounts(OpBuilder &builder, Location loc, unsigned numDims,
                         llvm::ArrayRef<int64_t> workgroupSize);

/// Computes subgroup ID and returns in (X, Y, Z) order.
///
/// Note that CUDA doesn't have a subgroupId equivalent so we are are computing
/// the subgroup ID based on the threadID. When tiling to warp we assume each
/// warp is full and we pick a workgroup size so that `workgroupSize.x %
/// warpSize == 0`. This is why we can have warpId = { threadId.x / warpSize,
/// threadId.y, threadId.z }.
llvm::SmallVector<linalg::ProcInfo, 2>
getSubgroupIdsAndCounts(OpBuilder &builder, Location loc, unsigned warpSize,
                        unsigned numDims, llvm::ArrayRef<int64_t> numSubgroups);

/// Indicates whether the given array of DeviceMappingAttrInterfaces is a
/// descending relative mapping, for example:
///  [#gpu.thread<z>, #gpu.thread<y>, #gpu.thread<x>]
/// or
///  [#gpu.thread<linear_dim_1>, #gpu.thread<linear_dim_0>]
bool isDescendingRelativeMappingIndices(ArrayRef<Attribute> array);

// Indicates whether the given `scf.forall` op has a processor ID mapping of
// the template type(s).
template <typename... Type>
bool forallOpHasMappingType(scf::ForallOp forallOp) {
  std::optional<ArrayAttr> mapping = forallOp.getMapping();
  if (!mapping || mapping.value().empty()) {
    return false;
  }

  return isa<Type...>(*mapping.value().begin());
}

// Indicates whether an operation is within a distributed context with the
// specified mapping type(s).
template <typename... Type>
bool operationHasParentForallOfMappingType(Operation *op) {
  auto parentForallOp = op->getParentOfType<scf::ForallOp>();
  while (parentForallOp) {
    if (forallOpHasMappingType<Type...>(parentForallOp)) {
      return true;
    }
    parentForallOp = parentForallOp->getParentOfType<scf::ForallOp>();
  }
  return false;
}

//===----------------------------------------------------------------------===//
// GPU vectorization
//===----------------------------------------------------------------------===//

/// Returns true if we can use all threads to perform vectorized load/store of
/// the given `shape`.
bool canPerformVectorAccessUsingAllThreads(ArrayRef<int64_t> shape,
                                           int64_t threadCount,
                                           int64_t vectorSize);

/// Pick an unrolling order that will allow tensorcore operation to reuse LHS
/// register. This is needed to get good performance on sm_80 target.
std::optional<SmallVector<int64_t>>
gpuMmaUnrollOrder(vector::ContractionOp contract);

//===----------------------------------------------------------------------===//
// GPU tiling and distribution
//===----------------------------------------------------------------------===//

/// Returns the attribute name carrying information about distribution.
const char *getGPUDistributeAttrName();

/// Returns the tile sizes at the given `tilingLevel` for compute ops in
/// `funcOp`.
FailureOr<SmallVector<int64_t>> getGPUTileSize(mlir::FunctionOpInterface funcOp,
                                               int tilingLevel);

/// Returns the functor to compute tile sizes at the given `tilingLevel` for
/// compute ops in `funcOp`.
FailureOr<scf::SCFTileSizeComputationFunction>
getGPUScfTileSizeComputeFn(mlir::FunctionOpInterface funcOp, int tilingLevel);

/// Returns true iff the rank of the input value 'val' is non-zero.
bool isNonZeroRank(TypedValue<VectorType> val);

//===----------------------------------------------------------------------===//
// GPU workgroup memory
//===----------------------------------------------------------------------===//

/// Allocates GPU workgroup memory matching the given `subview`. If there are
/// dynamic dimensions, the bounds are in `sizeBounds`.
std::optional<Value> allocateWorkgroupMemory(OpBuilder &builder,
                                             memref::SubViewOp subview,
                                             ArrayRef<Value> sizeBounds,
                                             DataLayout &);

/// Deallocates GPU workgroup memory behind `buffer`.
LogicalResult deallocateWorkgroupMemory(OpBuilder &, Value buffer);

/// Copies `src` value to `dst` in shared memory.
LogicalResult copyToWorkgroupMemory(OpBuilder &builder, Value src, Value dst);

/// Propagates shared memory copy to producer linalg.fill or consumer
/// linalg.generic when possible.
void propagateSharedMemoryCopy(mlir::FunctionOpInterface funcOp);

/// Inserts barriers before and after shared memory copy.
void insertBarriersAroundSharedMemoryCopy(mlir::FunctionOpInterface funcOp);

/// Emit reduction across a group for a given input. Emits `gpu.shuffle`
/// based reduction only when `expandSubgroupReduce` is set.
Value emitGPUGroupReduction(Location loc, OpBuilder &builder, Value input,
                            vector::CombiningKind kind, uint32_t size,
                            int warpSize, bool expandSubgroupReduce);

/// Return the native size of an operation used in contraction calculation.
// TODO: Make this take HW specific sizes.
std::optional<SmallVector<int64_t>> getWmmaNativeVectorSize(Operation *op);

/// Helper function to return native size for MMA.SYNC-based operations.
std::optional<SmallVector<int64_t>> getMmaNativeVectorSize(Operation *op);

/// Rerutn true if memref has #amdgpu.address_space<fat_raw_buffer> address
/// space.
bool hasAMDGPUFatRawBufferAddressSpace(MemRefType memrefType);

/// Return true if the given memref has one of the global address spaces - no
/// address space, explicit integer 0, #gpu.address_space<global>, or
/// #amdgpu.address_space<fat_raw_buffer>
bool hasGlobalMemoryAddressSpace(MemRefType memrefType);

/// Return true if the given memref has workgroup memory space.
bool hasSharedMemoryAddressSpace(MemRefType memrefType);

/// Packs vector of lower precision into a single 32-bit width element.
/// (i.e <2xf16> -> i32 and <4xi8> -> i32)
Value packVectorToSupportedWidth(Location loc, OpBuilder &builder, Value input);

/// Unpack single scalar element into a target vector type.
/// (i.e i32 -> vector<4xi8> or f32 -> vector<2xf16>)
Value unpackToVector(Location loc, OpBuilder &builder, Value packedInput,
                     VectorType targetVecType);

/// Emit identity constant based on combiningKind and type.
Value getCombiningIdentityValue(Location loc, OpBuilder &builder,
                                vector::CombiningKind kind, Type identityType);

/// Returns the matching GPU reduction operation.
mlir::gpu::AllReduceOperation
combiningKindToAllReduce(vector::CombiningKind kind);

/// This struct matches the parameters of the `iree_codegen.xor_shuffle`
/// attribute. Notably only the first two parameters are used in this struct
/// because the other parameters are not used in the current implementation of
/// the XOR shuffle.
/// The parameters are:
/// - rowElems: number of elements across which we apply a swizzle.
/// - accessElems: number of elements in one swizzle group.
/// So for example:
/// For data of length 8: [0][1][2][3][4][5][6][7]
/// and an XOR shuffle with rowElems = 4, accessElems = 2,
/// we can apply a swizzle to each logical row of 4 elements in groups of 2
/// elements.
///
/// Original data:
/// |               | swizzle group 0   | swizzle group 1   |
/// |---------------|-------------------|-------------------|
/// Logical row 0:  | [0][1]            | [2][3]            |
/// Logical row 1:  | [4][5]            | [6][7]            |
///
/// Swizzled data:
/// |               | swizzle group 0   | swizzle group 1   |
/// |---------------|-------------------|-------------------|
/// Logical row 0:  | [0][1]            | [2][3]            |
/// Logical row 1:  | [6][7]            | [4][5]            |
///
/// Final result: [0][1][2][3][6][7][4][5]
/// Note that when the swizzle was applied, the swaps were performed on
/// groups of 2 elements at a time (accessElems), across the 4 elements in the
/// row (rowElems).
struct XorShuffleParams {
  int64_t rowElems;
  int64_t accessElems;
};

/// Bounds for valid XOR shuffle parameters (min access elements per thread,
/// total elements in the tile). Used when sweeping over valid configs.
struct XorShuffleBounds {
  int64_t minAccessElems;
  int64_t totalTileElems;
};

/// For a given MMA intrinsic and operand, returns the lower bound and upper
/// bound for valid values of XOR shuffle attribute parameters, access width and
/// row width. For both parameters, the elements ingested per thread at a time
/// is used as the minimum bound and the total number of elements in the tile is
/// used as the upper bound. Note that if you want to do a sweep over valid XOR
/// swizzles, this is how the sweep should be done:
/// - sweep access elements over all multiples of the minimum bound, respecting
/// the upper bound.
/// - sweep row elements over all multiple of the access elements, respecting
/// the upper bound.
FailureOr<XorShuffleBounds>
getXorShuffleBounds(IREE::Codegen::InnerTileDescAttrInterface intrinsic,
                    int operandIndex);

/// Returns true if the XOR shuffle is valid for the given number of row
/// elements, number of access elements, and total tile elements.
/// Else returns false.
/// An XOR Shuffle can be invalid for many reasons:
/// - The number of row elements exceeds the total elements in the tile.
/// - The number of access elements exceeds the total elements in the row.
/// - The number of row elements does not evenly divide the total elements in
/// the tile.
/// - The number of access elements does not evenly divide the number of row
/// elements.
/// - The number of columns in a row does not evenly divide the total number of
/// tile elements.
bool isXORShuffleValid(int64_t numRowElems, int64_t numAccessElems,
                       int64_t totalTileElems);

/// Returns XOR shuffle parameters known to be optimal for the given target and
/// intrinsic based on profiling results.
FailureOr<XorShuffleParams> getXorShuffleParamsForTunedChipset(
    IREE::GPU::TargetAttr target,
    IREE::Codegen::InnerTileDescAttrInterface intrinsic, int operandIndex);

/// Note this generic heuristic for untuned cases is not guaranteed to be
/// optimal for all targets and intrinsics.
FailureOr<XorShuffleParams> getXorShuffleParamsForUntunedChipset(
    IREE::GPU::TargetAttr target,
    IREE::Codegen::InnerTileDescAttrInterface intrinsic,
    ArrayRef<int64_t> reductionTileSizes, int operandIndex);

/// Returns the XOR shuffle parameters (row elements and access elements) for
/// the given target, intrinsic, and operand index. For some targets and
/// intrinsics, the optimal XOR shuffle for a given operand might be known. For
/// these cases of "Tuned" targets and intrinsics, a look up table is used to
/// obtain values for the XOR shuffle attributes. Otherwise, these values are
/// computed by calculating the K tile size and LDS bank width. Note this
/// generic heuristic for untuned cases is not guaranteed to be optimal for all
/// targets and intrinsics.
FailureOr<XorShuffleParams>
getXorShuffleParams(IREE::GPU::TargetAttr target,
                    IREE::Codegen::InnerTileDescAttrInterface intrinsic,
                    ArrayRef<int64_t> reductionTileSizes, int operandIndex,
                    bool skipUntunedFallback = false);

/// Returns the XOR shuffle attribute for the given target, intrinsic, and
/// operand index.
FailureOr<Attribute>
getXorShuffleAttr(MLIRContext *context, Attribute baseConfigAttr,
                  IREE::GPU::TargetAttr target,
                  IREE::Codegen::InnerTileDescAttrInterface intrinsic,
                  ArrayRef<int64_t> reductionTileSizes, int operandIndex,
                  bool skipUntunedFallback = false);

/// Apply inverse XOR swizzle to a sub-tile-local source offset so that the
/// DMA write-side permutation matches the read-side (ResolveSwizzleHints).
///
/// When the subgroup transfer size is not a multiple of the swizzle period,
/// we add the subgroup's base offset within the full allocation before
/// swizzling and subtract it after.
///
///    Example: xor_shuffle<64, 8>, period = 64*64/8 = 512 elements.
///    Workgroup tile = 32x32, 4 subgroups of 8x32 = 256 each.
///    256 < 512, so the fix is needed:
///    ```
///      BUG: swizzle(local)
///           Subgroups 0 and 1 see the same local offsets but occupy
///           different rows in the full allocation.
///      FIX: swizzle(local + base) - base
///    ```
Value applyInverseXorSwizzleToDMASourceOffset(
    OpBuilder &builder, Location loc, Value srcLinearOffset,
    IREE::Codegen::XORShuffleAttr swizzle, Value dest);

//===----------------------------------------------------------------------===//
// GPU CodeGen op filter
//===----------------------------------------------------------------------===//

/// Returns true if the index map represents a transpose that benefits from
/// using shared memory when CodeGen towards the GPU.
bool sharedMemTransposeFilter(AffineMap indexMap);

//===----------------------------------------------------------------------===//
// GPU Target Information
//===----------------------------------------------------------------------===//
FailureOr<ArrayAttr> getSupportedMmaTypes(DictionaryAttr config);

FailureOr<ArrayAttr> getSupportedMmaTypes(mlir::FunctionOpInterface entryPoint);

/// Returns the GPU target attribute from `iree-gpu-test-target` if provided.
/// Returns null TargetAttr otherwise.
IREE::GPU::TargetAttr getCLGPUTarget(MLIRContext *context);

/// Returns the GPU target attribute from attribute `attr`.
IREE::GPU::TargetAttr getGPUTargetAttr(DictionaryAttr attr);
/// Returns the GPU target attribute from the executable target wrapping |op|
/// if found. Returns null TargetAttr otherwise.
IREE::GPU::TargetAttr getGPUTargetAttr(MLIRContext *context,
                                       IREE::HAL::ExecutableTargetAttr attr);
IREE::GPU::TargetAttr getGPUTargetAttr(Operation *op);

/// Check if the target architecture supports global load DMA.
/// Returns true only for CDNA4+ (gfx950+) architectures.
bool targetSupportsGlobalLoadDMA(IREE::GPU::TargetAttr target);

/// True if `gpu.shuffle` on `bitwidth`-bit elements is supported by `target`.
/// Widths up to `kShuffleNativeBits` are universally supported; wider widths
/// require a backend that decomposes them (today only AMDGPU). A null
/// `target` is treated conservatively (reject wide shuffles).
bool targetSupportsShuffleBitwidth(IREE::GPU::TargetAttr target,
                                   unsigned bitwidth);

// Methods to retrieve information association with `configuration` field
// of `hal.executable.target` attribute used commonly in GPU codegen pipelines.
std::optional<int64_t> getConfigWavesPerEu(DictionaryAttr targetAttr);
IntegerAttr getConfigWavesPerEuAttr(DictionaryAttr targetAttr);

/// Methods to add attributes to the `config` list.
void addConfigGPUTarget(MLIRContext *context, IREE::GPU::TargetAttr,
                        SmallVectorImpl<NamedAttribute> &config);
void addConfigWavesPerEu(MLIRContext *context, int64_t wavesPerEu,
                         SmallVectorImpl<NamedAttribute> &config);

/// Returns the GPU subgroup size chosen for the current CodeGen pipeline if
/// exists; otherwise returns the subgroup size from the GPU target description.
/// Returns std::nullopt if none found.
std::optional<int> getGPUSubgroupSize(mlir::FunctionOpInterface func);

/// Returns the top-level `IREE::HAL::ExecutableVariantOp` operations from the
/// given `moduleOp`, ensuring they are returned in their original IR order.
SmallVector<IREE::HAL::ExecutableVariantOp>
getExecutableVariantOps(mlir::ModuleOp moduleOp);

} // namespace mlir::iree_compiler

#endif // IREE_COMPILER_CODEGEN_UTILS_GPUUTILS_H_
