compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp

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

#include "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/TargetUtils/KnownTargets.h"
#include "iree/compiler/Codegen/Utils/MarkerUtils.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/InterleavedRange.h"
#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/ROCDLTargetInfo.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Matchers.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"

#include <cassert>
#include <cstdint>
#include <optional>

#define DEBUG_TYPE "iree-codegen-gpu-utils"
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
#define DBGSNL() (llvm::dbgs() << "\n")

constexpr unsigned kShuffleBitWidth = 32;
// TODO: These are AMD GPU specific. These need to find a better home.
constexpr char kWavesPerEuAttrName[] = "waves_per_eu";

static llvm::cl::opt<std::string> clTestTarget(
    "iree-gpu-test-target",
    llvm::cl::desc(
        "The target for IR LIT tests. Format is '<arch>:<feature>@<api>', "
        "where <feature> and <api> are optional; e.g., "
        "'gfx942:+sramecc,-xnack@hip'. If <api> is missing, it will be deduced "
        "from <arch>; e.g., 'gfx*' defaults to HIP, 'sm_*' defaults to CUDA"),
    llvm::cl::init(""));

namespace mlir::iree_compiler {

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

llvm::SmallVector<mlir::linalg::ProcInfo, 2>
getGPUThreadIdsAndCounts(mlir::OpBuilder &builder, mlir::Location loc,
                         unsigned numDims,
                         llvm::ArrayRef<int64_t> workgroupSize) {
  assert(numDims <= kNumGPUDims);
  llvm::SmallVector<mlir::linalg::ProcInfo, 2> procInfo(numDims);
  std::array<gpu::Dimension, kNumGPUDims> dimAttr{
      gpu::Dimension::x, gpu::Dimension::y, gpu::Dimension::z};
  mlir::Type indexType = builder.getIndexType();
  for (unsigned i = 0; i < numDims; ++i) {
    procInfo[numDims - 1 - i] = {
        mlir::gpu::ThreadIdOp::create(builder, loc, indexType, dimAttr[i]),
        mlir::arith::ConstantOp::create(builder, loc,
                                        builder.getIndexAttr(workgroupSize[i])),
        linalg::DistributionMethod::Cyclic};
  }
  return procInfo;
}

llvm::SmallVector<mlir::linalg::ProcInfo, 2>
getSubgroupIdsAndCounts(mlir::OpBuilder &builder, mlir::Location loc,
                        unsigned warpSize, unsigned numDims,
                        llvm::ArrayRef<int64_t> numSubgroups) {
  assert(numDims <= kNumGPUDims);
  llvm::SmallVector<mlir::linalg::ProcInfo, 2> procInfo(numDims);
  std::array<gpu::Dimension, kNumGPUDims> dimAttr{
      gpu::Dimension::x, gpu::Dimension::y, gpu::Dimension::z};
  mlir::Type indexType = builder.getIndexType();
  for (unsigned i = 0; i < numDims; ++i) {
    mlir::Value subgroupId =
        mlir::gpu::ThreadIdOp::create(builder, loc, indexType, dimAttr[i]);
    if (i == 0) {
      subgroupId = affine::AffineDelinearizeIndexOp::create(
                       builder, loc, subgroupId,
                       ArrayRef<int64_t>{numSubgroups[i], warpSize})
                       .getResult(0);
    }
    procInfo[numDims - 1 - i] = {
        subgroupId,
        mlir::arith::ConstantOp::create(builder, loc,
                                        builder.getIndexAttr(numSubgroups[i])),
        linalg::DistributionMethod::Cyclic};
  }
  return procInfo;
}

bool isDescendingRelativeMappingIndices(ArrayRef<Attribute> array) {
  int64_t prev = cast<DeviceMappingAttrInterface>(array[0]).getRelativeIndex();
  for (Attribute attr : array.drop_front()) {
    int64_t relativeIndex =
        cast<DeviceMappingAttrInterface>(attr).getRelativeIndex();
    if (relativeIndex != prev - 1) {
      return false;
    }
    prev = relativeIndex;
  }
  return true;
}

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

bool canPerformVectorAccessUsingAllThreads(ArrayRef<int64_t> shape,
                                           int64_t threadCount,
                                           int64_t vectorSize) {
  // Verify that each dimension of the shape can be distributed on the
  // threads
  // For zero dim tensor, consider it's too small to access using all threads.
  if (shape.empty()) {
    return false;
  }
  int64_t threadsAvailable = threadCount;
  for (const auto &[index, dim] : llvm::enumerate(llvm::reverse(shape))) {
    int64_t numElementPerThread = index == 0 ? vectorSize : 1;
    int64_t numThreads = dim / numElementPerThread;
    if (numThreads == 0) {
      return false;
    }
    if (numThreads > threadsAvailable) {
      // If there are no enough remaining threads to distribute the current
      // dimension, try to use all remaining threads. But we still need to make
      // sure all work can be distributed to these threads evenly.
      if (numThreads % threadsAvailable != 0) {
        return false;
      }
      numThreads = threadsAvailable;
    }
    if (threadsAvailable % numThreads != 0) {
      return false;
    }
    threadsAvailable = threadsAvailable / numThreads;
    if (threadsAvailable == 1) {
      break;
    }
  }
  return threadsAvailable == 1;
}

/// 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) {
  SmallVector<int64_t> order;
  // First make reduction the outer dimensions.
  for (auto [index, iter] : llvm::enumerate(contract.getIteratorTypes())) {
    if (vector::isReductionIterator(iter)) {
      order.push_back(index);
    }
  }

  llvm::SmallDenseSet<int64_t> dims;
  for (AffineExpr expr : contract.getIndexingMapsArray()[0].getResults()) {
    dims.insert(cast<AffineDimExpr>(expr).getPosition());
  }
  // Then parallel dimensions that are part of Lhs as we want to re-use Lhs.
  for (auto [index, iter] : llvm::enumerate(contract.getIteratorTypes())) {
    if (vector::isParallelIterator(iter) && dims.count(index)) {
      order.push_back(index);
    }
  }
  // Then the remaining parallel loops.
  for (auto [index, iter] : llvm::enumerate(contract.getIteratorTypes())) {
    if (vector::isParallelIterator(iter) && !dims.count(index)) {
      order.push_back(index);
    }
  }
  return order;
}

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

const char *getGPUDistributeAttrName() { return "iree.gpu.distribute_dim"; }

FailureOr<SmallVector<int64_t>> getGPUTileSize(mlir::FunctionOpInterface funcOp,
                                               int tilingLevel) {
  SmallVector<Operation *> computeOps = getComputeOps(funcOp);
  FailureOr<IREE::Codegen::LoweringConfigAttr> config =
      getFirstLoweringConfig<IREE::Codegen::LoweringConfigAttr>(computeOps);
  if (failed(config)) {
    return funcOp.emitOpError("failed to get lowering configuration");
  }

  return config->getTileSizeVals(tilingLevel);
}

FailureOr<scf::SCFTileSizeComputationFunction>
getGPUScfTileSizeComputeFn(mlir::FunctionOpInterface funcOp, int tilingLevel) {
  FailureOr<SmallVector<int64_t>> tileSizes =
      getGPUTileSize(funcOp, tilingLevel);
  if (failed(tileSizes)) {
    return failure();
  }
  scf::SCFTileSizeComputationFunction computeFn =
      [tileSizes](OpBuilder &builder,
                  Operation *op) -> SmallVector<OpFoldResult> {
    auto tileSizesOfr = getAsIndexOpFoldResult(op->getContext(), *tileSizes);
    auto zeroAttr = builder.getIndexAttr(0);
    int numLoops = cast<TilingInterface>(op).getLoopIteratorTypes().size();
    tileSizesOfr.resize(numLoops, zeroAttr);
    return tileSizesOfr;
  };
  return computeFn;
}

bool isNonZeroRank(TypedValue<VectorType> val) {
  return val.getType().getRank() != 0;
}

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

std::optional<Value> allocateWorkgroupMemory(OpBuilder &builder,
                                             memref::SubViewOp subview,
                                             ArrayRef<Value> sizeBounds,
                                             DataLayout &) {
  OpBuilder::InsertionGuard guard(builder);

  mlir::FunctionOpInterface funcOp =
      subview->getParentOfType<mlir::FunctionOpInterface>();
  if (!funcOp) {
    return std::nullopt;
  }

  // The subview size bounds are expected to be constant; they specify the shape
  // of the allocation.
  SmallVector<int64_t, 2> shape;
  for (Value bound : sizeBounds) {
    APInt value;
    if (!matchPattern(bound, m_ConstantInt(&value))) {
      return std::nullopt;
    }
    shape.push_back(value.getSExtValue());
  }

  builder.setInsertionPoint(&funcOp.front(), funcOp.front().begin());
  auto type = MemRefType::get(
      shape, subview.getType().getElementType(), MemRefLayoutAttrInterface{},
      gpu::AddressSpaceAttr::get(builder.getContext(),
                                 gpu::GPUDialect::getWorkgroupAddressSpace()));
  Value buffer = memref::AllocOp::create(builder, funcOp.getLoc(), type);
  return buffer;
}

LogicalResult deallocateWorkgroupMemory(OpBuilder &, Value /*buffer*/) {
  return success();
}

LogicalResult copyToWorkgroupMemory(OpBuilder &b, Value src, Value dst) {
  Operation *copyOp = memref::CopyOp::create(b, src.getLoc(), src, dst);
  setMarker(copyOp, getCopyToWorkgroupMemoryMarker());
  return success();
}

static bool propagateCopyDestIntoProducerFill(memref::CopyOp copyOp) {
  // Look for a fill Op writing into the copyOp source.
  Operation *prevOp = copyOp->getPrevNode();
  while (prevOp) {
    if (isMemoryEffectFree(prevOp)) {
      prevOp = prevOp->getPrevNode();
      continue;
    }

    auto fillOp = dyn_cast<linalg::FillOp>(prevOp);
    if (!fillOp) {
      break;
    }
    if (fillOp.output() != copyOp.getSource()) {
      break;
    }
    // Move the fillOp and change the destination to the copy destination.
    fillOp->moveBefore(copyOp);
    fillOp.getOutputsMutable().assign(copyOp.getTarget());
    return true;
  }
  return false;
}

// Split input/output operand from copy from shared memory into a separate
// input.
static void insertInputValueIntoGeneric(Value source, linalg::GenericOp op) {
  SmallVector<Value> newOperands;
  SmallVector<AffineMap> maps;
  for (OpOperand *in : op.getDpsInputOperands()) {
    newOperands.push_back(in->get());
    maps.push_back(op.getMatchingIndexingMap(in));
  }
  newOperands.push_back(source);
  assert(op.getNumDpsInits() == 1);
  OpOperand *outOperand = op.getDpsInitOperand(0);
  maps.push_back(op.getMatchingIndexingMap(outOperand));
  maps.push_back(op.getMatchingIndexingMap(outOperand));
  Location loc = op.getLoc();
  SmallVector<utils::IteratorType> iterTypes(op.getNumLoops(),
                                             utils::IteratorType::parallel);
  OpBuilder builder(op);
  auto newOp = linalg::GenericOp::create(builder, loc, newOperands,
                                         outOperand->get(), maps, iterTypes);
  newOp.getRegion().getBlocks().splice(newOp.getRegion().begin(),
                                       op.getRegion().getBlocks());

  Block &payload = newOp.getRegion().front();
  payload.addArgument(payload.getArguments().back().getType(), loc);
  setMarker(newOp, getCopyToWorkgroupMemoryMarker());
}

/// Propagate the shared memory copy into the consumer op if it's a fully
/// parallel linalg.generic.
static bool
propagateCopySourceIntoConsumerGeneric(memref::CopyOp copyOp,
                                       SmallVector<Operation *> &toDelete) {
  // Look for a generic Op reading the copyOp target.
  Operation *nextOp = copyOp->getNextNode();
  while (nextOp) {
    if (isMemoryEffectFree(nextOp)) {
      nextOp = nextOp->getNextNode();
      continue;
    }
    auto consumer = dyn_cast<linalg::GenericOp>(nextOp);
    if (!consumer || consumer.getNumDpsInits() != 1 ||
        !consumer.getMatchingIndexingMap(consumer.getDpsInitOperand(0))
             .isIdentity()) {
      break;
    }
    if (*consumer.getOutputs().begin() != copyOp.getTarget()) {
      break;
    }
    insertInputValueIntoGeneric(copyOp.getSource(), consumer);
    toDelete.push_back(consumer);
    return true;
  }
  return false;
}

/// This is needed because we are doing promotion to shared memory on buffers.
/// This is a fragile and temporary solution until we move to be able to do this
/// kind of transformations on tensors.
void propagateSharedMemoryCopy(mlir::FunctionOpInterface funcOp) {
  SmallVector<Operation *> toDelete;
  funcOp.walk([&toDelete](memref::CopyOp copyOp) {
    if (hasMarker(copyOp, getCopyToWorkgroupMemoryMarker())) {
      if (propagateCopyDestIntoProducerFill(copyOp) ||
          propagateCopySourceIntoConsumerGeneric(copyOp, toDelete)) {
        toDelete.push_back(copyOp.getOperation());
      }
    }
  });
  for (Operation *op : toDelete) {
    op->erase();
  }
}

void insertBarriersAroundSharedMemoryCopy(mlir::FunctionOpInterface funcOp) {
  OpBuilder builder(funcOp.getContext());
  // Insert barriers before and after copies to workgroup memory and skip
  // insert barriers between back to back copy to workgroup memory.
  funcOp.walk([&builder](Operation *copyOp) {
    if (hasMarker(copyOp, getCopyToWorkgroupMemoryMarker())) {
      Operation *prevOp = copyOp->getPrevNode();
      if (!prevOp || !hasMarker(prevOp, getCopyToWorkgroupMemoryMarker())) {
        builder.setInsertionPoint(copyOp);
        gpu::BarrierOp::create(builder, copyOp->getLoc(),
                               gpu::AddressSpace::Workgroup);
      }
      Operation *nextOp = copyOp->getNextNode();
      if (!nextOp || !hasMarker(nextOp, getCopyToWorkgroupMemoryMarker())) {
        builder.setInsertionPointAfter(copyOp);
        gpu::BarrierOp::create(builder, copyOp->getLoc(),
                               gpu::AddressSpace::Workgroup);
      }
    }
  });
}

//===----------------------------------------------------------------------===//
// Reduction utils
//===----------------------------------------------------------------------===//

/// Packs scalar element to its vector equivalent.
/// (i.e f16 -> vector<1xf16> and f32 -> vector<1xf32>)
static Value promoteElementToVector(Location loc, OpBuilder &builder,
                                    Value input) {
  VectorType vectorTypeBroadcast = VectorType::get({1}, input.getType());
  Value vectorInput =
      vector::BroadcastOp::create(builder, loc, vectorTypeBroadcast, input);
  return vectorInput;
}

Value packVectorToSupportedWidth(Location loc, OpBuilder &builder,
                                 Value input) {
  LLVM_DEBUG({
    auto vecType = cast<VectorType>(input.getType());
    Type elementType = vecType.getElementType();
    assert(vecType.getDimSize(0) * elementType.getIntOrFloatBitWidth() ==
               kShuffleBitWidth &&
           "vecSize * vecBitWidth needs to packable into 32-bitwidth.");
    assert(elementType.isIntOrFloat() &&
           "Only int and float packing is supported.");
  });
  VectorType packed32Type = VectorType::get({1}, builder.getI32Type());
  Value packedInputVec =
      vector::BitCastOp::create(builder, loc, packed32Type, input);
  Value packedInput =
      vector::ExtractOp::create(builder, loc, packedInputVec, 0);
  return packedInput;
}

Value unpackToVector(Location loc, OpBuilder &builder, Value packedInput,
                     VectorType targetVecType) {
  LLVM_DEBUG({
    Type packedType = packedInput.getType();
    assert(packedType.isIntOrFloat() && "Only ints and floats are unpackable.");
    Type elementType = targetVecType.getElementType();
    assert(targetVecType.getDimSize(0) * elementType.getIntOrFloatBitWidth() ==
               packedType.getIntOrFloatBitWidth() &&
           "packed width needs to be unpackable to vecSize * vecBitWidth.");
  });
  Value packedVector = promoteElementToVector(loc, builder, packedInput);
  Value unpackedVector =
      vector::BitCastOp::create(builder, loc, targetVecType, packedVector);
  return unpackedVector;
}

/// Emit warp reduction code sequence for a given scalar input value.
static Value warpReduction(Location loc, OpBuilder &builder, Value input,
                           vector::CombiningKind kind, uint32_t warpSize,
                           uint32_t numLaneToReduce,
                           bool expandSubgroupReduce) {
  assert(llvm::isPowerOf2_32(numLaneToReduce));
  assert((isa<IntegerType, FloatType>(input.getType())) &&
         "Input must be a scalar");
  IntegerType shuffleIntType = builder.getIntegerType(kShuffleBitWidth);
  Type origInputType = input.getType();
  const unsigned origBitWidth = origInputType.getIntOrFloatBitWidth();
  assert(origBitWidth <= kShuffleBitWidth && "Unsupported input type bitwidth");

  const bool needsPacking = kShuffleBitWidth != origBitWidth;
  IntegerType equivIntType = builder.getIntegerType(origBitWidth);

  // Defer expansion of subgroup reduction until later in pass pipeline to
  // enable conditional lowering to DPP ops, for potential perf gains over
  // gpu.shuffle ops.
  if (!expandSubgroupReduce && numLaneToReduce <= warpSize &&
      warpSize % numLaneToReduce == 0) {
    gpu::AllReduceOperation gpuReduceKind = combiningKindToAllReduce(kind);

    // SPIRV currently doesn't have a lowering for clustered reduction,
    // so if possible avoid adding problematic attribute until it is supported.
    if (numLaneToReduce == warpSize) {
      return gpu::SubgroupReduceOp::create(builder, loc, input, gpuReduceKind,
                                           /*uniform=*/false);
    }
    return gpu::SubgroupReduceOp::create(builder, loc, input, gpuReduceKind,
                                         /*uniform=*/false, numLaneToReduce);
  }

  // Otherwise, perform the shuffles over the supported scalar type. For inputs
  // of smaller bitwidth, perform packing and unpacking via the supported
  // integer type.
  auto unpack = [loc, &builder, needsPacking, equivIntType,
                 origInputType](Value packedVal) -> Value {
    if (!needsPacking) {
      return packedVal;
    }
    auto asInt = arith::TruncIOp::create(builder, loc, equivIntType, packedVal);
    return arith::BitcastOp::create(builder, loc, origInputType, asInt);
  };

  auto pack = [loc, &builder, needsPacking, equivIntType,
               shuffleIntType](Value unpackedVal) -> Value {
    if (!needsPacking) {
      return unpackedVal;
    }
    auto asInt =
        arith::BitcastOp::create(builder, loc, equivIntType, unpackedVal);
    return arith::ExtUIOp::create(builder, loc, shuffleIntType, asInt);
  };

  // Lane value always stays in the original type. We use it to perform arith
  // reductions.
  Value laneVal = input;
  // Parallel reduction using butterfly shuffles.
  for (uint64_t i = 1; i < numLaneToReduce; i <<= 1) {
    Value shuffled = gpu::ShuffleOp::create(builder, loc, pack(laneVal), i,
                                            /*width=*/warpSize,
                                            /*mode=*/gpu::ShuffleMode::XOR)
                         .getShuffleResult();
    laneVal = makeArithReduction(builder, loc, kind, laneVal, unpack(shuffled));
  }
  // Broadcast the result to all the lanes.
  if (warpSize != numLaneToReduce) {
    Value shuffled = gpu::ShuffleOp::create(builder, loc, pack(laneVal), 0,
                                            /*width=*/warpSize,
                                            /*mode=*/gpu::ShuffleMode::IDX)
                         .getShuffleResult();
    laneVal = unpack(shuffled);
  }

  return laneVal;
}

// List of identity elements by operation.
// https://en.wikipedia.org/wiki/Identity_element
static TypedAttr getCombiningKindIdentity(OpBuilder &builder,
                                          vector::CombiningKind combiningKind,
                                          Type type) {
  switch (combiningKind) {
  case vector::CombiningKind::ADD:
    return builder.getZeroAttr(type);
  case vector::CombiningKind::MUL: {
    if (type.isIntOrIndex()) {
      return builder.getIntegerAttr(type, 1);
    }
    return builder.getFloatAttr(type, 1);
  }
  case vector::CombiningKind::MINUI:
  case vector::CombiningKind::MINSI:
    return builder.getIntegerAttr(type, std::numeric_limits<int64_t>::max());
  case vector::CombiningKind::MAXUI:
  case vector::CombiningKind::MAXSI:
    return builder.getIntegerAttr(type, std::numeric_limits<int64_t>::min());
  case vector::CombiningKind::AND:
    return builder.getIntegerAttr(type, 1);
  case vector::CombiningKind::OR:
  case vector::CombiningKind::XOR:
    return builder.getZeroAttr(type);
  case vector::CombiningKind::MINIMUMNUMF:
  case vector::CombiningKind::MAXIMUMNUMF:
    return builder.getFloatAttr(
        type, APFloat::getQNaN(cast<FloatType>(type).getFloatSemantics()));
  case vector::CombiningKind::MINIMUMF:
  case vector::CombiningKind::MINNUMF: {
    auto posInfApFloat = APFloat::getInf(
        cast<FloatType>(type).getFloatSemantics(), /*Negative=*/false);
    return builder.getFloatAttr(type, posInfApFloat);
  }
  case vector::CombiningKind::MAXIMUMF:
  case vector::CombiningKind::MAXNUMF: {
    auto negInfApFloat = APFloat::getInf(
        cast<FloatType>(type).getFloatSemantics(), /*Negative=*/true);
    return builder.getFloatAttr(type, negInfApFloat);
  }
  }
  return TypedAttr();
}

/// Emit identity variable.
Value getCombiningIdentityValue(Location loc, OpBuilder &builder,
                                vector::CombiningKind kind, Type identityType) {
  auto vectorType = dyn_cast<VectorType>(identityType);
  Type elementType = identityType;
  if (vectorType) {
    elementType = vectorType.getElementType();
  }
  TypedAttr identityAttr = getCombiningKindIdentity(builder, kind, elementType);
  if (vectorType) {
    identityAttr = DenseElementsAttr::get(vectorType, identityAttr);
  }
  assert(identityAttr && "Unknown identity value for the reduction");
  Value identity =
      arith::ConstantOp::create(builder, loc, identityType, identityAttr);
  return identity;
}

/// Returns the matching GPU reduction operation.
gpu::AllReduceOperation combiningKindToAllReduce(vector::CombiningKind kind) {
  switch (kind) {
#define MAP_CASE(X)                                                            \
  case vector::CombiningKind::X:                                               \
    return gpu::AllReduceOperation::X

    MAP_CASE(ADD);
    MAP_CASE(MUL);
    MAP_CASE(MINUI);
    MAP_CASE(MINSI);
    MAP_CASE(MINNUMF);
    MAP_CASE(MAXSI);
    MAP_CASE(MAXUI);
    MAP_CASE(MAXNUMF);
    MAP_CASE(AND);
    MAP_CASE(OR);
    MAP_CASE(XOR);
    MAP_CASE(MINIMUMF);
    MAP_CASE(MAXIMUMF);
    MAP_CASE(MINIMUMNUMF);
    MAP_CASE(MAXIMUMNUMF);
#undef MAP_CASE
  }
  // Upstream LLVM has the same assertion for the reverse direction (see
  // mlir::gpu::convertReductionKind).
  llvm_unreachable("Vector and GPU reduction kinds should match 1:1");
}

/// Emit reduction across a group for a given input.
Value emitGPUGroupReduction(Location loc, OpBuilder &builder, Value input,
                            vector::CombiningKind kind, uint32_t size,
                            int warpSize, bool expandSubgroupReduce) {
  assert(
      size % warpSize == 0 &&
      "Group reduction only support for sizes aligned on warp size for now.");

  // First reduce on a single thread to get per lane reduction value.
  Value laneVal = vector::ReductionOp::create(builder, loc, kind, input);
  laneVal = warpReduction(loc, builder, laneVal, kind, warpSize, warpSize,
                          expandSubgroupReduce);
  // Simple case -- emit `gpu.subgroup_reduce` directly.
  if (!expandSubgroupReduce && size == warpSize) {
    return laneVal;
  }

  // More-involved case -- generate `gpu.shuffle` ops over i32 values (using the
  // butterfly shuffle algorithm).
  // if we have more than one warp, reduce across warps.
  if (size > warpSize) {
    uint32_t numWarp = size / warpSize;
    assert(numWarp <= warpSize &&
           "Only support 1 level, need to implement recursive/loop for this "
           "case.");
    auto addressSpaceAttr = gpu::AddressSpaceAttr::get(
        builder.getContext(), gpu::GPUDialect::getWorkgroupAddressSpace());
    MemRefType memrefType =
        MemRefType::get(numWarp, laneVal.getType(), MemRefLayoutAttrInterface{},
                        addressSpaceAttr);
    Value alloc = memref::AllocOp::create(builder, loc, memrefType);
    Value threadX = gpu::ThreadIdOp::create(
        builder, loc, builder.getIndexType(), gpu::Dimension::x);
    Value cstWarpSize = arith::ConstantIndexOp::create(builder, loc, warpSize);
    Value warpId = arith::DivUIOp::create(builder, loc, threadX, cstWarpSize);
    Value laneId = arith::RemUIOp::create(builder, loc, threadX, cstWarpSize);
    Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
    Value lane0 = arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::eq,
                                        laneId, zero);
    // Store the reduction for each warp.
    SmallVector<Value> indices = {warpId};
    scf::IfOp::create(builder, loc, lane0, [&](OpBuilder &b, Location l) {
      memref::StoreOp::create(b, l, laneVal, alloc, indices);
      scf::YieldOp::create(b, l);
    });
    gpu::BarrierOp::create(builder, loc, gpu::AddressSpace::Workgroup);
    // Further reduce the outputs from each warps with a single warp reduce.
    Value memrefSize =
        arith::ConstantIndexOp::create(builder, loc, numWarp - 1);
    Value laneIdInBounds =
        arith::MinUIOp::create(builder, loc, laneId, memrefSize);
    Value loadVal = memref::LoadOp::create(builder, loc, alloc, laneIdInBounds);
    Value cstNumWarp = arith::ConstantIndexOp::create(builder, loc, numWarp);
    if (!llvm::isPowerOf2_32(numWarp)) {
      // Pad with identity element if numel < warpSize for valid warp reduction.
      Value useIdentityElement = arith::CmpIOp::create(
          builder, loc, arith::CmpIPredicate::sge, laneId, cstNumWarp);
      numWarp = llvm::PowerOf2Ceil(numWarp);
      Value identity =
          getCombiningIdentityValue(loc, builder, kind, loadVal.getType());
      loadVal = arith::SelectOp::create(builder, loc, useIdentityElement,
                                        identity, loadVal);
    }
    laneVal = warpReduction(loc, builder, loadVal, kind, warpSize, numWarp,
                            /*expandSubgroupReduce=*/true);
  }

  return laneVal;
}

std::optional<SmallVector<int64_t>> getWmmaNativeVectorSize(Operation *op) {
  // Currently hardcode the size of wmma operation. When more cases are
  // supported this should be picked based on what the backend supports.
  int64_t m = 16;
  int64_t n = 16;
  if (auto contract = dyn_cast<vector::ContractionOp>(op)) {
    int64_t k = contract.getLhsType().getElementType().isF16() ? 16 : 8;
    SmallVector<int64_t> nativeSize(contract.getIteratorTypes().size() - 3, 1);
    nativeSize.append({m, n, k});
    return nativeSize;
  }
  if (auto writeOp = dyn_cast<vector::TransferWriteOp>(op)) {
    if (writeOp.getVectorType().getRank() < 2) {
      return std::nullopt;
    }
    SmallVector<int64_t> nativeSize(writeOp.getVectorType().getRank() - 2, 1);
    nativeSize.append({m, n});
    return nativeSize;
  }
  if (auto readOp = dyn_cast<vector::TransferReadOp>(op)) {
    // Transfer read ops may need different shapes based on how they are being
    // used. For simplicity just match the shape used by the extract strided op.
    VectorType sliceType;
    for (Operation *users : op->getUsers()) {
      auto extract = dyn_cast<vector::ExtractStridedSliceOp>(users);
      if (!extract) {
        return std::nullopt;
      }
      auto vecType = cast<VectorType>(extract.getResult().getType());
      if (sliceType && sliceType != vecType) {
        return std::nullopt;
      }
      sliceType = vecType;
    }
    return llvm::to_vector(sliceType.getShape());
  }
  if ((OpTrait::hasElementwiseMappableTraits(op) && op->getNumResults() == 1)) {
    if (auto vecType = dyn_cast<VectorType>(op->getResultTypes()[0])) {
      // TODO: The condition for unrolling elementwise should be restricted
      // only to operations that need unrolling (connected to the contract).
      if (vecType.getRank() < 2) {
        return std::nullopt;
      }

      // First check whether there is a slice to infer the shape from. This is
      // required for cases where the accumulator type differs from the input
      // types, in which case we will see an `arith.ext_` between the contract
      // and transfer_read which needs to be unrolled.
      VectorType sliceType;
      for (Operation *users : op->getUsers()) {
        auto extract = dyn_cast<vector::ExtractStridedSliceOp>(users);
        if (!extract) {
          return std::nullopt;
        }
        auto vecType = cast<VectorType>(extract.getResult().getType());
        if (sliceType && sliceType != vecType) {
          return std::nullopt;
        }
        sliceType = vecType;
      }
      if (sliceType) {
        return llvm::to_vector(sliceType.getShape());
      }

      // Else unroll for trailing elementwise.
      SmallVector<int64_t> nativeSize(vecType.getRank() - 2, 1);
      // Map elementwise ops to the output shape.
      nativeSize.append({m, n});
      return nativeSize;
    }
  }
  return std::nullopt;
}

static FailureOr<int64_t>
getOperandBitwidth(IREE::Codegen::InnerTileDescAttrInterface intrinsic,
                   int operandIndex) {
  SmallVector<Type> elementTypes;
  intrinsic.getElementTypes(elementTypes);
  assert(operandIndex >= 0 && "operand index must be non-negative");
  return elementTypes[operandIndex].getIntOrFloatBitWidth();
}

/// Returns the number of elements along the K dimensions (all reduction
/// dimensions) of the given intrinsic. For intrinsics with multiple K
/// dimensions, returns the product of all K dimensions. This is determined by
/// looking at all reduction dimensions in the undistributed tile shape.
static FailureOr<int64_t>
getKSize(IREE::Codegen::InnerTileDescAttrInterface intrinsic) {
  int64_t kSize = 1;
  int64_t lhsOperandIndex = 0;
  SmallVector<VectorType> undistributedTypes;
  intrinsic.getUndistributedTileTypes(undistributedTypes);
  ArrayRef<int64_t> shape = undistributedTypes[lhsOperandIndex].getShape();
  SmallVector<utils::IteratorType> lhsIteratorTypes =
      intrinsic.getOperandIteratorTypes()[lhsOperandIndex];
  for (auto [dim, iteratorType] : llvm::zip_equal(shape, lhsIteratorTypes)) {
    if (iteratorType == utils::IteratorType::reduction) {
      kSize *= dim;
    }
  }
  return kSize;
}

/// Returns the number of elements each thread accesses for the given intrinsic
/// and operand index. This is determined by the `element` field of the
/// subgroup layout.
static FailureOr<int64_t>
getNumAccessElems(IREE::Codegen::InnerTileDescAttrInterface intrinsic,
                  int operandIndex) {
  IREE::GPU::MMASingleSubgroupLayout layout =
      IREE::GPU::getSingleSubgroupLayout(intrinsic, operandIndex);
  return llvm::product_of(layout.element);
}

/// Returns the total number of elements in the shape of the given intrinsic
/// with the given operand index. This is the product of all dimensions in the
/// distributed shape: outer × thread × element.
static FailureOr<int64_t>
getTotalTileElems(IREE::Codegen::InnerTileDescAttrInterface intrinsic,
                  int operandIndex) {
  IREE::GPU::MMASingleSubgroupLayout layout =
      IREE::GPU::getSingleSubgroupLayout(intrinsic, operandIndex);
  return llvm::product_of(layout.outer) * llvm::product_of(layout.thread) *
         llvm::product_of(layout.element);
}

static FailureOr<XorShuffleParams> getXorShuffleParamsForGfx950(
    IREE::GPU::TargetAttr target,
    IREE::Codegen::InnerTileDescAttrInterface intrinsic) {
  if (auto smma = dyn_cast<IREE::GPU::ScaledMMAAttr>(intrinsic)) {
    switch (smma.getIntrinsic()) {
    case IREE::GPU::ScaledMMAIntrinsic::MFMA_SCALE_F32_16x16x128_B32:
      return XorShuffleParams({/*rowElems=*/256,
                               /*accessElems=*/32});
    default:
      // TODO(muzasyed): Add more intrinsics for gfx950.
      return failure();
    }
  }
  if (auto mma = dyn_cast<IREE::GPU::MMAAttr>(intrinsic)) {
    switch (mma.getIntrinsic()) {
    case IREE::GPU::MMAIntrinsic::MFMA_F32_16x16x32_F16:
    case IREE::GPU::MMAIntrinsic::MFMA_F32_32x32x16_F16:
    case IREE::GPU::MMAIntrinsic::MFMA_F32_16x16x32_BF16:
    case IREE::GPU::MMAIntrinsic::MFMA_F32_32x32x16_BF16:
      return XorShuffleParams({/*rowElems=*/64,
                               /*accessElems=*/8});
    default:
      return failure();
    }
  }
  return failure();
}

/// Validate the XOR shuffle parameters for the given intrinsic and operand
/// index. If the parameters produce an XOR Shuffle that is not valid, return
/// failure. Else, return the swizzle parameters.
static FailureOr<XorShuffleParams>
validateXorShuffle(FailureOr<XorShuffleParams> swizzle,
                   IREE::Codegen::InnerTileDescAttrInterface intrinsic,
                   int operandIndex) {
  FailureOr<int64_t> maybeTotalTileElems =
      getTotalTileElems(intrinsic, operandIndex);
  if (failed(maybeTotalTileElems)) {
    return failure();
  }
  int64_t totalTileElems = *maybeTotalTileElems;
  if (!isXORShuffleValid(swizzle->rowElems, swizzle->accessElems,
                         totalTileElems)) {
    return failure();
  }
  return swizzle;
}

// Disabling clang-tidy for the following functions, as it will be externally
// linked to the CAPI in a future PR.
// NOLINTBEGIN(misc-use-internal-linkage)
FailureOr<XorShuffleBounds>
getXorShuffleBounds(IREE::Codegen::InnerTileDescAttrInterface intrinsic,
                    int operandIndex) {
  FailureOr<int64_t> maybeMinimumAccessElems =
      getNumAccessElems(intrinsic, operandIndex);
  FailureOr<int64_t> maybeTotalTileElems =
      getTotalTileElems(intrinsic, operandIndex);
  if (failed(maybeMinimumAccessElems) || failed(maybeTotalTileElems)) {
    return failure();
  }
  return XorShuffleBounds{*maybeMinimumAccessElems, *maybeTotalTileElems};
}

bool isXORShuffleValid(int64_t numRowElems, int64_t numAccessElems,
                       int64_t totalTileElems) {
  // The number of total tile elements we want to swizzle must be greater than
  // or equal to the number of elements in a row.
  if (totalTileElems < numRowElems) {
    LDBG()
        << "The number of row elements exceed the total elements in the tile";
    return false;
  }
  // We need at least one column of access elements to swizzle within a row.
  if (numRowElems < numAccessElems) {
    LDBG()
        << "The number of access elements exceed the total elements in the row";
    return false;
  }
  // The size of a row must evenly divide the total number of tile elements.
  if (totalTileElems % numRowElems != 0) {
    LDBG() << "The number of row elements does not evenly divide the total "
              "elements in the tile";
    return false;
  }
  // The number of access elements must evenly divide the number of row
  // elements.
  if (numRowElems % numAccessElems != 0) {
    LDBG() << "The number of access elements does not evenly divide the number "
              "of row elements";
    return false;
  }
  // The number of columns in a row must evenly divide the total number of tile
  // elements. This is to avoid incomplete rows at the end of the tile.
  if (totalTileElems % (numRowElems / numAccessElems) != 0) {
    LDBG() << "The number of columns in a row does not evenly divide the total "
              "elements in the tile";
    return false;
  }
  return true;
}

FailureOr<XorShuffleParams> getXorShuffleParamsForTunedChipset(
    IREE::GPU::TargetAttr target,
    IREE::Codegen::InnerTileDescAttrInterface intrinsic, int operandIndex) {
  FailureOr<ROCDL::TargetInfo> maybeTargetInfo =
      ROCDL::TargetInfo::get(target.getArch());
  if (failed(maybeTargetInfo)) {
    return failure();
  }
  if (maybeTargetInfo->has(llvm::AMDGPU::FEAT_GFX950_INSTS)) {
    return validateXorShuffle(getXorShuffleParamsForGfx950(target, intrinsic),
                              intrinsic, operandIndex);
  }
  return failure();
}

FailureOr<XorShuffleParams> getXorShuffleParamsForUntunedChipset(
    IREE::GPU::TargetAttr target,
    IREE::Codegen::InnerTileDescAttrInterface intrinsic,
    ArrayRef<int64_t> reductionTileSizes, int operandIndex) {
  // Compute XOR shuffle swizzle parameters for bank conflict avoidance.
  // - rowElems: Select entirety of K Tile size, may not prevent bank
  //              conflicts if the K tile size is too small.
  // - accessElems: number of contiguous elements each thread accesses,
  //                 derived from the MMA intrinsic's element layout.
  int64_t numAccessElems = getNumAccessElems(intrinsic, operandIndex).value();

  // Calculate K tile size (total K elements in shared memory) to use as row
  // width. For small K tiles, this may not reduce bank conflicts effectively.
  int64_t kTileSize =
      llvm::product_of(reductionTileSizes) * getKSize(intrinsic).value();

  // Figure out how many elements can fit across all banks of LDS.
  IREE::GPU::TargetWgpAttr wgp = target.getWgp();
  FailureOr<int64_t> bitwidth = getOperandBitwidth(intrinsic, operandIndex);
  if (failed(bitwidth)) {
    return failure();
  }
  std::optional<int64_t> workgroupMemoryBankCount =
      wgp.getWorkgroupMemoryBankCount();
  if (!workgroupMemoryBankCount.has_value()) {
    return failure();
  }
  // Assuming each bank is 4 bytes wide (32 bits).
  int64_t ldsBankWidthBits =
      (workgroupMemoryBankCount.value() * int64_t(32)) / *bitwidth;

  // Row width must be less than or equal to the row size (in elements) of LDS
  // bank width to prevent bank conflicts.
  int64_t effectiverowElems = std::min(ldsBankWidthBits, kTileSize);

  // Ensure row width is at least access width (minimum 1 column).
  effectiverowElems = std::max(effectiverowElems, numAccessElems);
  return validateXorShuffle(XorShuffleParams({/*rowElems=*/effectiverowElems,
                                              /*accessElems=*/numAccessElems}),
                            intrinsic, operandIndex);
}

FailureOr<XorShuffleParams>
getXorShuffleParams(IREE::GPU::TargetAttr target,
                    IREE::Codegen::InnerTileDescAttrInterface intrinsic,
                    ArrayRef<int64_t> reductionTileSizes, int operandIndex,
                    bool skipUntunedFallback) {
  FailureOr<XorShuffleParams> xorShuffleAttr =
      getXorShuffleParamsForTunedChipset(target, intrinsic, operandIndex);
  if (failed(xorShuffleAttr) && !skipUntunedFallback) {
    xorShuffleAttr = getXorShuffleParamsForUntunedChipset(
        target, intrinsic, reductionTileSizes, operandIndex);
  }
  return xorShuffleAttr;
}
// NOLINTEND(misc-use-internal-linkage)

FailureOr<Attribute>
getXorShuffleAttr(MLIRContext *context, Attribute baseConfigAttr,
                  IREE::GPU::TargetAttr target,
                  IREE::Codegen::InnerTileDescAttrInterface intrinsic,
                  ArrayRef<int64_t> reductionTileSizes, int operandIndex,
                  bool skipUntunedFallback) {
  FailureOr<XorShuffleParams> xorShuffleParams = getXorShuffleParams(
      target, intrinsic, reductionTileSizes, operandIndex, skipUntunedFallback);
  if (failed(xorShuffleParams)) {
    return failure();
  }
  int64_t effectiverowElems = xorShuffleParams.value().rowElems;
  int64_t numAccessElems = xorShuffleParams.value().accessElems;
  auto swizzleAttr = IREE::Codegen::XORShuffleAttr::get(
      context, effectiverowElems, numAccessElems,
      /*row_stride=*/int64_t(0),
      /*per_phase=*/int64_t(0));
  return IREE::GPU::SwizzleOperandAttr::get(context, baseConfigAttr,
                                            swizzleAttr);
}

//===----------------------------------------------------------------------===//
// getMmaNativeVectorSize
//===----------------------------------------------------------------------===//
/// Returns vector::ContractionOp operand's index where the result is used.
static std::optional<int>
getVectorContractOpOperandId(vector::ContractionOp contractOp,
                             OpResult result) {
  if (contractOp.getLhs() == result) {
    return 0;
  }
  if (contractOp.getRhs() == result) {
    return 1;
  }
  if (contractOp.getAcc() == result) {
    return 2;
  }
  return std::nullopt;
}

/// Returns vector::ContractionOp operand's index  where the
/// vector::TransferReadOp is consumed either consumed directly or via
/// vector::ExtractStridedSliceOp.
static std::optional<int>
getVectorContractOpOperandIdForVectorReadOp(Operation *op) {
  vector::ContractionOp contractOp;

  // Check if the vector::TransferReadOp is consumed directly by
  // vector::ContractionOp.
  if (op->use_empty()) {
    return std::nullopt;
  }
  Operation *firstLevelUser = *((op->getUsers()).begin());
  if (!firstLevelUser) {
    return std::nullopt;
  }
  if (auto contractOp = dyn_cast<vector::ContractionOp>(firstLevelUser)) {
    return getVectorContractOpOperandId(contractOp, op->getResult(0));
  }

  // Check if the vector::TransferReadOp is consumed indirectly by
  // vector::ContractionOp. Only check until the second level of use-def chain.
  if (firstLevelUser->use_empty()) {
    return std::nullopt;
  }
  Operation *secondLevelUser = *((firstLevelUser->getUsers()).begin());
  if (!secondLevelUser) {
    return std::nullopt;
  }
  if (auto contractOp = dyn_cast<vector::ContractionOp>(secondLevelUser)) {
    return getVectorContractOpOperandId(contractOp,
                                        firstLevelUser->getResult(0));
  }
  return std::nullopt;
}

/// Helper function to return native size for MMA.SYNC-based operations.
std::optional<SmallVector<int64_t>> getMmaNativeVectorSize(Operation *op) {
  // Shape of native Tensor Core GPU mma.sync operations.
  int64_t mmaShapeM = 16;
  int64_t mmaShapeN = 8;
  int64_t mmaShapeK;

  // Shape the mma.sync warp-level operation.
  if (auto contract = dyn_cast<vector::ContractionOp>(op)) {
    Type sourceType = contract.getLhsType().getElementType();

    // Set mmaShapeK based on sourceType.
    if (sourceType.isInteger(4)) {
      mmaShapeK = 64;
    } else if (sourceType.isInteger(8)) {
      mmaShapeK = 32;
    } else if (sourceType.isF16() || sourceType.isBF16()) {
      mmaShapeK = 16;
    } else if (sourceType.isF32()) {
      mmaShapeK = 8;
    } else {
      LDBG() << "unsupported shape for vector.contract: ";
      return std::nullopt;
    }

    // Initialize/set the starting dims of the ranked shape, such as batch,
    // to 1.
    SmallVector<int64_t> mmaShape(contract.getIteratorTypes().size() - 3, 1);
    mmaShape.append({mmaShapeM, mmaShapeN, mmaShapeK});
    LDBG() << "shape for vector.contract: " << llvm::interleaved(mmaShape);
    return mmaShape;
  }

  // Shape of warp-level vector write operation.
  if (auto writeOp = dyn_cast<vector::TransferWriteOp>(op)) {
    if (writeOp.getVectorType().getRank() < 2) {
      return std::nullopt;
    }
    SmallVector<int64_t> outputShape(writeOp.getVectorType().getRank() - 2, 1);
    outputShape.append({mmaShapeM, mmaShapeN});
    LDBG() << "shape for vector.xfer_write: " << llvm::interleaved(outputShape);
    return outputShape;
  }

  // Shape of warp-level vector read (load) operation.
  if (auto readOp = dyn_cast<vector::TransferReadOp>(op)) {
    auto resultVectorType = cast<VectorType>(readOp.getVector().getType());
    Type resultElementType = resultVectorType.getElementType();

    std::optional<int> operandId =
        getVectorContractOpOperandIdForVectorReadOp(op);
    if (!operandId) {
      LDBG() << "Failed to get operandId for vector::xfer_read: " << *op;
      return std::nullopt;
    }

    // Loading F16 values from Shared Memory to Registers.
    if (resultElementType.isF16() || resultElementType.isBF16()) {
      // For matrixC.
      if (*operandId == 2) {
        SmallVector<int64_t> readShape;
        readShape.append({mmaShapeM, mmaShapeN});
        LLVM_DEBUG({
          llvm::interleaveComma(readShape,
                                DBGS() << "shape for vector.xfer_read: ");
          llvm::dbgs() << "\n";
        });
        return readShape;
      }

      // For matrixA and matrixB.
      if (*operandId == 0 || *operandId == 1) {
        // MmaSyncOp input operands: matrixA and matrixB.
        // LDSMx1, x2, x4:
        // - LDSMx1 loads a 1 tile  of 8x8.
        // - LDSMx2 loads a 2 tiles of 8x8.
        // - LDSMx4 loads a 4 tiles of 8x8. (in use)
        // IREE uses the largest tiled load, i.e., LDSMx4.

        // MmaSyncOp source operand: matrixC.
        // matrixC is also read/written in tiled block of 16x16. In the pass
        // OptimizeVectorTransfer, matrixC reads are moved above the mainloop
        // and writes are moved below the mainloop. Thus, mma.sync read/write
        // accumulator inplace.
        SmallVector<int64_t> readShape;
        readShape.append({16, 16});
        LLVM_DEBUG({
          llvm::interleaveComma(readShape,
                                DBGS() << "shape for vector.xfer_read: ");
          llvm::dbgs() << "\n";
        });
        return readShape;
      }
    }

    // Loading F32 values from Shared Memory to Registers.
    if (resultElementType.isF32()) {
      // Set mmaShapeK for F32 datatype mma.sync.f32.tf32.m16n8k8.
      mmaShapeK = 8;

      // For matrixC.
      if (*operandId == 2) {
        SmallVector<int64_t> readShape;
        readShape.append({mmaShapeM, mmaShapeN});
        LDBG() << "shape for vector.xfer_read: "
               << llvm::interleaved(readShape);
        return readShape;
      }
      // For matrixA.
      if (*operandId == 0) {
        SmallVector<int64_t> readShape;
        readShape.append({mmaShapeM, mmaShapeK});
        LDBG() << "shape for vector.xfer_read: "
               << llvm::interleaved(readShape);
        return readShape;
      }
      // For matrixB.
      if (*operandId == 1) {
        // Do not use ldmatrix for matrixB.
        // Transfer read ops may need different shapes based on how they are
        // being used. For simplicity just match the shape used by the extract
        // strided op.
        VectorType sliceType;
        for (Operation *users : op->getUsers()) {
          auto extract = dyn_cast<vector::ExtractStridedSliceOp>(users);
          if (!extract) {
            return std::nullopt;
          }
          auto vecType = cast<VectorType>(extract.getResult().getType());
          if (sliceType && sliceType != vecType) {
            return std::nullopt;
          }
          sliceType = vecType;
        }
        LDBG() << "shape for vector.xfer_read: "
               << llvm::interleaved(sliceType.getShape());
        return llvm::to_vector(sliceType.getShape());
      }
    }
  }
  LDBG() << "unsupported shape for " << op->getName().getStringRef();
  return std::nullopt;
}

bool hasGlobalMemoryAddressSpace(MemRefType memrefType) {
  Attribute addrSpace = memrefType.getMemorySpace();
  if (!addrSpace) {
    return true;
  }
  auto intAttr = dyn_cast<IntegerAttr>(addrSpace);
  // Accept both default numeric address space and HAL descriptor type address
  // space--the former is used by LLVMGPU while the latter is used by SPIR-V.
  if (intAttr && intAttr.getInt() == 0) {
    return true;
  }
  auto gpuAttr = dyn_cast<gpu::AddressSpaceAttr>(addrSpace);
  if (gpuAttr && gpuAttr.getValue() == gpu::AddressSpace::Global) {
    return true;
  }
  auto amdgpuAttr = dyn_cast<amdgpu::AddressSpaceAttr>(addrSpace);
  if (amdgpuAttr &&
      amdgpuAttr.getValue() == amdgpu::AddressSpace::FatRawBuffer) {
    return true;
  }
  return isa<IREE::HAL::DescriptorTypeAttr>(addrSpace);
}

bool hasAMDGPUFatRawBufferAddressSpace(MemRefType memrefType) {
  Attribute addrSpace = memrefType.getMemorySpace();
  if (!addrSpace) {
    return false;
  }
  auto amdgpuAttr = dyn_cast<amdgpu::AddressSpaceAttr>(addrSpace);
  if (amdgpuAttr &&
      amdgpuAttr.getValue() == amdgpu::AddressSpace::FatRawBuffer) {
    return true;
  }
  return false;
}

bool hasSharedMemoryAddressSpace(MemRefType memrefType) {
  auto addrSpace =
      dyn_cast_if_present<gpu::AddressSpaceAttr>(memrefType.getMemorySpace());
  return addrSpace &&
         addrSpace.getValue() == gpu::GPUDialect::getWorkgroupAddressSpace();
}

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

/// Returns true if the index map represents a transpose that benefits from
/// shared mem.
bool sharedMemTransposeFilter(AffineMap indexMap) {
  if (!indexMap.isEmpty() && indexMap.isPermutation()) {
    // Ensure that the fasted moving dimension (the last one) is permuted,
    // Otherwise shared memory promotion will not benefit the operation.
    if (indexMap.getDimPosition(indexMap.getNumDims() - 1) !=
        indexMap.getNumDims() - 1) {
      return true;
    }
  }
  return false;
}

//===----------------------------------------------------------------------===//
// GPU Target Information
//===----------------------------------------------------------------------===//

IREE::GPU::TargetAttr getCLGPUTarget(MLIRContext *context) {
  if (clTestTarget.empty()) {
    return nullptr;
  }

  auto [archAndFeatures, backend] = StringRef(clTestTarget).split("@");
  if (backend.empty()) {
    // Guess what the target API is based on common scheme. This does not work
    // for cases like "ampere" which can be accepted by both CUDA and Vulkan;
    // it's very limited. So it's targeting common cases to make writing tests
    // simpler.
    if (StringRef(clTestTarget).starts_with("sm_")) {
      backend = "cuda";
    } else if (StringRef(clTestTarget).starts_with("gfx")) {
      backend = "hip";
    } else if (StringRef(clTestTarget).starts_with("adreno")) {
      backend = "vulkan";
    } else if (StringRef(clTestTarget).starts_with("apple")) {
      backend = "vulkan";
    } else if (StringRef(clTestTarget).starts_with("valhall")) {
      backend = "vulkan";
    } else if (StringRef(clTestTarget).starts_with("videocore")) {
      backend = "vulkan";
    }
  }
  auto [arch, features] = StringRef(archAndFeatures).split(':');
  // Use the target specified in the command line for testing purposes.
  return IREE::GPU::getFullTarget(backend, arch, features, context);
}

IREE::GPU::TargetAttr getGPUTargetAttr(DictionaryAttr attr) {
  return dyn_cast_if_present<IREE::GPU::TargetAttr>(getConfigTargetInfo(attr));
}

IREE::GPU::TargetAttr getGPUTargetAttr(MLIRContext *context,
                                       IREE::HAL::ExecutableTargetAttr target) {
  IREE::GPU::TargetAttr gpuTargetAttr;
  if (target) {
    gpuTargetAttr = getGPUTargetAttr(target.getConfiguration());
  }
  if (!gpuTargetAttr) {
    gpuTargetAttr = getCLGPUTarget(context);
  }
  return gpuTargetAttr;
}

IREE::GPU::TargetAttr getGPUTargetAttr(Operation *op) {
  return getGPUTargetAttr(op->getContext(),
                          IREE::HAL::ExecutableTargetAttr::lookup(op));
}

bool targetSupportsGlobalLoadDMA(IREE::GPU::TargetAttr target) {
  if (!target) {
    return false;
  }
  FailureOr<ROCDL::TargetInfo> maybeTargetInfo =
      ROCDL::TargetInfo::get(target.getArch());
  if (failed(maybeTargetInfo)) {
    return false;
  }
  return maybeTargetInfo->has(llvm::AMDGPU::FEAT_GFX950_INSTS);
}

bool targetSupportsShuffleBitwidth(IREE::GPU::TargetAttr target,
                                   unsigned bitwidth) {
  if (bitwidth <= kShuffleNativeBits) {
    return true;
  }
  return target && target.isAMD();
}

void addConfigGPUTarget(MLIRContext *context,
                        IREE::GPU::TargetAttr gpuTargetAttr,
                        SmallVectorImpl<NamedAttribute> &config) {
  addConfigTargetInfo(context, gpuTargetAttr, config);
}

IntegerAttr getConfigWavesPerEuAttr(DictionaryAttr targetConfig) {
  return targetConfig.getAs<IntegerAttr>(kWavesPerEuAttrName);
}
std::optional<int64_t> getConfigWavesPerEu(DictionaryAttr targetConfig) {
  auto attr = getConfigWavesPerEuAttr(targetConfig);
  if (attr) {
    return attr.getInt();
  }
  return std::nullopt;
}
void addConfigWavesPerEu(MLIRContext *context, int64_t wavesPerEu,
                         SmallVectorImpl<NamedAttribute> &config) {
  config.emplace_back(
      StringAttr::get(context, kWavesPerEuAttrName),
      IntegerAttr::get(IntegerType::get(context, 64), wavesPerEu));
}

std::optional<int> getGPUSubgroupSize(mlir::FunctionOpInterface func) {
  // First try to see if there is a subgroup size chosen in the CodeGen pipeline
  // configuration.
  if (std::optional<int64_t> subgroupSize = getSubgroupSize(func)) {
    return subgroupSize.value();
  }
  // Then try to find the subgroup size from the target description.
  if (IREE::GPU::TargetAttr target = getGPUTargetAttr(func)) {
    return target.getPreferredSubgroupSize();
  }
  return std::nullopt;
}

SmallVector<IREE::HAL::ExecutableVariantOp>
getExecutableVariantOps(mlir::ModuleOp moduleOp) {
  // The variant ops must have a ExecutableOp parent. Thus we iterate on the
  // ExecutableOp using `getOps()` for efficiency. We do not need to walk
  // through all the ops in the `moduleOp` in a nested fashion.
  SmallVector<IREE::HAL::ExecutableVariantOp> executableVariantOps;
  for (auto executableOp : moduleOp.getOps<IREE::HAL::ExecutableOp>()) {
    auto iter = executableOp.getOps<IREE::HAL::ExecutableVariantOp>();
    executableVariantOps.append(iter.begin(), iter.end());
  }
  return executableVariantOps;
}

Value applyInverseXorSwizzleToDMASourceOffset(
    OpBuilder &builder, Location loc, Value srcLinearOffset,
    IREE::Codegen::XORShuffleAttr swizzle, Value dest) {
  // Compute the subgroup's base offset within the full allocation by tracing
  // through view-like ops and accumulating the linear element offset.
  Value destTrace = dest;
  OpFoldResult totalOffset = builder.getIndexAttr(0);
  while (auto viewOp = destTrace.getDefiningOp<ViewLikeOpInterface>()) {
    if (auto subviewOp = dyn_cast<memref::SubViewOp>(viewOp.getOperation())) {
      auto parentType = cast<MemRefType>(subviewOp.getSource().getType());
      SmallVector<int64_t> strides;
      int64_t offset;
      LogicalResult res = parentType.getStridesAndOffset(strides, offset);
      assert(succeeded(res) && "expected strided layout on subview source");
      (void)res;
      SmallVector<OpFoldResult> strideOFRs =
          getAsIndexOpFoldResult(builder.getContext(), strides);
      auto &&[expr, values] = computeLinearIndex(totalOffset, strideOFRs,
                                                 subviewOp.getMixedOffsets());
      totalOffset =
          affine::makeComposedFoldedAffineApply(builder, loc, expr, values);
    } else {
      assert((isa<memref::ExpandShapeOp, memref::CollapseShapeOp>(
                 viewOp.getOperation())) &&
             "unexpected view-like op in dest chain");
    }
    destTrace = viewOp.getViewSource();
  }

  // Only the offset unaligned to the swizzle period affects the XOR
  // computation.
  auto cst = getConstantIntValue(totalOffset);
  Value swizzleBaseOffset;
  int64_t swizzlePeriod = swizzle.getSwizzlePeriod();
  if (!cst || (*cst % swizzlePeriod != 0)) {
    swizzleBaseOffset =
        getValueOrCreateConstantIndexOp(builder, loc, totalOffset);
  }

  // Add the subgroup's base offset within the full allocation.
  Value swizzleInput = srcLinearOffset;
  if (swizzleBaseOffset) {
    swizzleInput =
        arith::AddIOp::create(builder, loc, swizzleBaseOffset, srcLinearOffset);
  }

  // Apply the swizzle (handles access-width alignment internally).
  Value swizzled = getValueOrCreateConstantIndexOp(
      builder, loc, swizzle.swizzleOffset(builder, loc, swizzleInput, dest));

  // Subtract the subgroup base offset.
  if (swizzleBaseOffset) {
    return arith::SubIOp::create(builder, loc, swizzled, swizzleBaseOffset);
  }
  return swizzled;
}

} // namespace mlir::iree_compiler