//===- GPUDialect.h - MLIR Dialect for GPU Kernels --------------*- C++ -*-===//
//
// Part of the LLVM Project, 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 defines the GPU kernel-related operations and puts them in the
// corresponding dialect.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_DIALECT_GPU_GPUDIALECT_H
#define MLIR_DIALECT_GPU_GPUDIALECT_H

#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
#include "mlir/Dialect/DLTI/Traits.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/FunctionSupport.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"

namespace mlir {
class FuncOp;

namespace gpu {

/// Utility class for the GPU dialect to represent triples of `Value`s
/// accessible through `.x`, `.y`, and `.z` similarly to CUDA notation.
struct KernelDim3 {
  Value x;
  Value y;
  Value z;
};

class AsyncTokenType
    : public Type::TypeBase<AsyncTokenType, Type, TypeStorage> {
public:
  // Used for generic hooks in TypeBase.
  using Base::Base;
};

/// MMAMatrixType storage and uniquing. Array is uniqued based on its shape
/// and type.
struct MMAMatrixStorageType : public TypeStorage {
  MMAMatrixStorageType(unsigned numDims, const int64_t *dimShapes,
                       Type elementType, StringRef operand)
      : dimShapes(dimShapes), numDims(numDims), elementType(elementType),
        operand(operand) {}

  /// The hash key for uniquing.
  using KeyTy = std::tuple<ArrayRef<int64_t>, Type, StringRef>;
  bool operator==(const KeyTy &key) const {
    return key == KeyTy(getShape(), elementType, operand);
  }

  /// Construction.
  static MMAMatrixStorageType *construct(TypeStorageAllocator &allocator,
                                         const KeyTy &key) {
    ArrayRef<int64_t> shape = allocator.copyInto(std::get<0>(key));
    StringRef operand = allocator.copyInto(std::get<2>(key));

    return new (allocator.allocate<MMAMatrixStorageType>())
        MMAMatrixStorageType(shape.size(), shape.data(), std::get<1>(key),
                             operand);
  }

  ArrayRef<int64_t> getShape() const {
    return ArrayRef<int64_t>(dimShapes, numDims);
  }

  StringRef getOperand() const { return operand; }

  /// Reference to the shape of the MMA matrix.
  const int64_t *dimShapes;

  /// Number of dimensions in the MMA matrix.
  unsigned numDims;

  /// Element type of elements held in the MMA matrix.
  Type elementType;

  /// MMA operand that this MMAMatrix holds. The general form of operation this
  /// type supports is given by the equation C += A*B. This field specifies
  /// which operand in the given equation is held by this type. The valid values
  /// are "AOp", "BOp" and "COp".
  StringRef operand;
};

/// MMAMatrix represents a matrix held by a subgroup for matrix-matrix multiply
/// accumulate operations. MMAMatrices are taken as direct operands by these
/// operations and are also produced as results. These matrices are meant to
/// reside in the registers. A limited number of pointwise operations can be
/// performed on these matrices, i.e., operations which operate uniformly on
/// all the elements in the matrix and do not change the order of matrix
/// elements. The above conditions exist because the layout of matrix elements
/// inside the matrix is opaque i.e., the elements may be present in the
/// matrix in any order. The general usage of this type is shown as follows:-
///
///   %0 = gpu.subgroup_mma_load_matrix %arg0[%c0, %c0] {leadDimension = 16 :
///           index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "AOp">
///
/// The MMAMatrixType describes the shape of the matrix being loaded and the
/// operand being loaded too. The operand needs to be specified to aid the
/// lowering of this type to dialects such as NVVM where each workitem may
/// hold different amount of elements depending on the elementType of the
/// matrix. For e.g., Each workitem holds 4 vector<2xf16>s for f16 data type
/// and 8 f32s for f32 data type of MMAMatrix. Some other instances of usage
/// are:-
///
///   %3 = gpu.subgroup_mma_compute %0, %1, %2 :
///   !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp">
///    -> !gpu.mma_matrix<16x16xf32, "COp">
///
///
///   gpu.subgroup_mma_store_matrix %3, %arg22[%c0, %c0] {leadDimension = 16
///           : index}: !gpu.mma_matrix<16x16xf32, "COp">, memref<16x16xf32>
// TODO: consider moving this to ODS.
class MMAMatrixType
    : public Type::TypeBase<MMAMatrixType, Type, MMAMatrixStorageType> {
public:
  using Base::Base;

  /// Get MMAMatrixType and verify construction Invariants.
  static MMAMatrixType get(ArrayRef<int64_t> shape, Type elementType,
                           StringRef operand);

  /// Get MMAMatrixType at a particular location and verify construction
  /// Invariants.
  static MMAMatrixType getChecked(function_ref<InFlightDiagnostic()> emitError,
                                  ArrayRef<int64_t> shape, Type elementType,
                                  StringRef operand);

  /// Check if a type is valid a MMAMatrixType elementType.
  static bool isValidElementType(Type elementType);

  /// Verify that shape and elementType are actually allowed for the
  /// MMAMatrixType.
  static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
                              ArrayRef<int64_t> shape, Type elementType,
                              StringRef operand);

  /// Get number of dims.
  unsigned getNumDims() const;

  /// Get shape of the matrix.
  ArrayRef<int64_t> getShape() const;

  /// Get elementType of a single element.
  Type getElementType() const;

  /// The general form of operation this type supports is given by the equation
  /// C += A*B. This function returns which operand in the given equation is
  /// held by this type. String returned can be one of"AOp", "BOp" and "COp".
  StringRef getOperand() const;
};

// Adds a `gpu.async.token` to the front of the argument list.
void addAsyncDependency(Operation *op, Value token);

} // end namespace gpu
} // end namespace mlir

#include "mlir/Dialect/GPU/GPUOpsEnums.h.inc"

#include "mlir/Dialect/GPU/GPUOpsDialect.h.inc"

#include "mlir/Dialect/GPU/GPUOpInterfaces.h.inc"

#define GET_OP_CLASSES
#include "mlir/Dialect/GPU/GPUOps.h.inc"

#endif // MLIR_DIALECT_GPU_GPUDIALECT_H