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

#include <cstdint>
#include "iree/compiler/Codegen/Common/GPU/GPUPatterns.h"
#include "iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.h"
#include "iree/compiler/Codegen/Common/VectorLayoutAnalysis.h"
#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
#include "iree/compiler/Codegen/Utils/GPUUtils.h"
#include "iree/compiler/Utils/Permutation.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/GPU/Utils/GPUUtils.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/Verifier.h"
#include "mlir/Rewrite/PatternApplicator.h"

namespace mlir::iree_compiler {

using namespace mlir::iree_compiler::IREE::VectorExt;
using VectorValue = TypedValue<VectorType>;

/// Helper to linearize the given |ids| with maximum values given as |sizes|.
/// Gets the element ID in terms of |elementCount| and adds the element
/// |offset|. For example,
///
/// IDs = [d0, d1, d2, d3]
/// sizes = [s0, s1, s2, s3]
/// linear_index = d0 * (s1 * s2 * s3)
///              + d1 * (s2 * s3)
///              + d2 * (s3)
///              + d3
/// return element_index = linear_index * |elementCount| + |offset|;
static Value linearizeIndex(OpBuilder &builder, Value offset,
                            ArrayRef<OpFoldResult> ids, ArrayRef<int64_t> sizes,
                            int64_t elementCount) {
  SmallVector<AffineExpr> exprs(ids.size() + 1);
  bindSymbolsList(builder.getContext(), MutableArrayRef{exprs});
  AffineExpr idExpr = builder.getAffineConstantExpr(0);

  for (int i = 0, e = ids.size(); i < e; ++i) {
    if (sizes[i] > 1) {
      // Multiply by the residual threads along this dimension (which must be
      // faster changing than all previous dimensions) and add the id for this
      // dimension.
      idExpr = idExpr * builder.getAffineConstantExpr(sizes[i]) + exprs[i];
    }
  }
  idExpr = idExpr * builder.getAffineConstantExpr(elementCount);
  idExpr = idExpr + exprs.back();
  SmallVector<OpFoldResult> mapArgs(ids);
  mapArgs.push_back(offset);
  return affine::makeComposedAffineApply(
             builder, offset.getLoc(),
             AffineMap::get(0, mapArgs.size(), idExpr), mapArgs)
      .getResult();
}

/// Given a set of base transfer |indices|, |offsets| for the batch/outer
/// dimensions, and distributed warp and thread indices, computes the indices
/// of the distributed transfer operation based on the |vectorLayout|.
static SmallVector<Value> getTransferIndicesFromNestedLayout(
    OpBuilder &b, ValueRange indices, ArrayRef<int64_t> offsets,
    NestedLayoutAttr vectorLayout, AffineMap permutationMap,
    ArrayRef<Value> warpIndices, ArrayRef<Value> threadIndices) {
  auto isBroadcast = [](AffineExpr expr) {
    if (auto constExpr = dyn_cast<AffineConstantExpr>(expr))
      return constExpr.getValue() == 0;
    return false;
  };
  int64_t rank = vectorLayout.getRank();
  // Permute the batch and outer vector offsets to match the order of
  // the vector dimensions using the inverse of the batch/offset order.
  ArrayRef<int64_t> batchOffsets(offsets.begin(), rank);
  ArrayRef<int64_t> outerVectorOffsets(offsets.begin() + rank, rank);

  SmallVector<Value> slicedIndices(indices.begin(), indices.end());
  for (const auto &[i, dim] : llvm::enumerate(permutationMap.getResults())) {
    // Broadcasted dimension offsets can be used as-is; the read index is
    // invariant of the thread in such cases (and illegal for writes).
    if (isBroadcast(dim)) {
      continue;
    }
    unsigned pos = cast<AffineDimExpr>(dim).getPosition();
    SmallVector<OpFoldResult> ids = {
        warpIndices[i], b.getIndexAttr(batchOffsets[i]),
        b.getIndexAttr(outerVectorOffsets[i]), threadIndices[i]};
    // The order in which a vector dimension is "tiled" is
    // subgroups -> batches -> outer vectors -> threads -> elements
    SmallVector<int64_t> sizes = {
        vectorLayout.getSubgroupTile()[i], vectorLayout.getBatchTile()[i],
        vectorLayout.getOuterTile()[i], vectorLayout.getThreadTile()[i]};
    slicedIndices[pos] = linearizeIndex(b, indices[pos], ids, sizes,
                                        vectorLayout.getElementTile()[i]);
  }
  return slicedIndices;
}

static SmallVector<int64_t>
getElementVectorTileShape(NestedLayoutAttr vectorLayout) {
  int64_t rank = vectorLayout.getRank();
  SmallVector<int64_t> tileShape = vectorLayout.getDistributedShape();
  // We tile to a vector with BATCH, OUTER, and ELEMENT dimensions. So to access
  // the subvector only containing elements, we need indices in all BATCH and
  // OUTER (rank * 2) dimensions to have tile size 1.
  for (int i = 0, e = rank * 2; i < e; ++i) {
    tileShape[i] = 1;
  }
  return tileShape;
}

/// Computes the warp and thread indices for the given vector layout from a
/// single linearized thread ID.
static void populateWarpAndThreadIndices(RewriterBase &rewriter, Value threadId,
                                         int64_t subgroupSize,
                                         NestedLayoutAttr vectorLayout,
                                         SmallVector<Value> &warpIndices,
                                         SmallVector<Value> &threadIndices) {
  // The delinearized thread IDs are returned from outer most to inner most,
  // i.e. before applying the layout described dimensions ordering.
  int64_t rank = vectorLayout.getRank();
  SmallVector<Value> threadIds =
      vectorLayout.computeThreadIds(threadId, subgroupSize, rewriter);
  warpIndices = SmallVector<Value>(threadIds.begin(), threadIds.begin() + rank);
  threadIndices = SmallVector<Value>(threadIds.begin() + rank,
                                     threadIds.begin() + 2 * rank);
}

namespace {

/// Pattern to distribute `vector.transfer_read` ops with nested layouts.
struct DistributeTransferRead final
    : OpDistributionPattern<vector::TransferReadOp> {
  using OpDistributionPattern::OpDistributionPattern;

  DistributeTransferRead(MLIRContext *context, Value threadId,
                         int64_t subgroupSize)
      : OpDistributionPattern(context), threadId(threadId),
        subgroupSize(subgroupSize) {}

  LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
                                DistributionSignature &signature,
                                PatternRewriter &rewriter) const override {
    // TODO: Support masking.
    if (readOp.getMask()) {
      return rewriter.notifyMatchFailure(readOp, "unimplemented: masked read");
    }
    NestedLayoutAttr vectorLayout =
        dyn_cast<NestedLayoutAttr>(signature[readOp.getResult()]);
    if (!vectorLayout) {
      return rewriter.notifyMatchFailure(readOp,
                                         "non-nested transfer_read layout");
    }

    // Guard on memrefs for distribution. In isolation this pattern is agnostic
    // to tensors or memrefs.
    if (!isa<MemRefType>(readOp.getSource().getType())) {
      return rewriter.notifyMatchFailure(readOp,
                                         "distribution expects memrefs");
    }

    SmallVector<int64_t> distShape = vectorLayout.getDistributedShape();
    SmallVector<int64_t> tileShape = getElementVectorTileShape(vectorLayout);
    int64_t rank = vectorLayout.getRank();

    Type elementType = readOp.getSource().getType().getElementType();
    auto vectorType = VectorType::get(distShape, elementType);
    // The shape of the vector we read is pre-permutation. The permutation is
    // a transpose on the resulting read vector.
    auto innerVectorType =
        VectorType::get(vectorLayout.getElementTile(), elementType);

    // Initialize the full distributed vector for unrolling the batch/outer
    // vector dimensions.
    Value zero = rewriter.create<arith::ConstantOp>(
        readOp.getLoc(), vectorType, rewriter.getZeroAttr(vectorType));
    VectorValue acc = cast<VectorValue>(zero);

    SmallVector<Value> warpIndices, threadIndices;
    populateWarpAndThreadIndices(rewriter, threadId, subgroupSize, vectorLayout,
                                 warpIndices, threadIndices);

    ValueRange indices = readOp.getIndices();
    SmallVector<int64_t> strides(rank, 1);
    for (SmallVector<int64_t> offsets :
         StaticTileOffsetRange(distShape, tileShape)) {
      SmallVector<Value> slicedIndices = getTransferIndicesFromNestedLayout(
          rewriter, indices, offsets, vectorLayout, readOp.getPermutationMap(),
          warpIndices, threadIndices);

      VectorValue slicedRead = rewriter.create<vector::TransferReadOp>(
          readOp.getLoc(), innerVectorType, readOp.getSource(), slicedIndices,
          readOp.getPermutationMapAttr(), readOp.getPadding(), readOp.getMask(),
          readOp.getInBoundsAttr());

      if (acc.getType().getRank() == 0) {
        // TODO: This should really be a folding pattern in
        // insert_strided_slice, but instead insert_strided_slice just doesn't
        // support 0-d vectors...
        acc = slicedRead;
      } else {
        acc = rewriter.create<vector::InsertStridedSliceOp>(
            readOp.getLoc(), slicedRead, acc, offsets, strides);
      }
    }

    replaceOpWithDistributedValues(rewriter, readOp, acc);
    return success();
  }

  Value threadId;
  int64_t subgroupSize;
};

/// Pattern to distribute `vector.transfer_write` ops with nested layouts.
struct DistributeTransferWrite final
    : OpDistributionPattern<vector::TransferWriteOp> {
  using OpDistributionPattern::OpDistributionPattern;

  DistributeTransferWrite(MLIRContext *context, Value threadId,
                          int64_t subgroupSize)
      : OpDistributionPattern(context), threadId(threadId),
        subgroupSize(subgroupSize) {}

  LogicalResult matchAndRewrite(vector::TransferWriteOp writeOp,
                                DistributionSignature &signature,
                                PatternRewriter &rewriter) const override {
    // TODO: Support masking.
    if (writeOp.getMask()) {
      return rewriter.notifyMatchFailure(writeOp,
                                         "unimplemented: masked write");
    }
    NestedLayoutAttr vectorLayout =
        dyn_cast<NestedLayoutAttr>(signature[writeOp.getVector()]);
    if (!vectorLayout) {
      return rewriter.notifyMatchFailure(writeOp,
                                         "non-nested transfer_write layout");
    }

    if (!isa<MemRefType>(writeOp.getSource().getType())) {
      return rewriter.notifyMatchFailure(writeOp,
                                         "distribution expects memrefs");
    }

    SmallVector<int64_t> distShape = vectorLayout.getDistributedShape();
    SmallVector<int64_t> tileShape = getElementVectorTileShape(vectorLayout);
    int64_t rank = vectorLayout.getRank();

    SmallVector<Value> warpIndices, threadIndices;
    populateWarpAndThreadIndices(rewriter, threadId, subgroupSize, vectorLayout,
                                 warpIndices, threadIndices);

    Value distributedVector =
        getDistributed(rewriter, writeOp.getVector(), vectorLayout);

    ValueRange indices = writeOp.getIndices();
    for (SmallVector<int64_t> offsets :
         StaticTileOffsetRange(distShape, tileShape)) {
      SmallVector<Value> slicedIndices = getTransferIndicesFromNestedLayout(
          rewriter, indices, offsets, vectorLayout, writeOp.getPermutationMap(),
          warpIndices, threadIndices);

      // Extract the "element vector" from the inner most dimensions. All outer
      // dimensions are either unrolled or distributed such that this is a
      // contiguous slice.
      ArrayRef<int64_t> offsetArray(offsets);
      Value slicedVector = rewriter.create<vector::ExtractOp>(
          writeOp.getLoc(), distributedVector,
          offsetArray.take_front(rank * 2));
      // Promote the slicedVector to 0-d vector if it is a scalar.
      if (!isa<VectorType>(slicedVector.getType())) {
        auto promotedType =
            VectorType::get({}, getElementTypeOrSelf(slicedVector));
        slicedVector = rewriter.create<vector::BroadcastOp>(
            writeOp.getLoc(), promotedType, slicedVector);
      }
      rewriter.create<vector::TransferWriteOp>(
          writeOp.getLoc(), slicedVector, writeOp.getSource(), slicedIndices,
          writeOp.getPermutationMapAttr(), writeOp.getMask(),
          writeOp.getInBoundsAttr());
    }

    rewriter.eraseOp(writeOp);
    return success();
  }

  Value threadId;
  int64_t subgroupSize;
};

struct DistributeBroadcast final : OpDistributionPattern<vector::BroadcastOp> {
  using OpDistributionPattern::OpDistributionPattern;

  LogicalResult matchAndRewrite(vector::BroadcastOp broadcastOp,
                                DistributionSignature &signature,
                                PatternRewriter &rewriter) const override {
    Location loc = broadcastOp.getLoc();
    VectorValue dstVector = broadcastOp.getVector();
    auto vectorLayout = dyn_cast<NestedLayoutAttr>(signature[dstVector]);
    if (!vectorLayout) {
      return rewriter.notifyMatchFailure(broadcastOp,
                                         "non-nested result vector layout");
    }
    SmallVector<int64_t> distShape = vectorLayout.getDistributedShape();
    Type elementType =
        llvm::cast<ShapedType>(dstVector.getType()).getElementType();
    auto vectorType = VectorType::get(distShape, elementType);

    VectorValue srcVector = dyn_cast<VectorValue>(broadcastOp.getSource());
    // If the srcVector is a scalar (like f32) we proceed with the scalar
    // distribution branch.
    if (!srcVector) {
      // The way distribution currently works, there is no partial thread
      // distribution, so a scalar is available to all threads. Scalar
      // distribution is simply a broadcast from scalar to the distributed
      // result shape.
      Value source = broadcastOp.getSource();
      VectorValue accumulator =
          rewriter.create<vector::BroadcastOp>(loc, vectorType, source);
      replaceOpWithDistributedValues(rewriter, broadcastOp, accumulator);
      return success();
    }

    auto sourceLayout = dyn_cast<NestedLayoutAttr>(signature[srcVector]);
    if (!sourceLayout) {
      return rewriter.notifyMatchFailure(broadcastOp,
                                         "non-nested source vector layout");
    }

    Value accumulator = rewriter.create<arith::ConstantOp>(
        loc, vectorType, rewriter.getZeroAttr(vectorType));

    int64_t rank = vectorLayout.getRank();
    // We unroll along both the batch and outer dimensions for a similar reason
    // to the transfer ops. `vector.broadcast` can only broadcast along outer
    // dims, so mixing broadcasted and un-broadcasted element/outer dims can't
    // be represented with a single `vector.broadcast`.
    SmallVector<int64_t> resultVectorUnrollShape =
        getElementVectorTileShape(vectorLayout);

    Value distributedSource = getDistributed(rewriter, srcVector, sourceLayout);

    VectorType broadcastTargetType =
        VectorType::get(vectorLayout.getElementTile(), elementType);

    int64_t sourceRank = sourceLayout.getRank();

    for (SmallVector<int64_t> offsets :
         StaticTileOffsetRange(distShape, resultVectorUnrollShape)) {
      ArrayRef<int64_t> offsetsRef(offsets);

      // Slice out the last |sourceRank| dimensions which is the inner
      // broadcasted shape.
      ArrayRef<int64_t> batchSourceOffsets =
          offsetsRef.slice(rank - sourceRank, sourceRank);
      ArrayRef<int64_t> outerSourceOffsets =
          offsetsRef.slice(2 * rank - sourceRank, sourceRank);

      // Construct the list of source offsets based on the batch/outer order of
      // the broadcasted vector. This is because we need to compute the offsets
      // into the distributed source vector with the distributed permutation.
      SmallVector<int64_t> sourceOffsets;
      sourceOffsets.append(batchSourceOffsets.begin(),
                           batchSourceOffsets.end());
      sourceOffsets.append(outerSourceOffsets.begin(),
                           outerSourceOffsets.end());

      // Extract a slice of the input to be broadcasted.
      Value slice = rewriter.create<vector::ExtractOp>(loc, distributedSource,
                                                       sourceOffsets);
      // TODO: Support non-trivial element orders.
      if (vector::isBroadcastableTo(slice.getType(), broadcastTargetType) !=
          vector::BroadcastableToResult::Success) {
        return rewriter.notifyMatchFailure(
            broadcastOp,
            "unimplemented: non-trivial broadcast source element order");
      }
      Value broadcastedSlice =
          rewriter.create<vector::BroadcastOp>(loc, broadcastTargetType, slice);
      // Insert into the broadcasted destination vector.
      accumulator = rewriter.create<vector::InsertOp>(
          loc, broadcastedSlice, accumulator, offsetsRef.take_front(rank * 2));
    }

    replaceOpWithDistributedValues(rewriter, broadcastOp, accumulator);
    return success();
  }
};

static int64_t getShuffleOffset(NestedLayoutAttr layout, int64_t dim) {
  return layout.getThreadStrides()[dim];
}

static int64_t getShuffleWidth(NestedLayoutAttr layout, int64_t dim) {
  return layout.getThreadTile()[dim];
}

/// The lowering for multi_reduction is done in two steps:
///   1. Local Reduce: Each thread reduces all elements carried by it along
///      the reduction dimensions. This is the batch, outer and element dims.
///   2. Thread Reduce: Each thread reduces result of step 1 across threads
///      by doing a butterfly shuffle.
///   3. Accumulator Reduce: Each thread reduces it's intermediate reduced
///      results with the accumulator it holds.
/// Currently, reduction across warps is not supported, but it would just add
/// another step, Warp Reduce, where threads do an atomic addition on a buffer.
struct DistributeMultiReduction final
    : OpDistributionPattern<vector::MultiDimReductionOp> {
  using OpDistributionPattern::OpDistributionPattern;

  DistributeMultiReduction(MLIRContext *context, int64_t subgroupSize,
                           int64_t maxBitsPerShuffle, int64_t benefit = 1)
      : OpDistributionPattern(context, benefit), subgroupSize(subgroupSize),
        maxBitsPerShuffle(maxBitsPerShuffle) {}

  LogicalResult matchAndRewrite(vector::MultiDimReductionOp multiReduceOp,
                                DistributionSignature &signature,
                                PatternRewriter &rewriter) const override {
    VectorValue srcVector = multiReduceOp.getSource();
    Value acc = multiReduceOp.getAcc();
    Value res = multiReduceOp.getResult();
    auto accVector = dyn_cast<VectorValue>(acc);
    auto resVector = dyn_cast<VectorValue>(res);

    auto srcLayout = dyn_cast_or_null<NestedLayoutAttr>(signature[srcVector]);
    if (!srcLayout) {
      return rewriter.notifyMatchFailure(multiReduceOp,
                                         "expected nested layout attr");
    }

    Type elemTy = srcVector.getType().getElementType();
    unsigned elemBitwidth = elemTy.getIntOrFloatBitWidth();
    if (elemBitwidth != maxBitsPerShuffle) {
      return rewriter.notifyMatchFailure(
          multiReduceOp, llvm::formatv("unimplemented: packed shuffle",
                                       elemBitwidth, maxBitsPerShuffle));
    }

    VectorValue disSrc =
        getDistributed(rewriter, srcVector, signature[srcVector]);

    Value disAcc;
    if (accVector) {
      disAcc = getDistributed(rewriter, accVector, signature[accVector]);
    } else {
      // Scalars are always distributed to all threads already.
      disAcc = multiReduceOp.getAcc();
    }

    Location loc = multiReduceOp.getLoc();

    SmallVector<bool> reducedDims = multiReduceOp.getReductionMask();
    int64_t rank = srcVector.getType().getRank();

    // Do thread local reduce.

    // The distributed reduction mask is simply the same mask appended
    // thrice.
    SmallVector<bool> distributedReductionMask;
    distributedReductionMask.reserve(3 * rank);
    for (int i = 0; i < 3; ++i) {
      distributedReductionMask.append(reducedDims.begin(), reducedDims.end());
    }
    Value localInit = getCombiningIdentityValue(
        loc, rewriter, multiReduceOp.getKind(), disAcc.getType());
    auto localReduction = rewriter.create<vector::MultiDimReductionOp>(
        loc, disSrc, localInit, distributedReductionMask,
        multiReduceOp.getKind());

    VectorValue locallyReduced;
    if (accVector) {
      locallyReduced = dyn_cast<VectorValue>(localReduction.getResult());
    } else {
      // Broadcast scalar accumulator to vector.
      VectorType vecType = VectorType::get(ArrayRef{int64_t(1)}, elemTy);
      locallyReduced = rewriter.create<vector::BroadcastOp>(
          loc, vecType, localReduction.getResult());
    }

    assert(locallyReduced && "result should have been a vector");

    // Flatten the locally reduced value.
    VectorType shaped = locallyReduced.getType();
    int64_t numElements = shaped.getNumElements();
    SmallVector<int64_t> flatShape(1, numElements);
    VectorType flatVecType = VectorType::get(flatShape, elemTy);
    VectorValue flat =
        rewriter.create<vector::ShapeCastOp>(loc, flatVecType, locallyReduced);

    // Do inter-thread/warp reduce.
    FailureOr<VectorValue> threadReduced = doThreadReduction(
        rewriter, srcLayout, flat, multiReduceOp.getKind(), reducedDims);
    if (failed(threadReduced)) {
      return failure();
    }

    // Do reduction against accumulator, which needs to be done after thread
    // reduction.
    VectorValue unflattened = rewriter.create<vector::ShapeCastOp>(
        loc, shaped, threadReduced.value());

    if (!accVector) {
      // Broadcast the scalar (e.g., f32) to a vector type (e.g., vector<f32>)
      // because the following implementation requires the operand to be a
      // vector.
      disAcc = rewriter.create<vector::BroadcastOp>(loc, shaped, disAcc);
    }

    Value accReduction = vector::makeArithReduction(
        rewriter, loc, multiReduceOp.getKind(), unflattened, disAcc);
    auto accReduced = dyn_cast<VectorValue>(accReduction);
    if (!accReduced) {
      return failure();
    }

    if (resVector) {
      replaceOpWithDistributedValues(rewriter, multiReduceOp, accReduced);
    } else {
      Value accReducedVal = rewriter.create<vector::ExtractOp>(
          loc, accReduction, ArrayRef{int64_t(0)});
      replaceOpWithDistributedValues(rewriter, multiReduceOp, accReducedVal);
    }

    return success();
  }

  FailureOr<VectorValue> doThreadReduction(RewriterBase &rewriter,
                                           NestedLayoutAttr layout,
                                           VectorValue flat,
                                           vector::CombiningKind kind,
                                           ArrayRef<bool> reductionMask) const {
    VectorType flatVecType = flat.getType();
    int64_t numElements = flatVecType.getNumElements();
    Location loc = flat.getLoc();

    auto constOp = rewriter.create<arith::ConstantOp>(
        loc, rewriter.getZeroAttr(flatVecType));
    auto res = llvm::cast<VectorValue>(constOp.getResult());

    for (unsigned i = 0; i < numElements; ++i) {
      Value extracted = rewriter.create<vector::ExtractOp>(loc, flat, i);

      // Reduce across all reduction dimensions 1-by-1.
      for (unsigned i = 0, e = reductionMask.size(); i != e; ++i) {
        if (reductionMask[i]) {
          int64_t offset = getShuffleOffset(layout, i);
          int64_t width = getShuffleWidth(layout, i);
          assert(offset <= std::numeric_limits<uint32_t>::max() &&
                 width <= std::numeric_limits<uint32_t>::max());

          extracted = rewriter.create<gpu::SubgroupReduceOp>(
              loc, extracted, combiningKindToAllReduce(kind),
              /*uniform=*/false, /*cluster_size=*/width,
              /*cluster_stride=*/offset);
        }
      }

      res = rewriter.create<vector::InsertOp>(loc, extracted, res, i);
    }

    return res;
  }

  int64_t subgroupSize;
  int64_t maxBitsPerShuffle;
};

/// The lowering for Contract is performed in three steps (similar to above
/// multi_reduction):
///   1. Local Contract: Each thread performs operations on its locally
///   distributed elements.
///   2. Subgroup Reduction: Threads in each subgroup reduce the results from
///   step 1 across threads using a subgroup reduction if distribution occurs
///   along the reduction dimension.
///   3. Accumulator Reduction: Each thread combines its intermediate results
///   with its held accumulator.
///
/// Currently, reduction across multiple warps is not supported.
struct DistributeContract final : OpDistributionPattern<vector::ContractionOp> {
  using OpDistributionPattern::OpDistributionPattern;

  DistributeContract(MLIRContext *context, int64_t subgroupSize,
                     int64_t maxBitsPerShuffle, int64_t benefit = 1)
      : OpDistributionPattern(context, benefit), subgroupSize(subgroupSize),
        maxBitsPerShuffle(maxBitsPerShuffle) {}

  LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
                                DistributionSignature &signature,
                                PatternRewriter &rewriter) const override {
    FailureOr<VectorContractOpInfo> maybeOpInfo =
        VectorContractOpInfo::inferFromIndexingMaps(
            contractOp.getIndexingMapsArray());
    if (failed(maybeOpInfo)) {
      return rewriter.notifyMatchFailure(contractOp, "not a contraction");
    }
    // If mmaAttr exists, defer the lowering to use MMA.
    // Notify failure if the "iree.amdgpu.mma" intrinsic attribute is present.
    auto mmaAttr = contractOp->getAttrOfType<IREE::GPU::MmaInterfaceAttr>(
        "iree.amdgpu.mma");
    if (mmaAttr) {
      return rewriter.notifyMatchFailure(
          contractOp, "iree.amdgpu.mma intrinsic attribute exists");
    }

    auto lhsLayout = dyn_cast<NestedLayoutAttr>(signature[contractOp.getLhs()]);
    if (!lhsLayout) {
      return rewriter.notifyMatchFailure(
          contractOp, "missing nested layout for contraction lhs");
    }
    auto rhsLayout = dyn_cast<NestedLayoutAttr>(signature[contractOp.getRhs()]);
    if (!rhsLayout) {
      return rewriter.notifyMatchFailure(
          contractOp, "missing nested layout for contraction rhs");
    }

    Value disLhs = getDistributed(rewriter, contractOp.getLhs(), lhsLayout);
    Value disRhs = getDistributed(rewriter, contractOp.getRhs(), rhsLayout);

    Value acc = contractOp.getAcc();
    Value res = contractOp.getResult();
    auto accVector = dyn_cast<VectorValue>(acc);
    auto resVector = dyn_cast<VectorValue>(res);
    Value disAcc;
    if (accVector) {
      disAcc = getDistributed(rewriter, accVector, signature[accVector]);
    } else {
      disAcc = contractOp.getAcc();
    }

    Type accElemTy = getElementTypeOrSelf(acc.getType());

    MLIRContext *ctx = contractOp.getContext();
    Location loc = contractOp.getLoc();

    // Step 1: local contraction
    vector::ContractionOp localContractOp = doDistributedContraction(
        rewriter, loc, ctx, contractOp, disLhs, disRhs, disAcc);

    int64_t rank = lhsLayout.getRank();
    SmallVector<bool> reducedDims(rank, false);

    // Identify the reduction dimension and apply it for subgroup reduction.
    for (auto [index, iteratorType] :
         llvm::enumerate(contractOp.getIteratorTypes())) {
      if (vector::isReductionIterator(iteratorType)) {
        auto map = contractOp.getIndexingMapsArray()[0];
        int64_t redIdx = *(map.getResultPosition(getAffineDimExpr(index, ctx)));
        reducedDims[redIdx] = true;
      }
    }

    VectorValue localContractValue;
    if (accVector) {
      localContractValue = dyn_cast<VectorValue>(localContractOp.getResult());
    } else {
      VectorType vecType = VectorType::get(ArrayRef{int64_t(1)}, accElemTy);
      localContractValue = rewriter.create<vector::BroadcastOp>(
          loc, vecType, localContractOp.getResult());
    }

    assert(localContractValue && "result should have been a vector");

    // Flatten the locally result value.
    VectorType shaped = localContractValue.getType();
    int64_t numElements = shaped.getNumElements();
    SmallVector<int64_t> flatShape(1, numElements);
    VectorType flatVecType = VectorType::get(flatShape, accElemTy);
    VectorValue flat = rewriter.create<vector::ShapeCastOp>(loc, flatVecType,
                                                            localContractValue);

    // Step 2: Do subgroup reduction.
    FailureOr<VectorValue> threadReduced = doThreadReduction(
        rewriter, lhsLayout, flat, contractOp.getKind(), reducedDims);
    if (failed(threadReduced)) {
      return failure();
    }

    // Do reduction against accumulator, which needs to be done after thread
    // reduction.
    VectorValue unflattened = rewriter.create<vector::ShapeCastOp>(
        loc, shaped, threadReduced.value());

    if (!accVector) {
      disAcc = rewriter.create<vector::BroadcastOp>(loc, shaped, disAcc);
    }

    // Step 3: Accumulator Reduction
    Value accReduction = vector::makeArithReduction(
        rewriter, loc, contractOp.getKind(), unflattened, disAcc);
    auto accReduced = dyn_cast<VectorValue>(accReduction);
    if (!accReduced) {
      return failure();
    }

    if (resVector) {
      replaceOpWithDistributedValues(rewriter, contractOp, accReduced);
    } else {
      Value accReducedVal = rewriter.create<vector::ExtractOp>(
          loc, accReduction, SmallVector<int64_t>{0});
      replaceOpWithDistributedValues(rewriter, contractOp, accReducedVal);
    }

    return success();
  }

  vector::ContractionOp
  doDistributedContraction(RewriterBase &rewriter, Location loc,
                           MLIRContext *ctx, vector::ContractionOp contractOp,
                           Value lhs, Value rhs, Value acc) const {
    SmallVector<AffineMap> maps = contractOp.getIndexingMapsArray();
    ArrayRef<Attribute> iteratorTypes =
        contractOp.getIteratorTypes().getValue();

    // Given that the distribution format is <BATCH x OUTER x ELEMENT>,
    // the iterations and affine maps need to be replicated three times.

    SmallVector<Attribute> newIterators;
    // Replicate the iterators for local vector.contract
    for (int i = 0; i < 3; ++i) {
      newIterators.append(iteratorTypes.begin(), iteratorTypes.end());
    }

    // Replicate the affine maps for local vector.contract
    SmallVector<AffineMap> newMaps;
    for (AffineMap map : maps) {
      int64_t numDims = map.getNumDims();
      int64_t numResults = map.getNumResults();
      SmallVector<AffineExpr> exprs;
      for (int i = 0; i < 3; ++i) {
        AffineMap shiftedMap = map.shiftDims(i * numDims);
        for (int j = 0; j < numResults; ++j) {
          exprs.push_back(shiftedMap.getResult(j));
        }
      }
      AffineMap newMap =
          AffineMap::get(/*dimCount=*/3 * numDims,
                         /*symbolCount=*/map.getNumSymbols(), exprs, ctx);
      newMaps.push_back(newMap);
    }

    Value localInit = getCombiningIdentityValue(
        loc, rewriter, contractOp.getKind(), acc.getType());

    auto localContractOp = rewriter.create<vector::ContractionOp>(
        loc, lhs, rhs, localInit, rewriter.getAffineMapArrayAttr(newMaps),
        rewriter.getArrayAttr(newIterators), contractOp.getKind());
    localContractOp->setDiscardableAttrs(
        contractOp->getDiscardableAttrDictionary());

    return localContractOp;
  }

  FailureOr<VectorValue> doThreadReduction(RewriterBase &rewriter,
                                           NestedLayoutAttr layout,
                                           VectorValue flat,
                                           vector::CombiningKind kind,
                                           ArrayRef<bool> reductionMask) const {
    VectorType flatVecType = flat.getType();
    int64_t numElements = flatVecType.getNumElements();
    Location loc = flat.getLoc();

    auto constOp = rewriter.create<arith::ConstantOp>(
        loc, rewriter.getZeroAttr(flatVecType));
    auto res = llvm::cast<VectorValue>(constOp.getResult());

    for (unsigned i = 0; i < numElements; ++i) {
      Value extracted = rewriter.create<vector::ExtractOp>(loc, flat, i);

      // Reduce across all reduction dimensions 1-by-1.
      for (unsigned i = 0, e = reductionMask.size(); i != e; ++i) {
        if (reductionMask[i]) {
          int64_t offset = getShuffleOffset(layout, i);
          int64_t width = getShuffleWidth(layout, i);
          assert(offset <= std::numeric_limits<uint32_t>::max() &&
                 width <= std::numeric_limits<uint32_t>::max());

          extracted = rewriter.create<gpu::SubgroupReduceOp>(
              loc, extracted, combiningKindToAllReduce(kind),
              /*uniform=*/false, /*cluster_size=*/width,
              /*cluster_stride=*/offset);
        }
      }

      res = rewriter.create<vector::InsertOp>(loc, extracted, res, i);
    }

    return res;
  }

  int64_t subgroupSize;
  int64_t maxBitsPerShuffle;
};

struct DistributeTranspose final : OpDistributionPattern<vector::TransposeOp> {
  using OpDistributionPattern::OpDistributionPattern;

  LogicalResult matchAndRewrite(vector::TransposeOp transposeOp,
                                DistributionSignature &signature,
                                PatternRewriter &rewriter) const override {
    VectorValue value = transposeOp.getVector();
    VectorLayoutInterface layout = dyn_cast<NestedLayoutAttr>(signature[value]);
    if (!layout) {
      return rewriter.notifyMatchFailure(transposeOp,
                                         "layout must be NestedLayoutAttr");
    }

    /// Transpose only changes the notion of where the data carried by each
    /// thread comes from in the transposed SIMD vector. The data carried by
    /// each thread is still the same, transposed as requested by the operation.
    /// So, for distributed dimensions (thread and subgroup) transpose is a
    /// no-op.
    ///
    /// Example (indices [0-3] represent ids of the threads carrying the data):
    ///
    /// input: vector<2x4xf16>
    ///
    /// 0 0 1 1
    /// 2 2 3 3
    ///
    /// after transpose,
    ///
    /// transp: vector<4x2xf16>
    ///
    /// 0 2
    /// 0 2
    /// 1 3
    /// 1 3
    ///
    /// As it can be seen, each thread is still carrying the same data but
    /// just holds a transposed version of it.

    VectorValue input = getDistributed(rewriter, value, layout);
    // Permute batch, outer and element based on the given permutation.
    int64_t rank = value.getType().getRank();
    SmallVector<int64_t> permutation;
    for (int i = 0; i < 3; ++i) {
      for (auto it : transposeOp.getPermutation()) {
        permutation.push_back(it + (i * rank));
      }
    }
    VectorValue transposed = rewriter.create<vector::TransposeOp>(
        transposeOp.getLoc(), input, permutation);
    replaceOpWithDistributedValues(rewriter, transposeOp, transposed);
    return success();
  }
};

struct DistributeBatchOuterToLayoutConversions final
    : OpDistributionPattern<IREE::VectorExt::ToLayoutOp> {
  using OpDistributionPattern::OpDistributionPattern;

  LogicalResult matchAndRewrite(IREE::VectorExt::ToLayoutOp toLayoutOp,
                                DistributionSignature &signature,
                                PatternRewriter &rewriter) const override {
    Location loc = toLayoutOp.getLoc();
    auto input = cast<VectorValue>(toLayoutOp.getInput());
    auto output = cast<VectorValue>(toLayoutOp.getOutput());
    auto layoutA = dyn_cast<NestedLayoutAttr>(signature[input]);
    auto layoutB = dyn_cast<NestedLayoutAttr>(signature[output]);

    if (!layoutA || !layoutB) {
      return rewriter.notifyMatchFailure(toLayoutOp, "non-nested layout");
    }

    // Check if everything other than batch and outer tile matches.
    if (layoutA.getSubgroupTile() != layoutB.getSubgroupTile()) {
      return failure();
    }
    if (layoutA.getSubgroupStrides() != layoutB.getSubgroupStrides()) {
      return failure();
    }
    if (layoutA.getThreadTile() != layoutB.getThreadTile()) {
      return failure();
    }
    if (layoutA.getThreadStrides() != layoutB.getThreadStrides()) {
      return failure();
    }
    if (layoutA.getElementTile() != layoutB.getElementTile()) {
      return failure();
    }

    auto batchTileA = SmallVector<int64_t>(layoutA.getBatchTile());
    auto outerTileA = SmallVector<int64_t>(layoutA.getOuterTile());
    auto batchTileB = SmallVector<int64_t>(layoutB.getBatchTile());
    auto outerTileB = SmallVector<int64_t>(layoutB.getOuterTile());

    // Check if there is a batch/outer tile mismatch.
    if (batchTileA == batchTileB && outerTileA == outerTileB) {
      return rewriter.notifyMatchFailure(toLayoutOp,
                                         "trivial layout conversion");
    }

    SmallVector<int64_t> shapeA = layoutA.getDistributedShape();
    SmallVector<int64_t> shapeB = layoutB.getDistributedShape();
    int64_t rank = layoutA.getRank();

    // Interleave batch and outer dims by transposing.

    // Build a permutation for interleaving.
    auto interleavePermutation =
        llvm::to_vector(llvm::seq<int64_t>(shapeA.size()));
    for (int i = 0; i < rank; ++i) {
      // Batch tile : [0...rank]
      // OuterTile : [rank+1...2*rank]
      // Interleave : [batch0, outer0, batch1, outer1,...]
      interleavePermutation[2 * i] = i;
      interleavePermutation[2 * i + 1] = i + rank;
    }

    auto interleaved = rewriter.create<vector::TransposeOp>(
        loc, getDistributed(rewriter, input, layoutA), interleavePermutation);

    // Shape cast to match the new layout.

    SmallVector<int64_t> transposedShapeB(shapeB);
    applyPermutationToVector(transposedShapeB, interleavePermutation);
    Type reshapedType = VectorType::get(
        transposedShapeB, interleaved.getResultVectorType().getElementType());

    auto reshaped =
        rewriter.create<vector::ShapeCastOp>(loc, reshapedType, interleaved);

    // Inverse transpose to preserve original order.
    SmallVector<int64_t> invertedPermutation =
        invertPermutationVector(interleavePermutation);

    auto layouted = rewriter.create<vector::TransposeOp>(loc, reshaped,
                                                         invertedPermutation);

    replaceOpWithDistributedValues(rewriter, toLayoutOp, layouted.getResult());
    return success();
  }
};

struct DistributeStep final : OpDistributionPattern<vector::StepOp> {
  using OpDistributionPattern::OpDistributionPattern;

  // This is a helper aggregate
  // to hold the information about
  // a dimension.
  // For e.g. : 3x4x2 shape will
  // have lengths = [3, 4, 2]
  // and strides = [8, 2, 1]
  struct DimInfo {
    std::optional<Value> dimIdx;
    int64_t dimLen;
    int64_t dimStride;
  };

  // This is a helper function to extract the remaining
  // dimensions with their original strides once the
  // distributed dimensions are extracted out
  //         threads
  //          V
  // E.g. 3 x 4 x 2
  // This will return back remaining dimensions that
  // have lengths = [3, 2] and strides = [8, 1]
  SmallVector<DimInfo> getRemainingDims(ArrayRef<DimInfo> distributedStrides,
                                        int64_t originalLen) const {
    SmallVector<DimInfo> remainingDims;
    int64_t currLen = originalLen;
    for (const DimInfo &dInfo : distributedStrides) {
      if (dInfo.dimStride != 0) {
        int64_t dStride = dInfo.dimStride;
        int64_t dLen = dInfo.dimLen;
        int64_t higherStride = dLen * dStride;
        if (higherStride < currLen) {
          remainingDims.push_back(
              {std::nullopt, currLen / higherStride, higherStride});
        }
        currLen = dStride;
      }
    }
    remainingDims.push_back({std::nullopt, currLen, 1});
    return remainingDims;
  }

  // This is a helper to extract lengths of all dimensions
  SmallVector<int64_t> getLens(ArrayRef<DimInfo> dimInfos) const {
    SmallVector<int64_t> lens;
    lens.reserve(dimInfos.size());
    for (const DimInfo &dInfo : dimInfos) {
      lens.push_back(dInfo.dimLen);
    }
    return lens;
  }

  // This is a helper to extract strides from a given shape
  // E.g. : a shape of 2x3x4 will return strides [12, 4, 1]
  SmallVector<int64_t> getStrides(ArrayRef<int64_t> shape) const {
    int64_t elementCount = ShapedType::getNumElements(shape);
    SmallVector<int64_t> strides;
    int64_t currStride = elementCount;
    for (int64_t len : shape) {
      currStride = currStride / len;
      strides.push_back(currStride);
    }
    return strides;
  }

  // Once we are in the realm of remaining dimensions,
  // the strides are not packed. This is a helper to
  // obtain the packed strides of the remaining dimensions.
  // (See above for an example of remaining dimensions under
  //  getRemainingDims)
  SmallVector<int64_t> getPackedStrides(ArrayRef<DimInfo> dims) const {
    SmallVector<int64_t> lens = getLens(dims);
    return getStrides(lens);
  }

  // This function emulates the slicing of otherwise large constant
  // across threads and subgroups.
  VectorValue generateSlicedStep(OpBuilder &builder, Location loc,
                                 ArrayRef<DimInfo> distributedDims,
                                 int64_t distributedLen,
                                 int64_t originalLen) const {
    SmallVector<DimInfo> remainingDims =
        getRemainingDims(distributedDims, originalLen);
    SmallVector<int64_t> remainingPackedStrides =
        getPackedStrides(remainingDims);

    SmallVector<APInt> offsets;
    offsets.reserve(distributedLen);
    // As for a complex example what the following
    // maths would achieve:
    //    wave
    //     |   threads
    //     V   V
    // 2 x 3 x 4 x 2 = 0 1 2 .... 48
    // say vector.step : vector<48xindex> is to be distributed.
    // --------------------------------------------------------
    // The the distribution should be as follows:
    // wave0:
    // t0: 0 1 24 25
    // t1: 2 3 26 27
    // t2: 4 5 28 29
    // t4: 6 7 30 31
    //
    // wave1:
    // t0: 8 9 32 33
    // t1: 10 11 34 35
    // t2: 12 13 36 37
    // t4: 14 15 38 39
    // ... etc
    //
    // So wave0 & t0 value this constant offset that we generate
    // below initially. Then followed by thread and subgroup weighted
    // addition that is weighted by their stride.
    for (size_t i = 0; i < distributedLen; i++) {
      int64_t offset = 0;
      for (const auto &[dimInfo, packedStride] :
           zip(remainingDims, remainingPackedStrides)) {
        offset += ((i / packedStride) % dimInfo.dimLen) * dimInfo.dimStride;
      }
      offsets.push_back(APInt(/*width=*/64, offset));
    }
    VectorType offsetType =
        VectorType::get({distributedLen}, builder.getIndexType());
    auto constOffset = builder.create<arith::ConstantOp>(
        loc, DenseElementsAttr::get(offsetType, offsets));
    Value finalOffset = constOffset;
    for (const DimInfo &dimInfo : distributedDims) {
      assert(dimInfo.dimIdx.has_value());
      if (dimInfo.dimStride != 0) {
        auto strideVal =
            builder.create<arith::ConstantIndexOp>(loc, dimInfo.dimStride);
        auto dimIdxOffsetPerElem = builder.create<arith::MulIOp>(
            loc, strideVal, dimInfo.dimIdx.value());
        auto dimIdxOffset = builder.create<vector::BroadcastOp>(
            loc, offsetType, dimIdxOffsetPerElem);
        finalOffset =
            builder.create<arith::AddIOp>(loc, finalOffset, dimIdxOffset);
      }
    }
    return cast<VectorValue>(finalOffset);
  }

  DistributeStep(MLIRContext *context, Value threadId, int64_t subgroupSize)
      : OpDistributionPattern(context), threadId(threadId),
        subgroupSize(subgroupSize) {}
  LogicalResult matchAndRewrite(vector::StepOp stepOp,
                                DistributionSignature &signature,
                                PatternRewriter &rewriter) const override {
    Location loc = stepOp.getLoc();
    VectorValue result = stepOp.getResult();
    NestedLayoutAttr resultLayout =
        dyn_cast<NestedLayoutAttr>(signature[result]);
    if (!resultLayout) {
      return rewriter.notifyMatchFailure(
          stepOp, "missing nested layout for step op result");
    }
    SmallVector<Value> subgroupIndices, threadIndices;
    populateWarpAndThreadIndices(rewriter, threadId, subgroupSize, resultLayout,
                                 subgroupIndices, threadIndices);

    SmallVector<int64_t> undistributedShape =
        resultLayout.getUndistributedPackedShape();
    SmallVector<int64_t> undistributedStrides = getStrides(undistributedShape);
    constexpr int64_t subgroupIdx = 0;
    constexpr int64_t threadIdx = 3;

    ArrayRef<int64_t> subgroupLengths = resultLayout.getSubgroupTile();
    ArrayRef<int64_t> threadLengths = resultLayout.getThreadTile();
    // Step op by definition should be single dimensional.
    SmallVector<int64_t> distributedShape =
        signature[result].getDistributedShape();

    int64_t distributedElements = ShapedType::getNumElements(distributedShape);
    int64_t originalElements = result.getType().getNumElements();
    SmallVector<DimInfo, 2> distributedDims{
        {subgroupIndices[0], subgroupLengths[0],
         undistributedStrides[subgroupIdx]},
        {threadIndices[0], threadLengths[0], undistributedStrides[threadIdx]}};
    llvm::sort(distributedDims, [](const DimInfo &lhs, const DimInfo &rhs) {
      return lhs.dimStride > rhs.dimStride;
    });
    VectorValue slicedStepOp = generateSlicedStep(
        rewriter, loc, distributedDims, distributedElements, originalElements);
    VectorType finalSlicedStepOpType =
        VectorType::get({distributedShape}, result.getType().getElementType());
    auto finalSlicedStepOp = rewriter.create<vector::ShapeCastOp>(
        loc, finalSlicedStepOpType, slicedStepOp);
    replaceOpWithDistributedValues(rewriter, stepOp, {finalSlicedStepOp});
    return success();
  }

  Value threadId;
  int64_t subgroupSize;
};

} // namespace

void populateGPUDistributeNestedLayoutAttrPatterns(RewritePatternSet &patterns,
                                                   Value threadId,
                                                   int64_t subgroupSize,
                                                   int64_t maxBitsPerShuffle) {
  patterns.add<DistributeTransferRead, DistributeTransferWrite>(
      patterns.getContext(), threadId, subgroupSize);
  patterns.add<DistributeBroadcast, DistributeTranspose>(patterns.getContext());
  patterns.add<DistributeMultiReduction>(patterns.getContext(), subgroupSize,
                                         maxBitsPerShuffle);
  patterns.add<DistributeContract>(patterns.getContext(), subgroupSize,
                                   maxBitsPerShuffle);
  patterns.add<DistributeBatchOuterToLayoutConversions>(patterns.getContext());
  patterns.add<DistributeStep>(patterns.getContext(), threadId, subgroupSize);
}

}; // namespace mlir::iree_compiler
