// Conversions from TritonGPU layouts (e.g. BlockedEncodingAttr) to
// LinearLayout.

#ifndef TRITON_DIALECT_TRITONGPU_IR_LINEARLAYOUTCONVERSIONS_H
#define TRITON_DIALECT_TRITONGPU_IR_LINEARLAYOUTCONVERSIONS_H

#include <optional>

#include "triton/Tools/LinearLayout.h"

namespace mlir::triton {
enum class ScaleDotElemType : uint32_t;
} // namespace mlir::triton

namespace mlir::triton::gpu {
class SwizzledSharedEncodingAttr;
class NVMMASharedEncodingAttr;
class TensorOrMemDesc;
class MemDescType;
class CGAEncodingAttr;
enum class TMAMode;

// - BlockedEncodingAttrs have the following input dimensions.
//
//   "register": elements in one thread
//   "lane": threads in a warp
//   "warp": warps in a block/CTA
//   "block": blocks in a cluster
//
// - An n-dimensional SwizzledSharedEncodingAttr has the following input
// dimensions.
//
//   "offset": the n'th element in the allocation, within a particular thread
//      block (i.e. within a CTA).  The offset is measured in elements, not
//      bytes.
//   "block": blocks in a cluster
//
// All layouts have the following output dimensions.
//
//  "dimi" for i in 0..n-1: the location in the n'th logical dimension of the
//  output tensor.  These also are not reordered according to the layout's
//  `order`.
//
// You can flatten the input or output dimensions into a single dimension using
// LinearLayout::flattenIns/Outs().
//
// elemBitWidth is the bit width of one element in the layout.  This is required
// to compute the linear layout for MMAv3 (i.e. Hopper) shared layouts (i.e.
// shared layouts with nvmma_shared layout) but is otherwise unused.
LinearLayout toLinearLayout(RankedTensorType type);
LinearLayout toLinearLayout(MemDescType type);
LinearLayout toLinearLayout(TensorOrMemDesc type);
// UNSAFE OVERLOAD!
// If you call this with a SharedMemoryEncodingAttr, you should call it
// with the allocShape as the shape, otherwise the layout will be incorrect!
LinearLayout toLinearLayout(ArrayRef<int64_t> shape, Attribute layout);

// Returns the linear component of a padded shared encoding. The encoding must
// satisfy isPaddedEncoding (asserts otherwise).
//
// Unlike toLinearLayout, this makes explicit that the resulting linear layout
// is incomplete — the padding information is not captured in the linear layout.
LinearLayout paddedLinearLayout(MemDescType type);
LinearLayout paddedLinearLayout(ArrayRef<int64_t> shape, Attribute encoding);

// Convert to a linear layout, returning only the linear component of padded
// encodings because padding cannot be represented by LinearLayout.
LinearLayout toLinearLayoutIgnoringPadding(MemDescType type);
LinearLayout toLinearLayoutIgnoringPadding(ArrayRef<int64_t> shape,
                                           Attribute encoding);

// Convert the shared encoding of a tensor with `nvmma_shared` layout to a
// LinearLayout that maps from a linear shared memory offset to tensor index.
//
// If `disableSwizzle` is set, then the resulting layout does not include
// swizzling.
LinearLayout nvmmaSharedToLinearLayout(ArrayRef<int64_t> shape,
                                       NVMMASharedEncodingAttr shared,
                                       TMAMode mode,
                                       bool disableSwizzle = false);
FailureOr<LinearLayout>
nvmmaSharedToLinearLayout(ArrayRef<int64_t> shape,
                          NVMMASharedEncodingAttr shared, TMAMode mode,
                          bool disableSwizzle, bool emitErrors);

// Given a linear layout where the input dimensions contain a "block" dimension,
// this method sets the "block" dimension to 0 and removes the corresponding
// output dimensions.
//
// Note that this behavior differs from calling
// `LinearLayout::sublayout(inDimNames, outDimNames)` when "block" is not in
// `inDimNames`. The latter does not modify the output sizes.
LinearLayout getLayoutWithinBlock(const LinearLayout &layout);

// Combines the layout of a CTA (input dims [register, lane, warp]) with the
// layout of a CGA (i.e. a block), and ensures that the resulting layout has the
// given shape.
//
// See the nomenclature note at the top of LinearLayoutConversions.cpp for why
// the variable with type CGAEncodingAttr is called cgaLayoutAttr.
LinearLayout combineCtaCgaWithShape(LinearLayout ctaLayout,
                                    CGAEncodingAttr cgaLayoutAttr,
                                    ArrayRef<int64_t> shape);

LinearLayout chooseWmmaCTALinearLayout(MLIRContext *ctx, unsigned rank,
                                       ArrayRef<unsigned> warpsPerCTA,
                                       ArrayRef<unsigned> tilesPerWarp);
// In this function, we construct a linear layout representing the
// <shared memory offset, iteration, block> -> <tensor element index> mapping
// for entire `src` and `dst` tensors.  We determine the shape of the
// intermediate shared memory buffer needed for a register-to-register
// conversion using the maximum size accessed in each dimension from `src`'s
// layout and `dst`'s layout.  See the getRepShapeForCvt function in
// Allocation.cpp for details. Note that the buffer might be smaller than the
// tensor being converted, so we need multiple "iterations" to move a subregion
// of the `src` tensor to the corresponding subregion of the `dst` tensor.  The
// pesudo code of layout conversion is as follows:
//
// for iter in 0..numIterations:
//   sync threads
//   for vecIdx in [0..numRegisters/storeVec]:
//     registers <- get registers used in iter
//     offsets <- get offsets using the intermediate linear layout
//     store registers[vecIdx * storeVec, (vecIdx + 1) * storeVec)] to shared
//     memory
//   sync threads
//   for vecIdx in [0..numRegisters/loadVec]:
//     registers <- get registers used in iter
//     offsets <- get offsets using the intermediate linear layout
//     load registers[vecIdx * loadVec, (vecIdx + 1) * loadVec)] from shared
//     memory
LinearLayout chooseShemLayoutForRegToRegConversion(
    MLIRContext *ctx, ArrayRef<unsigned> tensorShape,
    ArrayRef<unsigned> repShape, ArrayRef<unsigned> order);

// Create LinearLayout for scale in scaled mfma.
LinearLayout chooseScaledMfmaScaleLayout(MLIRContext *ctx, int dotOperandIdx,
                                         ArrayRef<int64_t> dotOperandShape,
                                         unsigned mfmaMDim,
                                         ArrayRef<unsigned> tilesPerWarp,
                                         ArrayRef<unsigned> warpsPerCTA);

LinearLayout chooseScaledWmmaScaleLayout(
    MLIRContext *ctx, int dotOperandIdx, ArrayRef<int64_t> dotOperandShape,
    unsigned wmmaMDim, unsigned wmmaNDim, bool isTransposed,
    unsigned scaleFactor, LinearLayout ctaLayout, CGAEncodingAttr cgaLayout);

LinearLayout getSM120DotScaledScaleLayout(MLIRContext *ctx,
                                          ArrayRef<int64_t> shape, int opIdx,
                                          ArrayRef<unsigned> warpsPerCTA,
                                          CGAEncodingAttr cgaLayout);

// Create LinearLayout for nvidia mma tile.
LinearLayout nvidiaMmaTile(MLIRContext *ctx, ArrayRef<unsigned> tileShape,
                           unsigned kWidth, ArrayRef<unsigned> order,
                           ArrayRef<unsigned> repOrder);

// Create a LinearLayout similar to mfmaLayout, but changing each thread to hold
// 8 elements. This layout is useful for emitting the widest 128-bit global
// store instructions. Since it closely resembles mfmaLayout, conversion between
// the two can be done using transferWithinWarp, without involving LDS
std::optional<LinearLayout> chooseMfmaLikeStoreLayout(RankedTensorType valType);

// Create the core layout (atom in the PTX manual) a given nvmma shared encoding
LinearLayout getCoreMatrixLinearLayout(NVMMASharedEncodingAttr shared,
                                       bool disableSwizzle);

// Create the canonical 2D SMEM linear layout for scales.
// This is the shared-memory layout for the 32x4x4 TMEM scale-tile convention
// recognized between the user kernel and the compiler.
LinearLayout getScaleSmemLayoutForTMEMCopy(MLIRContext *ctx,
                                           ArrayRef<int64_t> shape,
                                           CGAEncodingAttr cgaLayout);

// Create a TDM (Tensor DMA) LinearLayout: (message, warp, block) ->
// (dim0, dim1, ...).  TDM is warp-granular.  The "warp" sublayout is an
// identity over `warpsPerCTA`, zero-padded up to log2(numWarps) so the
// full module warpId is covered; padded rows expose the redundant bits
// as free variables (via getFreeVariableMasks("warp")) so partial copies
// (K < numWarps) can pred-off inactive warps.  "message" covers the
// per-warp tile (surjectivity); "block" comes from `cgaLayout`.
//
// `warpUsedHint`: power-of-two-popcount bitmask whose K = popcount(hint)
// set bits select active warps.  The varying warpId bit positions in the
// active set are the warp bits that contribute to per-warp offsets; all
// other warpId bits become free variables for predicating inactive warps.
// Empty = no-hint default (lowest log2(K) bits, K = prod(warpsPerCTA)).
LinearLayout getTDMLinearLayout(ArrayRef<int64_t> blockShape,
                                ArrayRef<unsigned> warpsPerCTA,
                                const LinearLayout &cgaLayout, int totalWarps,
                                std::optional<uint32_t> warpUsedHint = {});

} // namespace mlir::triton::gpu

namespace mlir {

// Conversion from `srcLayout` to `dstLayout` involving the minimum amount of
// data transfer. The output will be such that layout.getInDimNames() ==
// layout.getOutDimNames() and the conversion will not include block (resp.
// warp or lane) if it can be avoided.
triton::LinearLayout minimalCvtLayout(const triton::LinearLayout &srcLayout,
                                      const triton::LinearLayout &dstLayout);

// Type-based convenience overload for layouts that can be converted to a
// linear layout.
triton::LinearLayout minimalCvtLayout(Type srcTy, Type dstTy);

} // namespace mlir

#endif // TRITON_DIALECT_TRITONGPU_IR_LINEARLAYOUTCONVERSIONS_H
