// 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 "iree/compiler/Codegen/Dialect/GPU/Transforms/Transforms.h"
#include <cstdint>

#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.h"
#include "iree/compiler/Codegen/Utils/GPUUtils.h"
#include "iree/compiler/Codegen/Utils/MarkerUtils.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/MatchUtils.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "llvm/Support/Debug.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/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Support/LLVM.h"

#define DEBUG_TYPE "iree-codegen-gpu-transforms"

namespace mlir::iree_compiler::IREE::GPU {

//===---------------------------------------------------------------------===//
// Forall Fusion
//===---------------------------------------------------------------------===//

static FailureOr<SmallVector<scf::ForallOp>>
getEquivalentMappingConsumerLoopNest(scf::ForallOp producer,
                                     scf::ForallOp consumer) {
  auto compareMappingTypes = [&](ArrayRef<Attribute> l, ArrayRef<Attribute> r) {
    return (llvm::all_of(l, llvm::IsaPred<gpu::GPUThreadMappingAttr>) &&
            llvm::all_of(r, llvm::IsaPred<gpu::GPUThreadMappingAttr>)) ||
           (llvm::all_of(l, llvm::IsaPred<gpu::GPUWarpMappingAttr>) &&
            llvm::all_of(r, llvm::IsaPred<gpu::GPUWarpMappingAttr>));
  };

  ArrayRef<Attribute> producerMapping = producer.getMappingAttr().getValue();
  ArrayRef<Attribute> consumerMapping = consumer.getMappingAttr().getValue();

  if (producerMapping.empty() || consumerMapping.empty()) {
    return failure();
  }

  // Require descending relative indices so that the linearization and
  // delinearization done in subsequent steps are valid.
  if (!isDescendingRelativeMappingIndices(producerMapping) ||
      !isDescendingRelativeMappingIndices(consumerMapping)) {
    return failure();
  }

  // If both loops share the same kind of mapping, return the sole consumer.
  if (compareMappingTypes(producerMapping, consumerMapping)) {
    return SmallVector<scf::ForallOp>({consumer});
  }

  // The only other supported case is fusing a thread mapped loop into a nest
  // of a warp and lane forall.
  if (!llvm::all_of(producerMapping,
                    llvm::IsaPred<gpu::GPUThreadMappingAttr>) ||
      !llvm::all_of(consumerMapping, llvm::IsaPred<IREE::GPU::LaneIdAttr>)) {
    return failure();
  }
  auto outerWarpLoop = consumer->getParentOfType<scf::ForallOp>();
  if (!outerWarpLoop || !llvm::all_of(outerWarpLoop.getMappingAttr().getValue(),
                                      llvm::IsaPred<gpu::GPUWarpMappingAttr>)) {
    return failure();
  }
  return SmallVector<scf::ForallOp>({outerWarpLoop, consumer});
}

static FailureOr<Value> createSharedAllocDestination(RewriterBase &rewriter,
                                                     scf::ForallOp forallOp) {
  if (forallOp->getNumResults() != 1) {
    return failure();
  }

  // Skip swizzle hint ops.
  Operation *destination = forallOp.getDpsInits()[0].getDefiningOp();
  auto swizzleOp = dyn_cast<IREE::Codegen::SwizzleHintOp>(destination);
  if (swizzleOp) {
    destination = swizzleOp->getOperand(0).getDefiningOp();
  }

  // Fail if the destination is not a `tensor.empty` op and cannot be trivially
  // converted to a `bufferization.alloc_tensor`.
  auto empty = dyn_cast<tensor::EmptyOp>(destination);
  if (!empty) {
    return failure();
  }

  // Create a `bufferization.alloc_tensor` op with memory space
  // `#gpu.addresses_space<workgroup>`.
  Location loc = empty->getLoc();
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(empty);
  Attribute sharedMemoryAddrSpace = gpu::AddressSpaceAttr::get(
      rewriter.getContext(), gpu::GPUDialect::getWorkgroupAddressSpace());
  auto allocTensor = bufferization::AllocTensorOp::create(
      rewriter, loc, cast<TensorType>(empty.getResult().getType()),
      empty.getDynamicSizes(),
      /*copy=*/Value(), /*size_hint=*/Value(),
      /*memory_space=*/sharedMemoryAddrSpace);

  // If the original `tensor.empty` has a swizzle hint, apply it to the new
  // allocation. Note that if there is a swizzle hint, it will be the only user
  // of the `tensor.empty` op.
  if (swizzleOp) {
    auto newSwizzle = IREE::Codegen::SwizzleHintOp::create(
        rewriter, loc, allocTensor.getResult(), swizzleOp.getSwizzle());
    return newSwizzle.getResult();
  }
  return allocTensor.getResult();
}

LogicalResult fuseForallIntoConsumer(RewriterBase &rewriter,
                                     scf::ForallOp producer,
                                     scf::ForallOp consumer,
                                     SmallVector<Operation *> consumerChain) {
  // TODO: Support multi-result producer loops.
  if (producer->getNumResults() != 1) {
    return failure();
  }

  FailureOr<SmallVector<scf::ForallOp>> consumerLoopNest =
      getEquivalentMappingConsumerLoopNest(producer, consumer);
  if (failed(consumerLoopNest)) {
    return failure();
  }

  // Verify that all loops are normalized.
  if (!producer.isNormalized() ||
      !llvm::all_of(consumerLoopNest.value(), [](scf::ForallOp forall) {
        return forall.isNormalized();
      })) {
    return failure();
  }

  // Step 1. Get the destination of the producer loop as a shared memory
  // allocation.
  rewriter.setInsertionPointToStart(consumer.getBody());
  FailureOr<Value> maybeDest = createSharedAllocDestination(rewriter, producer);
  if (failed(maybeDest)) {
    return failure();
  }
  Value sharedDest = maybeDest.value();

  // Step 2. Move the consumer chain to right before the last user in the
  // chain.
  if (!consumerChain.empty()) {
    Operation *base = consumerChain.back();
    for (Operation *op : consumerChain) {
      if (op == base) {
        continue;
      }
      rewriter.moveOpBefore(op, base);
    }
  }

  // Step 3. Create the `iree_gpu.barrier_region` to wrap the fused producer.
  auto barrierOp = IREE::GPU::BarrierRegionOp::create(
      rewriter, producer.getLoc(), sharedDest.getType(), sharedDest);
  rewriter.setInsertionPointToStart(barrierOp.getBody());

  // Step 4. Compute the producer IDs in terms of the consumer IDs.
  // The producer IDs are computed as follows:
  //
  // producer = [p0, ..., pn] ∈ [0, ..., 0] to [P0, ..., Pn]
  // consumer = [c0, ..., cn] ∈ [0, ..., 0] to [C0, ..., Cn]
  //
  //                   Not a real op
  //                         |
  // %ub = P0 * ... * Pn     |
  // %step = C0 * ... * Cn   v
  // %flatc = affine.linearize_index %c0, ..., %cn
  // scf.for %id = %flatc to %ub step %step {
  //   %p:n = affine.delinearize_index %id into [%P0, ..., %Pn]
  //   ...
  // }
  //
  // Note: We use 0 as the loop lower bound instead of the linearized consumer
  // loop ID if possible to make later loop promotion patterns easier.

  MLIRContext *context = rewriter.getContext();
  Location loc = producer.getLoc();

  // Compute the linearize consumer loop ID and total consumer loop worker
  // count (C0 * ... * Cn).
  AffineExpr d0, d1, d2;
  bindDims(context, d0, d1, d2);
  AffineExpr mulAdd = d0 * d1 + d2;
  OpFoldResult linearId = rewriter.getIndexAttr(0);
  OpFoldResult consumerWorkerCount = rewriter.getIndexAttr(1);
  for (auto loop : *consumerLoopNest) {
    for (auto [inductionVar, workerCount] :
         llvm::zip_equal(getAsOpFoldResult(loop.getInductionVars()),
                         loop.getMixedUpperBound())) {
      linearId = affine::makeComposedFoldedAffineApply(
          rewriter, loc, mulAdd, {linearId, workerCount, inductionVar});
      consumerWorkerCount = affine::makeComposedFoldedAffineApply(
          rewriter, loc, d0 * d1, {consumerWorkerCount, workerCount});
    }
  }

  // Compute the total producer loop worker count (P0 * ... * Pn).
  Value linearConsumerIdVal =
      getValueOrCreateConstantIndexOp(rewriter, loc, linearId);
  SmallVector<OpFoldResult> producerRanges;
  OpFoldResult producerWorkerCount = rewriter.getIndexAttr(1);
  for (auto workerCount : producer.getMixedUpperBound()) {
    producerRanges.push_back(workerCount);
    producerWorkerCount = affine::makeComposedFoldedAffineApply(
        rewriter, loc, d0 * d1, {producerWorkerCount, workerCount});
  }

  std::optional<int64_t> staticProducerCount =
      getConstantIntValue(producerWorkerCount);
  std::optional<int64_t> staticConsumerCount =
      getConstantIntValue(consumerWorkerCount);
  bool perfectlyDivides =
      staticConsumerCount && staticProducerCount &&
      staticProducerCount.value() % staticConsumerCount.value() == 0;

  // Step 5. Create the `scf.for` loop for the producer.
  // If the consumer worker count perfectly divides the producer worker count,
  // then we can use a lower bound of 0 and keep the loop bounds static.
  Value lb = perfectlyDivides ? arith::ConstantIndexOp::create(rewriter, loc, 0)
                              : linearConsumerIdVal;
  Value ub =
      getValueOrCreateConstantIndexOp(rewriter, loc, producerWorkerCount);
  Value step =
      getValueOrCreateConstantIndexOp(rewriter, loc, consumerWorkerCount);
  auto newProducer = scf::ForOp::create(rewriter, loc, lb, ub, step,
                                        barrierOp.getBody()->getArgument(0));
  setLoopUnrollMarker(newProducer);
  Block *loopBody = newProducer.getBody();

  // Get the replacement IDs for the producer loop.
  rewriter.setInsertionPointToStart(loopBody);
  Value newFlatProducerId =
      perfectlyDivides
          ? affine::makeComposedAffineApply(
                rewriter, loc, d0 + d1,
                {newProducer.getInductionVar(), linearConsumerIdVal})
          : newProducer.getInductionVar();

  // We require a descending relative mapping and scf.forall loop ranges are
  // listed from outer most to inner most, so we can use the ranges directly
  // for the delinearization basis.
  auto delinearize = affine::AffineDelinearizeIndexOp::create(
      rewriter, loc, newFlatProducerId, llvm::to_vector(producerRanges));

  SmallVector<Value> newBlockArgs = delinearize.getResults();
  newBlockArgs.append(newProducer.getRegionIterArgs().begin(),
                      newProducer.getRegionIterArgs().end());

  // Step 6. Inline the region of the producer and replace the terminator.
  scf::InParallelOp terminator = producer.getTerminator();
  rewriter.mergeBlocks(producer.getBody(), loopBody, newBlockArgs);

  rewriter.setInsertionPointAfter(terminator);
  if (isa<tensor::ParallelInsertSliceOp>(
          *terminator.getYieldingOps().begin())) {
    auto parallelInsert = cast<tensor::ParallelInsertSliceOp>(
        *terminator.getYieldingOps().begin());

    // Create an insert_slice to yield from the loop body.
    SmallVector<OpFoldResult, 4> sourceOffsets =
        parallelInsert.getMixedOffsets();
    SmallVector<OpFoldResult, 4> sourceSizes = parallelInsert.getMixedSizes();
    SmallVector<OpFoldResult, 4> sourceStrides =
        parallelInsert.getMixedStrides();
    Value insertedSlice = tensor::InsertSliceOp::create(
        rewriter, loc, parallelInsert.getSource(), parallelInsert.getDest(),
        parallelInsert.getMixedOffsets(), parallelInsert.getMixedSizes(),
        parallelInsert.getMixedStrides());
    scf::YieldOp::create(rewriter, loc, insertedSlice);
    rewriter.eraseOp(parallelInsert);
    rewriter.eraseOp(terminator);
  } else {
    assert(isa<IREE::GPU::CoalescedGatherDMAOp>(
               *terminator.getYieldingOps().begin()) &&
           "expected coalesced gather dma op");
    // Create the new CoalescedGatherDMAOp outside the in_parallel region
    // and remove the in_parallel terminator.
    auto coalescedGather = cast<IREE::GPU::CoalescedGatherDMAOp>(
        *terminator.getYieldingOps().begin());

    // Create the new CoalescedGatherDMAOp with a result outside the in_parallel
    rewriter.setInsertionPoint(terminator);
    auto newGatherOp = IREE::GPU::CoalescedGatherDMAOp::create(
        rewriter, loc, coalescedGather.getInit().getType(),
        coalescedGather.getSource(), coalescedGather.getIndices(),
        coalescedGather.getInit(), coalescedGather.getLane(),
        coalescedGather.getInBoundsAttr());
    Value gatherResult = newGatherOp.getResult();

    // Use a tensor.insert_slice to insert the gather result back into the
    // shared memory destination. Extract offsets from the init's extract_slice.
    SmallVector<OpFoldResult> destIndices;
    SmallVector<OpFoldResult> sizes;
    SmallVector<OpFoldResult> strides;

    // Get offsets, sizes, and strides from the init's extract_slice op
    auto extractSlice =
        coalescedGather.getInit().getDefiningOp<tensor::ExtractSliceOp>();
    if (extractSlice) {
      destIndices = extractSlice.getMixedOffsets();
      sizes = extractSlice.getMixedSizes();
      strides = extractSlice.getMixedStrides();
    } else {
      // Fallback: use offset 0 if init is not from an extract_slice
      auto initType = cast<ShapedType>(coalescedGather.getInit().getType());
      int64_t rank = initType.getRank();
      destIndices.assign(rank, rewriter.getIndexAttr(0));
      sizes =
          getAsIndexOpFoldResult(rewriter.getContext(), initType.getShape());
      strides.assign(rank, rewriter.getIndexAttr(1));
    }

    // Get the destination tensor from the loop's iter_args
    Value dest = newProducer.getRegionIterArgs()[0];
    Value insertedSlice = tensor::InsertSliceOp::create(
        rewriter, loc, gatherResult, dest, destIndices, sizes, strides);

    // Yield the result with the inserted slice
    scf::YieldOp::create(rewriter, loc, insertedSlice);

    // Erase the old coalesced gather op and the in_parallel terminator
    rewriter.eraseOp(coalescedGather);
    rewriter.eraseOp(terminator);
  }

  // Step 7. Yield the result of the loop from the barrier op and replace the
  // producer.
  rewriter.setInsertionPointToEnd(barrierOp.getBody());
  IREE::GPU::YieldOp::create(rewriter, loc, newProducer.getResults());

  rewriter.replaceOp(producer, barrierOp);
  return success();
}

/// Collect pairs of combining ops from warp and lane forall operations.
template <typename WarpOpT, typename LaneOpT>
static SmallVector<std::pair<WarpOpT, LaneOpT>>
collectInsertPairs(scf::ForallOp warpForallOp, scf::ForallOp laneForallOp) {
  SmallVector<std::pair<WarpOpT, LaneOpT>> insertPairs;
  for (auto [warpArg, laneArg] : llvm::zip_equal(
           warpForallOp.getRegionOutArgs(), laneForallOp.getRegionOutArgs())) {
    auto warpInsert = cast<WarpOpT>(warpForallOp.getCombiningOps(warpArg)[0]);
    auto laneInsert = cast<LaneOpT>(laneForallOp.getCombiningOps(laneArg)[0]);
    insertPairs.push_back({warpInsert, laneInsert});
  }
  return insertPairs;
}

/// Inline warp and lane forall bodies into the thread forall.
static void inlineForallBodies(RewriterBase &rewriter,
                               scf::ForallOp threadForallOp,
                               scf::ForallOp warpForallOp,
                               scf::ForallOp laneForallOp) {
  // The terminator will be replaced with the terminator of the inlined block.
  rewriter.eraseOp(threadForallOp.getTerminator());
  SmallVector<Value> replacementArgs(threadForallOp.getInductionVars());
  replacementArgs.pop_back();
  llvm::append_range(replacementArgs, threadForallOp.getRegionOutArgs());
  rewriter.mergeBlocks(warpForallOp.getBody(), threadForallOp.getBody(),
                       replacementArgs);
  SmallVector<Value> innerReplacementArgs{
      threadForallOp.getInductionVars().back()};
  llvm::append_range(innerReplacementArgs, laneForallOp.getOutputs());
  rewriter.inlineBlockBefore(laneForallOp.getBody(), laneForallOp,
                             innerReplacementArgs);
}

/// Compose and fuse parallel insert slice operations from warp and lane
/// foralls.
static void composeParallelInsertSlices(
    RewriterBase &rewriter,
    SmallVector<std::pair<tensor::ParallelInsertSliceOp,
                          tensor::ParallelInsertSliceOp>> &insertPairs) {
  for (auto [warpInsert, laneInsert] : insertPairs) {
    SmallVector<OpFoldResult> composedOffsets;
    for (auto [warpOffset, laneOffset] : llvm::zip_equal(
             warpInsert.getMixedOffsets(), laneInsert.getMixedOffsets())) {
      SmallVector<Value> offsets;
      if (auto warpOffsetVal = dyn_cast<Value>(warpOffset)) {
        offsets.push_back(warpOffsetVal);
      }
      if (auto laneOffsetVal = dyn_cast<Value>(laneOffset)) {
        offsets.push_back(laneOffsetVal);
      }
      OpBuilder::InsertionGuard g(rewriter);
      rewriter.setInsertionPoint(warpInsert);
      if (!offsets.empty()) {
        (void)setInsertionPointAfterLastValue(rewriter, offsets);
      }
      composedOffsets.push_back(IREE::LinalgExt::addOfrs(
          rewriter, laneInsert.getLoc(), warpOffset, laneOffset));
    }
    OpBuilder::InsertionGuard g(rewriter);
    rewriter.setInsertionPoint(warpInsert);
    rewriter.replaceOpWithNewOp<tensor::ParallelInsertSliceOp>(
        warpInsert, laneInsert.getSource(), warpInsert.getDest(),
        composedOffsets, laneInsert.getMixedSizes(),
        laneInsert.getMixedStrides());
    rewriter.eraseOp(laneInsert);
  }
}

/// Compose and fuse coalesced gather DMA operations with parallel insert
/// slices. Creates a DMA op with a result, then uses parallel_insert_slice
/// to insert the result into the forall's shared_out at the composed offsets.
static void composeCoalescedGatherDMA(
    RewriterBase &rewriter,
    SmallVector<std::pair<tensor::ParallelInsertSliceOp,
                          IREE::GPU::CoalescedGatherDMAOp>> &insertPairs) {
  for (auto [warpInsert, laneInsert] : insertPairs) {
    OpBuilder::InsertionGuard g(rewriter);

    // Create an extract_slice for the DMA's init to match the source shape.
    // This is done outside the in_parallel region.
    rewriter.setInsertionPoint(warpInsert->getParentOp());
    auto destSlice = tensor::ExtractSliceOp::create(
        rewriter, warpInsert.getLoc(), warpInsert.getDest(),
        warpInsert.getMixedOffsets(), warpInsert.getMixedSizes(),
        warpInsert.getMixedStrides());

    // Create new CoalescedGatherDMAOp with a result (outside in_parallel).
    auto dmaOp = IREE::GPU::CoalescedGatherDMAOp::create(
        rewriter, warpInsert.getLoc(), destSlice.getType(),
        laneInsert.getSource(), laneInsert.getIndices(), destSlice,
        laneInsert.getLane(), laneInsert.getInBoundsAttr());

    // Replace the warp parallel_insert_slice with one that inserts the DMA
    // result.
    rewriter.setInsertionPoint(warpInsert);
    rewriter.replaceOpWithNewOp<tensor::ParallelInsertSliceOp>(
        warpInsert, dmaOp.getResult(), warpInsert.getDest(),
        warpInsert.getMixedOffsets(), warpInsert.getMixedSizes(),
        warpInsert.getMixedStrides());
    rewriter.eraseOp(laneInsert);
  }
}

/// Check if a forall op uses iree_gpu.coalesced_gather_dma operations
/// instead of tensor.parallel_insert_slice in its terminator.
static bool isCoalescedGatherForallOp(scf::ForallOp forallOp) {
  return llvm::any_of(forallOp.getRegionOutArgs(), [&](BlockArgument outArg) {
    return llvm::any_of(forallOp.getCombiningOps(outArg),
                        llvm::IsaPred<IREE::GPU::CoalescedGatherDMAOp>);
  });
}

FailureOr<scf::ForallOp>
fuseNestedLaneAndWarpForalls(RewriterBase &rewriter, scf::ForallOp warpForallOp,
                             scf::ForallOp laneForallOp) {
  // Verify mappings.
  if (!warpForallOp.getMapping() ||
      !llvm::all_of(*warpForallOp.getMapping(),
                    llvm::IsaPred<gpu::GPUWarpMappingAttr>)) {
    return rewriter.notifyMatchFailure(warpForallOp, "not a warp forall op");
  }
  if (!laneForallOp.getMapping() || laneForallOp.getMapping()->size() != 1 ||
      !isa<IREE::GPU::LaneIdAttr>(laneForallOp.getMapping()->getValue()[0])) {
    return rewriter.notifyMatchFailure(
        laneForallOp, "inner forall op is not mapped to a single lane id");
  }
  // Verify the lane forall is the only nested forall op.
  SmallVector<scf::ForallOp> innerForallOps(
      warpForallOp.getBody()->getOps<scf::ForallOp>());
  if (innerForallOps.size() != 1) {
    return rewriter.notifyMatchFailure(warpForallOp,
                                       "expected a single inner forall op");
  }
  if (warpForallOp.getOperation() != laneForallOp->getParentOp()) {
    return rewriter.notifyMatchFailure(
        laneForallOp, "expected warp forall op to be the lane forall's parent");
  }
  // Only allow arith/affine ops and tensor.extract_slice. We would want other
  // ops to be fused into the lane forall before this transformation happens.
  for (Operation &op : warpForallOp.getBody()->getOperations()) {
    if (!isa_and_nonnull<arith::ArithDialect, affine::AffineDialect>(
            op.getDialect()) &&
        !isa<tensor::ExtractSliceOp, scf::ForallOp, scf::InParallelOp>(op)) {
      return rewriter.notifyMatchFailure(
          warpForallOp,
          "Warp forall body has non arith, affine, or extract_slice ops");
    }
  }
  if (!warpForallOp.isNormalized() || !laneForallOp.isNormalized()) {
    return rewriter.notifyMatchFailure(warpForallOp,
                                       "forall ops are not normalized");
  }
  if (laneForallOp.getNumResults() != warpForallOp.getNumResults()) {
    return rewriter.notifyMatchFailure(
        laneForallOp,
        "lane and warp foralls have a different number of results");
  }
  auto hasSingleNonStridedFullRankParallelInsert = [](scf::ForallOp forall) {
    for (BlockArgument initArg : forall.getRegionOutArgs()) {
      SmallVector<Operation *> combiningOps = forall.getCombiningOps(initArg);
      if (combiningOps.size() != 1) {
        return false;
      }
      if (isa<IREE::GPU::CoalescedGatherDMAOp>(combiningOps[0])) {
        continue;
      }
      auto parallelInsertOp =
          dyn_cast<tensor::ParallelInsertSliceOp>(combiningOps[0]);
      if (!parallelInsertOp ||
          !areAllConstantIntValue(parallelInsertOp.getMixedStrides(), 1)) {
        return false;
      }
      if (parallelInsertOp.getSourceType().getRank() !=
          parallelInsertOp.getDestType().getRank()) {
        return false;
      }
    }
    return true;
  };
  if (!hasSingleNonStridedFullRankParallelInsert(warpForallOp) ||
      !hasSingleNonStridedFullRankParallelInsert(laneForallOp)) {
    return rewriter.notifyMatchFailure(warpForallOp,
                                       "forall op has strided combining op");
  }

  // Verify that the source of all combining ops of the warpForallOp are
  // produced by the laneForallOp, and that the laneForallOp combining ops
  // have the same rank as the corresponding combining op in the warpForallOp.
  // This is the requirement for composing the combining ops into a single
  // thread distributed combining op.
  for (auto [warpResult, laneResult] :
       llvm::zip_equal(warpForallOp.getResults(), laneForallOp.getResults())) {
    unsigned int resultIdx = warpResult.getResultNumber();
    BlockArgument warpInitArg = warpForallOp.getTiedBlockArgument(
        &warpForallOp.getOutputsMutable()[resultIdx]);
    auto warpInsertOp = cast<tensor::ParallelInsertSliceOp>(
        warpForallOp.getCombiningOps(warpInitArg)[0]);
    if (warpInsertOp.getSource() != laneResult) {
      return rewriter.notifyMatchFailure(
          warpForallOp, "combining op source is not inner lane forall result");
    }
  }

  // Create the thread mapped forall to replace the nested foralls.
  SmallVector<Attribute> threadMappings = llvm::map_to_vector(
      llvm::seq<unsigned>(warpForallOp.getRank() + laneForallOp.getRank()),
      [&](unsigned mappingDim) -> Attribute {
        unsigned mappingId =
            static_cast<unsigned>(gpu::MappingId::LinearDim0) + mappingDim;
        return gpu::GPUThreadMappingAttr::get(
            rewriter.getContext(), static_cast<gpu::MappingId>(mappingId));
      });
  std::reverse(threadMappings.begin(), threadMappings.end());
  SmallVector<OpFoldResult> lbs(threadMappings.size(),
                                rewriter.getIndexAttr(0));
  SmallVector<OpFoldResult> steps(threadMappings.size(),
                                  rewriter.getIndexAttr(1));
  SmallVector<OpFoldResult> ubs = llvm::to_vector(llvm::concat<OpFoldResult>(
      warpForallOp.getMixedUpperBound(), laneForallOp.getMixedUpperBound()));
  scf::ForallOp threadForallOp = scf::ForallOp::create(
      rewriter, warpForallOp.getLoc(), lbs, ubs, steps,
      warpForallOp.getOutputs(),
      ArrayAttr::get(rewriter.getContext(), threadMappings));

  bool hasCoalescedGatherDMA = isCoalescedGatherForallOp(laneForallOp);

  // The lane forall's terminator also needs to be erased after inlining.
  Operation *laneForallTerminator = laneForallOp.getBody()->getTerminator();
  if (!hasCoalescedGatherDMA) {
    // Collect pairs of combining ops to compose after inlining everything.
    auto insertPairs = collectInsertPairs<tensor::ParallelInsertSliceOp,
                                          tensor::ParallelInsertSliceOp>(
        warpForallOp, laneForallOp);

    // Inline the warp and lane forall bodies into the thread forall, and remap
    // the induction variables.
    inlineForallBodies(rewriter, threadForallOp, warpForallOp, laneForallOp);
    composeParallelInsertSlices(rewriter, insertPairs);
  } else {
    // Collect pairs of combining ops to compose after inlining everything.
    auto insertPairs = collectInsertPairs<tensor::ParallelInsertSliceOp,
                                          IREE::GPU::CoalescedGatherDMAOp>(
        warpForallOp, laneForallOp);

    // Inline the warp and lane forall bodies into the thread forall, and remap
    // the induction variables.
    inlineForallBodies(rewriter, threadForallOp, warpForallOp, laneForallOp);
    composeCoalescedGatherDMA(rewriter, insertPairs);
  }
  rewriter.eraseOp(laneForallTerminator);
  rewriter.eraseOp(laneForallOp);
  return threadForallOp;
}

/// Return whether a parallel insert slice operation can be collapsed with
/// the given reassociation indices. For a slice to be collapsible, each group
/// of collapsed dimensions must be fully contiguous in the destination type.
/// For example, the following parallel_insert_slice can be collapsed:
/// ```
///   tensor.parallel_insert_slice ...
///     tensor<1x2x3x4x5xf32> into tensor<?x?x3x?x5xf32>
///   ...
///   // Collapsed by:
///   tensor.collapse_shape %result [[0, 1, 2], [3, 4]]
///     tensor<?x?x3x?x5> into tensor<?x?x?xf32>
/// ```
/// Since the slice <1x2x3> is contiguous in <?x?x3> for the first group, and
/// the slice <4x5> is contiguous in <?x5> for the second group.
static LogicalResult
collapsibleSlicePrecondition(RewriterBase &rewriter,
                             tensor::ParallelInsertSliceOp sliceOp,
                             SmallVector<ReassociationIndices> reassociations) {
  if (!areAllConstantIntValue(sliceOp.getMixedStrides(), 1)) {
    return rewriter.notifyMatchFailure(sliceOp, "strides are not all 1");
  }
  RankedTensorType sliceType = sliceOp.getSourceType();
  if (sliceOp.getMixedSizes().size() != sliceType.getRank()) {
    return rewriter.notifyMatchFailure(
        sliceOp, "parallel insert slice is rank reducing");
  }

  SmallVector<OpFoldResult> sizes = sliceOp.getMixedSizes();
  RankedTensorType fullTensorType = sliceOp.getDestType();
  ArrayRef<int64_t> destShape = fullTensorType.getShape();
  for (auto group : reassociations) {
    bool isFullSlice = true;
    for (auto idx : llvm::reverse(group)) {
      std::optional<int64_t> constSize = getConstantIntValue(sizes[idx]);
      // If the size is dynamic, then conservatively assume it is not full.
      if (!constSize.has_value()) {
        if (!isFullSlice) {
          return rewriter.notifyMatchFailure(
              sliceOp,
              "parallel insert slice is not contiguous in the destination");
        }
        isFullSlice = false;
        continue;
      }
      if (constSize.value() == 1) {
        if (destShape[idx] == 1) {
          continue;
        }
        // Unit slices are okay as long as they are the outermost sliced dims.
        // Any other non-unit sliced dimension that is more outer than the unit
        // slice would be invalid, so we set `isFullSlice` to false.
        isFullSlice = false;
        continue;
      }
      // If the size is not unit, then the slice must be full so far.
      if (!isFullSlice) {
        return rewriter.notifyMatchFailure(
            sliceOp,
            "parallel insert slice is not contiguous in the destination");
      }
      if (constSize.value() != destShape[idx]) {
        isFullSlice = false;
      }
    }
  }

  return success();
}

/// Collapse all `ops` with the given `reassociations`. All `ops` are expected
/// to have equivalent offsets, sizes, and strides. All strides are expected to
/// be 1. This function assumes that the parallelInsertOp passes the
/// collapsibleSlicePrecondition.
static tensor::ParallelInsertSliceOp
collapseParallelInsertOp(RewriterBase &rewriter,
                         tensor::ParallelInsertSliceOp parallelInsertOp,
                         SmallVector<ReassociationIndices> reassociations) {
  OpBuilder::InsertionGuard g(rewriter);
  // Compute the collapsed offsets, sizes, and strides.
  auto subsetOp =
      cast<SubsetInsertionOpInterface>(parallelInsertOp.getOperation());
  Operation *lastOp = setInsertionPointAfterLastNeededValue(rewriter, subsetOp);
  Location loc = lastOp->getLoc();
  int64_t resultIdx = parallelInsertOp.getTiedOpResult().getResultNumber();
  auto forallOp = parallelInsertOp->getParentOfType<scf::ForallOp>();
  Value loopInit = forallOp.getOutputs()[resultIdx];
  SmallVector<OpFoldResult> mixedInitSizes =
      tensor::getMixedSizes(rewriter, loc, loopInit);
  SmallVector<OpFoldResult> offsets = parallelInsertOp.getMixedOffsets();
  SmallVector<OpFoldResult> sizes = parallelInsertOp.getMixedSizes();
  SmallVector<OpFoldResult> newSizes, newOffsets;
  for (auto group : reassociations) {
    if (group.size() == 1) {
      newOffsets.push_back(offsets[group[0]]);
      newSizes.push_back(sizes[group[0]]);
      continue;
    }
    ArrayRef<OpFoldResult> basis(mixedInitSizes.begin() + group.front(),
                                 mixedInitSizes.begin() + group.back() + 1);
    ArrayRef<OpFoldResult> groupOffsets(offsets.begin() + group.front(),
                                        offsets.begin() + group.back() + 1);
    SmallVector<Value> offsetVals =
        llvm::map_to_vector(groupOffsets, [&](OpFoldResult ofr) {
          return getValueOrCreateConstantIndexOp(rewriter, loc, ofr);
        });
    OpFoldResult collapsedOffset =
        affine::AffineLinearizeIndexOp::create(rewriter, loc, offsetVals, basis,
                                               /*disjoint=*/true)
            .getResult();
    ArrayRef<OpFoldResult> groupSizes(sizes.begin() + group.front(),
                                      sizes.begin() + group.back() + 1);
    OpFoldResult collapsedSize =
        IREE::LinalgExt::computeProductUsingAffine(rewriter, loc, groupSizes);
    newOffsets.push_back(collapsedOffset);
    newSizes.push_back(collapsedSize);
  }
  SmallVector<OpFoldResult> newStrides(newSizes.size(),
                                       rewriter.getIndexAttr(1));

  // Collapse the slice source.
  loc = parallelInsertOp.getParallelCombiningParent()->getLoc();
  rewriter.setInsertionPoint(parallelInsertOp.getParallelCombiningParent());
  auto newCollapse = tensor::CollapseShapeOp::create(
      rewriter, loc, parallelInsertOp.getSource(), reassociations);

  // Collapse the parallel insert slice.
  rewriter.setInsertionPoint(parallelInsertOp);
  auto newInsertOp = rewriter.replaceOpWithNewOp<tensor::ParallelInsertSliceOp>(
      parallelInsertOp, newCollapse, parallelInsertOp.getDest(), newOffsets,
      newSizes, newStrides);
  return newInsertOp;
}

FailureOr<scf::ForallOp>
fuseCollapseShapeIntoProducerForall(RewriterBase &rewriter,
                                    scf::ForallOp forallOp,
                                    tensor::CollapseShapeOp collapseOp) {
  // Check that there is a single user of the collapsed result.
  auto forallResult = cast<OpResult>(collapseOp.getSrc());
  if (!forallResult.hasOneUse()) {
    return rewriter.notifyMatchFailure(forallOp,
                                       "forall result has multiple uses");
  }

  // Get the result's corresponding parallel_insert_slice op.
  SmallVector<Operation *> parallelInsertOps = forallOp.getCombiningOps(
      forallOp.getRegionIterArgs()[forallResult.getResultNumber()]);
  if (parallelInsertOps.size() != 1) {
    return rewriter.notifyMatchFailure(
        forallOp, "Expected a single parallel_insert_slice");
  }

  auto parallelInsertOp =
      dyn_cast<tensor::ParallelInsertSliceOp>(parallelInsertOps.front());
  if (!parallelInsertOp) {
    return rewriter.notifyMatchFailure(
        forallOp, "Expected parallel_insert_slice combining op");
  }

  // Collapse the parallel insert slice op.
  SmallVector<ReassociationIndices> reassociations =
      collapseOp.getReassociationIndices();
  if (failed(collapsibleSlicePrecondition(rewriter, parallelInsertOp,
                                          reassociations))) {
    return failure();
  }
  tensor::ParallelInsertSliceOp newParallelInsertOp =
      collapseParallelInsertOp(rewriter, parallelInsertOp, reassociations);

  // At this point, the newParallelInsertOp still has the destination of the
  // original parallel insert op, so the destination is the original expanded
  // init block argument, and we can use it to get the sizes for the expand.
  // The block argument will be corrected later, when the forall op is replaced.
  Value initArg = newParallelInsertOp.getDest();
  Value forallOutput = forallOp.getOutputs()[forallResult.getResultNumber()];
  Location loc = forallOutput.getLoc();
  rewriter.setInsertionPointAfterValue(forallOutput);
  SmallVector<OpFoldResult> initSizes =
      tensor::getMixedSizes(rewriter, loc, forallOutput);
  loc = initArg.getLoc();
  rewriter.setInsertionPointToStart(forallOp.getBody());
  auto expandedInitArg = tensor::ExpandShapeOp::create(
      rewriter, loc, initArg.getType(), initArg, reassociations, initSizes);

  // The new parallel insert slice is collapsed, so don't use the expanded init.
  // Also don't replace the expand shape src with its own result.
  rewriter.replaceUsesWithIf(
      initArg, expandedInitArg.getResult(), [&](OpOperand &operand) {
        return operand != expandedInitArg.getSrcMutable() &&
               operand != newParallelInsertOp.getDestMutable();
      });

  // Now create a new scf::Forall with a collapsed loop init.
  loc = forallOp->getLoc();
  rewriter.setInsertionPoint(forallOp);
  SmallVector<Value> newForallOutputs(forallOp.getOutputs());
  Value collapsedLoopInit = tensor::CollapseShapeOp::create(
      rewriter, loc, newForallOutputs[forallResult.getResultNumber()],
      reassociations);
  newForallOutputs[forallResult.getResultNumber()] = collapsedLoopInit;

  scf::ForallOp newForallOp = scf::ForallOp::create(
      rewriter, loc, forallOp.getMixedLowerBound(),
      forallOp.getMixedUpperBound(), forallOp.getMixedStep(), newForallOutputs,
      forallOp.getMappingAttr());

  SmallVector<Value> argReplacements(newForallOp.getInductionVars());
  argReplacements.append(newForallOp.getRegionIterArgs().begin(),
                         newForallOp.getRegionIterArgs().end());
  newForallOp.getTerminator()->erase();
  rewriter.mergeBlocks(forallOp.getBody(), newForallOp.getBody(),
                       argReplacements);

  // Replaces the uses of the old scf.forall with the new scf.forall
  rewriter.replaceOp(collapseOp,
                     newForallOp.getResult(forallResult.getResultNumber()));
  for (int idx = 0; idx < forallOp->getNumResults(); ++idx) {
    if (idx == forallResult.getResultNumber()) {
      continue;
    }
    forallOp->getResult(idx).replaceAllUsesWith(newForallOp->getResult(idx));
  }
  rewriter.eraseOp(forallOp);
  return newForallOp;
}

/// Return whether the `parallelInsertOp` can be clamped along the sliced
/// dimensions of `extractSliceOp`. The dimensions of the extractSliceOp source
/// are expected to match the dimensions of the parallelInsertOp destination.
/// This function checks that the parallelInsertOp is not rank reducing along
/// any of the sliced dimensions of the extractSliceOp.
static LogicalResult canClampParallelInsertSlice(
    RewriterBase &rewriter, tensor::ParallelInsertSliceOp parallelInsertOp,
    tensor::ExtractSliceOp extractSliceOp,
    llvm::SmallDenseSet<unsigned int> insertRankReductionMask) {
  // Find the dimensions that are sliced by the extractSliceOp
  llvm::SmallDenseSet<unsigned int> slicedDims;
  ArrayRef<int64_t> sliceStaticSizes = extractSliceOp.getStaticSizes();
  ArrayRef<int64_t> sliceSourceSizes =
      extractSliceOp.getSourceType().getShape();
  for (int dim = 0; dim < sliceStaticSizes.size(); ++dim) {
    if (ShapedType::isDynamic(sliceStaticSizes[dim]) ||
        sliceStaticSizes[dim] != sliceSourceSizes[dim]) {
      slicedDims.insert(dim);
    }
  }
  for (int dim = 0; dim < parallelInsertOp.getDestType().getRank(); ++dim) {
    if (insertRankReductionMask.contains(dim) && slicedDims.contains(dim)) {
      return rewriter.notifyMatchFailure(
          parallelInsertOp, "parallel insert reduces sliced dimensions");
    }
  }
  return success();
}

/// Clamps the source of a parallel_insert_slice op to fit within the
/// `upperBoundSizes`. This function computes the upper bound sizes, and creates
/// an extract slice op on the parallel insert source, which is then used in a
/// new parallel insert slice to replace the old one. This function assumes that
/// the parallel insert op passes `canClampParallelInsertSlice` precondition.
static FailureOr<tensor::ParallelInsertSliceOp>
clampParallelInsertSliceOp(RewriterBase &rewriter,
                           tensor::ParallelInsertSliceOp parallelInsertOp,
                           SmallVector<OpFoldResult> upperBoundSizes) {
  OpBuilder::InsertionGuard g(rewriter);
  auto subsetOp =
      cast<SubsetInsertionOpInterface>(parallelInsertOp.getOperation());
  Operation *lastOp = setInsertionPointAfterLastNeededValue(rewriter, subsetOp);
  Location loc = lastOp->getLoc();

  // Clamp the parallel_insert_slice sizes to fit within the full result tensor.
  SmallVector<OpFoldResult> offsets = parallelInsertOp.getMixedOffsets();
  SmallVector<OpFoldResult> sizes = parallelInsertOp.getMixedSizes();
  SmallVector<OpFoldResult> clampedSizes;
  for (auto [offset, size, ub] :
       llvm::zip_equal(offsets, sizes, upperBoundSizes)) {
    AffineExpr d0, d1, d2;
    MLIRContext *ctx = rewriter.getContext();
    bindDims(ctx, d0, d1, d2);
    auto lbClampMap = AffineMap::get(3, 0, {d0 - d1, d2}, ctx);
    auto ubClampMap = rewriter.getMultiDimIdentityMap(2);
    OpFoldResult lbClamped = affine::makeComposedFoldedAffineMax(
        rewriter, loc, lbClampMap, {ub, offset, rewriter.getIndexAttr(0)});
    OpFoldResult ubClamped = affine::makeComposedFoldedAffineMin(
        rewriter, loc, ubClampMap, {lbClamped, size});
    clampedSizes.push_back(ubClamped);
  }

  // Compute the clamped type. This could be rank reduced, but rank reduced
  // dimensions will never be potentially zero by construction. The earlier
  // matchers ensure that all sliceable users are not rank reduced along a
  // dimensions that is being sliced by the loop consumer.
  llvm::SmallDenseSet<unsigned int> rankReductionMask =
      computeRankReductionMask(parallelInsertOp.getStaticSizes(),
                               parallelInsertOp.getSourceType().getShape(),
                               /*matchDynamic=*/true)
          .value();
  SmallVector<int64_t> clampedShape;
  SmallVector<OpFoldResult> rankReducedClampedSizes;
  SmallVector<Value> d;
  for (auto [idx, clampedSize] : llvm::enumerate(clampedSizes)) {
    if (rankReductionMask.contains(idx)) {
      continue;
    }
    dispatchIndexOpFoldResult(clampedSize, d, clampedShape);
    rankReducedClampedSizes.push_back(clampedSize);
  }
  RankedTensorType clampedType =
      parallelInsertOp.getSourceType().clone(clampedShape);
  // Create an extract_slice to extract the correct size from the parallel
  // insert source.
  SmallVector<OpFoldResult> zeros(clampedType.getRank(),
                                  rewriter.getIndexAttr(0));
  SmallVector<OpFoldResult> ones(clampedType.getRank(),
                                 rewriter.getIndexAttr(1));
  Operation *combiningOp =
      parallelInsertOp.getParallelCombiningParent().getOperation();
  rewriter.setInsertionPoint(combiningOp);
  loc = combiningOp->getLoc();
  auto extractOp = tensor::ExtractSliceOp::create(
      rewriter, loc, clampedType, parallelInsertOp.getSource(), zeros,
      rankReducedClampedSizes, ones);

  // Replace the parallel insert op with the clamped version, and return the
  // new parallel insert slice.
  rewriter.setInsertionPoint(parallelInsertOp);
  loc = parallelInsertOp->getLoc();
  return rewriter.replaceOpWithNewOp<tensor::ParallelInsertSliceOp>(
      parallelInsertOp, extractOp.getResult(), parallelInsertOp.getDest(),
      parallelInsertOp.getMixedOffsets(), clampedSizes,
      parallelInsertOp.getMixedStrides());
}

FailureOr<scf::ForallOp>
fuseExtractSliceIntoProducerForall(RewriterBase &rewriter,
                                   scf::ForallOp forallOp,
                                   tensor::ExtractSliceOp extractSliceOp) {
  auto forallResult = cast<OpResult>(extractSliceOp.getSource());
  if (!forallResult.hasOneUse()) {
    return rewriter.notifyMatchFailure(forallOp,
                                       "forall result has multiple uses");
  }
  BlockArgument initBbarg =
      forallOp.getRegionIterArgs()[forallResult.getResultNumber()];
  SmallVector<Operation *> parallelInsertOps =
      forallOp.getCombiningOps(initBbarg);
  if (parallelInsertOps.size() != 1) {
    return rewriter.notifyMatchFailure(
        forallOp, "Expected a single parallel_insert_slice");
  }

  auto parallelInsertOp =
      dyn_cast<tensor::ParallelInsertSliceOp>(parallelInsertOps.front());
  if (!parallelInsertOp) {
    return rewriter.notifyMatchFailure(
        forallOp, "Expected parallel_insert_slice combining op");
  }

  // Only zero offset extract_slice ops are supported.
  if (!areAllConstantIntValue(extractSliceOp.getMixedOffsets(), 0)) {
    return rewriter.notifyMatchFailure(forallOp,
                                       "extract_slice has non-zero offsets");
  }

  // The extract_slice index operands must dominate the forall loop in order
  // to extract a slice of the init operand later.
  DominanceInfo domInfo;
  int64_t indexOperandStartIdx =
      extractSliceOp.getOffsetSizeAndStrideStartOperandIndex();
  SmallVector<Value> indexOperands(extractSliceOp->getOperands().begin() +
                                       indexOperandStartIdx,
                                   extractSliceOp->getOperands().end());
  if (!llvm::all_of(indexOperands,
                    [&](Value v) { return domInfo.dominates(v, forallOp); })) {
    return rewriter.notifyMatchFailure(
        extractSliceOp,
        "Extract slice index operands do not dominate the forall op");
  }

  // Compute the rank reduction mask of the extract_slice for resolving rank
  // reduction at the end. For rank reducing slices, the extract_slice is
  // fused into the loop as a non rank reducing slice, and then a collapse
  // shape is added on the result of the loop. This simplifies the logic in
  // this pattern, and other patterns for collapse shape fusion can then fuse
  // this collapse shape into the loop if needed.
  auto maybeRankReductionMask = computeRankReductionMask(
      extractSliceOp.getStaticSizes(), extractSliceOp.getType().getShape(),
      /*matchDynamic=*/true);
  if (!maybeRankReductionMask) {
    return rewriter.notifyMatchFailure(extractSliceOp,
                                       "Could not compute rank reduction mask");
  }

  std::optional<llvm::SmallDenseSet<unsigned int>>
      maybeInsertRankReductionMask =
          computeRankReductionMask(parallelInsertOp.getStaticSizes(),
                                   parallelInsertOp.getSourceType().getShape(),
                                   /*matchDynamic=*/true);
  if (!maybeInsertRankReductionMask) {
    return rewriter.notifyMatchFailure(parallelInsertOp,
                                       "Could not compute rank reduction mask");
  }
  llvm::SmallDenseSet<unsigned int> insertRankReductionMask =
      maybeInsertRankReductionMask.value();

  // Verify that the parallelInsertOp can be clamped to the sizes of the
  // extractSliceOp.
  if (failed(canClampParallelInsertSlice(rewriter, parallelInsertOp,
                                         extractSliceOp,
                                         insertRankReductionMask))) {
    return failure();
  }
  int64_t resultIdx = forallResult.getResultNumber();

  // Clamp the parallel insert slice source to fit within the extracted slice.
  SmallVector<OpFoldResult> newInitSizes = extractSliceOp.getMixedSizes();
  FailureOr<tensor::ParallelInsertSliceOp> maybeClampedParallelInsertSliceOp =
      clampParallelInsertSliceOp(rewriter, parallelInsertOp, newInitSizes);
  if (failed(maybeClampedParallelInsertSliceOp)) {
    return failure();
  }
  tensor::ParallelInsertSliceOp clampedParallelInsertSliceOp =
      maybeClampedParallelInsertSliceOp.value();

  // Now replace users of the forall loop init argument with the output operand
  // from outside the loop. Do not replace the clamped parallel insert dest.
  Value forallOutput = forallOp.getOutputs()[forallResult.getResultNumber()];
  rewriter.replaceUsesWithIf(initBbarg, forallOutput, [&](OpOperand &operand) {
    return operand != clampedParallelInsertSliceOp.getDestMutable();
  });

  // Clone the extract_slice, and replace the source with the forall init
  // operand.
  Value forallInit = forallOp.getOutputs()[resultIdx];
  rewriter.setInsertionPoint(forallOp);
  auto extractedInit = tensor::ExtractSliceOp::create(
      rewriter, forallOp->getLoc(), forallInit,
      extractSliceOp.getMixedOffsets(), extractSliceOp.getMixedSizes(),
      extractSliceOp.getMixedStrides());

  // Clone the forall op with the extracted init operand to replace the
  // original forall op.
  Location loc = forallOp->getLoc();
  rewriter.setInsertionPoint(forallOp);
  SmallVector<Value> newForallOutputs(forallOp.getOutputs());
  newForallOutputs[resultIdx] = extractedInit.getResult();

  scf::ForallOp newForallOp = scf::ForallOp::create(
      rewriter, loc, forallOp.getMixedLowerBound(),
      forallOp.getMixedUpperBound(), forallOp.getMixedStep(), newForallOutputs,
      forallOp.getMappingAttr());

  SmallVector<Value> argReplacements(newForallOp.getInductionVars());
  argReplacements.append(newForallOp.getRegionIterArgs().begin(),
                         newForallOp.getRegionIterArgs().end());
  newForallOp.getTerminator()->erase();
  rewriter.mergeBlocks(forallOp.getBody(), newForallOp.getBody(),
                       argReplacements);

  // Create a collapse_shape to handle rank reduction.
  Value extractedResult = newForallOp->getResult(resultIdx);
  auto forallResultType = cast<ShapedType>(extractedResult.getType());
  SmallVector<ReassociationIndices> reassociations;
  ReassociationIndices reassociation;
  for (int i = 0; i < forallResultType.getRank(); ++i) {
    if (maybeRankReductionMask->contains(i)) {
      reassociation.push_back(i);
      continue;
    }
    reassociation.push_back(i);
    reassociations.push_back(reassociation);
    reassociation = {};
  }
  auto collapseShape = tensor::CollapseShapeOp::create(
      rewriter, extractSliceOp->getLoc(), extractedResult, reassociations);

  // Replace forall and extract_slice ops with the new operations.
  rewriter.replaceAllOpUsesWith(extractSliceOp, collapseShape);
  rewriter.replaceOp(forallOp, newForallOp);
  return newForallOp;
}

//===----------------------------------------------------------------------===//
// Conversion to MultiMmaOp
//===----------------------------------------------------------------------===//

static AffineMap
dropDims(MLIRContext *context, int64_t newDimCount, AffineMap map,
         llvm::SmallDenseMap<int64_t, int64_t> &oldDimsToNewDimsMap) {
  assert(map.isProjectedPermutation() && "expected projected permutation");

  SmallVector<AffineExpr> newResults;
  for (auto expr : map.getResults()) {
    int64_t dimPos = cast<AffineDimExpr>(expr).getPosition();
    if (!oldDimsToNewDimsMap.contains(dimPos)) {
      continue;
    }
    newResults.push_back(
        getAffineDimExpr(oldDimsToNewDimsMap[dimPos], context));
  }
  return AffineMap::get(/*dimCount=*/newDimCount, /*symbolCount=*/0, newResults,
                        context);
}

static FailureOr<IREE::Codegen::InnerTiledOp>
convertScaledContractionToInnerTiledMma(
    RewriterBase &rewriter, linalg::LinalgOp linalgOp,
    IREE::Codegen::InnerTileDescAttrInterface kind) {
  auto smmaKind = dyn_cast<IREE::GPU::ScaledMMAAttr>(kind);
  FailureOr<IREE::LinalgExt::ScaledContractionDimensions> contractionDims =
      IREE::LinalgExt::inferScaledContractionDims(linalgOp);
  if (failed(contractionDims)) {
    return failure();
  }

  if (contractionDims->m.empty() || contractionDims->n.empty() ||
      contractionDims->k.empty() || contractionDims->kB.empty()) {
    return failure();
  }

  MLIRContext *context = rewriter.getContext();

  int64_t innerM = contractionDims->m.back();
  int64_t innerN = contractionDims->n.back();
  int64_t innerK = contractionDims->k.back();
  int64_t innerKb = contractionDims->kB.back();

  AffineExpr mExpr = rewriter.getAffineDimExpr(innerM);
  AffineExpr nExpr = rewriter.getAffineDimExpr(innerN);
  AffineExpr kExpr = rewriter.getAffineDimExpr(innerK);
  AffineExpr kBExpr = rewriter.getAffineDimExpr(innerKb);

  SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();
  AffineMap lhsMap = indexingMaps[0];
  AffineMap rhsMap = indexingMaps[1];
  AffineMap sc1Map = indexingMaps[2];
  AffineMap sc2Map = indexingMaps[3];
  AffineMap accMap = indexingMaps[4];

  auto getNormalizedPermutation =
      [&](AffineMap map,
          ArrayRef<AffineExpr> expectedDimOrder) -> SmallVector<int64_t> {
    llvm::SmallDenseMap<AffineExpr, int64_t> dimMap;
    for (auto [i, expr] : llvm::enumerate(expectedDimOrder)) {
      dimMap[expr] = i;
    }
    SmallVector<int64_t> permutation;
    for (AffineExpr resExpr : map.getResults()) {
      if (!dimMap.contains(resExpr)) {
        return {};
      }
      permutation.push_back(dimMap[resExpr]);
    }
    return permutation;
  };

  // TODO: Enable batched intrinsics and get the appropriate sub-map here.
  SmallVector<int64_t> lhsInnerPerm = getNormalizedPermutation(
      lhsMap.getMinorSubMap(3), {mExpr, kExpr, kBExpr});
  SmallVector<int64_t> sc1InnerPerm =
      getNormalizedPermutation(sc1Map.getMinorSubMap(2), {mExpr, kExpr});
  SmallVector<int64_t> rhsInnerPerm = getNormalizedPermutation(
      rhsMap.getMinorSubMap(3), {kExpr, kBExpr, nExpr});
  SmallVector<int64_t> sc2InnerPerm =
      getNormalizedPermutation(sc2Map.getMinorSubMap(2), {kExpr, nExpr});
  SmallVector<int64_t> accInnerPerm =
      getNormalizedPermutation(accMap.getMinorSubMap(2), {mExpr, nExpr});
  if (lhsInnerPerm.empty() || sc1InnerPerm.empty() || rhsInnerPerm.empty() ||
      sc2InnerPerm.empty() || accInnerPerm.empty()) {
    return failure();
  }

  SmallVector<int64_t> bounds = linalgOp.getStaticLoopRanges();
  auto [intrinsicM, intrinsicN, intrinsicK, intrinsicKB] =
      smmaKind.getScaledMNKShape();
  if (intrinsicM != bounds[innerM] || intrinsicN != bounds[innerN] ||
      intrinsicK != bounds[innerK]) {
    return failure();
  }

  ValueRange inputs = linalgOp->getOperands();

  SmallVector<Type> eltTypes;
  smmaKind.getElementTypes(eltTypes);
  for (int i :
       {kScaledMMAOperandLhs, kScaledMMAOperandRhs, kScaledMMAOperandAcc}) {
    if (cast<RankedTensorType>(inputs[i].getType()).getElementType() !=
        eltTypes[i]) {
      return failure();
    }
  }

  SmallVector<utils::IteratorType> linalgIteratorTypes =
      linalgOp.getIteratorTypesArray();
  llvm::SmallDenseSet<int64_t> droppedDims = {innerM, innerN, innerK, innerKb};
  llvm::SmallDenseMap<int64_t, int64_t> oldDimsToNewDimsMap;
  int64_t currentDim = 0;
  int64_t numDims = lhsMap.getNumDims();
  SmallVector<utils::IteratorType> iteratorTypes;
  for (int64_t dim = 0, e = numDims; dim < e; ++dim) {
    if (droppedDims.contains(dim)) {
      continue;
    }
    iteratorTypes.push_back(linalgIteratorTypes[dim]);
    oldDimsToNewDimsMap[dim] = currentDim++;
  }
  AffineMap outerLhsMap =
      dropDims(context, numDims - 4, lhsMap, oldDimsToNewDimsMap);
  AffineMap outerRhsMap =
      dropDims(context, numDims - 4, rhsMap, oldDimsToNewDimsMap);
  AffineMap outerSc1Map =
      dropDims(context, numDims - 4, sc1Map, oldDimsToNewDimsMap);
  AffineMap outerSc2Map =
      dropDims(context, numDims - 4, sc2Map, oldDimsToNewDimsMap);
  AffineMap outerAccMap =
      dropDims(context, numDims - 4, accMap, oldDimsToNewDimsMap);
  std::optional<SmallVector<SmallVector<int64_t>>> perms =
      SmallVector<SmallVector<int64_t>>{
          lhsInnerPerm, rhsInnerPerm, sc1InnerPerm, sc2InnerPerm, accInnerPerm};
  SmallVector<int64_t> identityPerm = {0, 1};
  if (lhsInnerPerm == identityPerm && rhsInnerPerm == identityPerm &&
      accInnerPerm == identityPerm) {
    perms = std::nullopt;
  }

  IREE::Codegen::LoweringConfigAttrInterface maybeLoweringConfig =
      getLoweringConfig(linalgOp);
  auto semantics = InnerTiledSemanticsAttr::get(context, /*distributed=*/false,
                                                /*opaque=*/true);
  auto newMmaOp = rewriter.replaceOpWithNewOp<IREE::Codegen::InnerTiledOp>(
      linalgOp, /*inputs=*/ValueRange{inputs}.drop_back(),
      /*inits=*/ValueRange{inputs}.back(),
      ArrayRef<AffineMap>{outerLhsMap, outerRhsMap, outerSc1Map, outerSc2Map,
                          outerAccMap},
      iteratorTypes, smmaKind, semantics, perms);
  if (maybeLoweringConfig) {
    setLoweringConfig(newMmaOp, maybeLoweringConfig);
  }
  return newMmaOp;
}

// Helper to convert a contraction-like linalg op to an iree_codegen.inner_tiled
// op with a MMA-like intrinsic descriptor.
FailureOr<IREE::Codegen::InnerTiledOp> convertContractionToInnerTiledMma(
    RewriterBase &rewriter, linalg::LinalgOp linalgOp,
    IREE::Codegen::InnerTileDescAttrInterface kind) {
  if (!linalgOp.hasPureTensorSemantics()) {
    return failure();
  }

  FailureOr<IREE::LinalgExt::ScaledContractionDimensions> maybeScaledContrDims =
      IREE::LinalgExt::inferScaledContractionDims(linalgOp);
  if (succeeded(maybeScaledContrDims)) {
    return convertScaledContractionToInnerTiledMma(rewriter, linalgOp, kind);
  }

  IREE::GPU::MmaInterfaceAttr mmaKind =
      dyn_cast<IREE::GPU::MmaInterfaceAttr>(kind);
  FailureOr<linalg::ContractionDimensions> maybeContractionDims =
      linalg::inferContractionDims(linalgOp);
  if (failed(maybeContractionDims)) {
    return failure();
  }

  linalg::ContractionDimensions contractionDims = *maybeContractionDims;
  if (contractionDims.m.empty() || contractionDims.n.empty() ||
      contractionDims.k.empty()) {
    return failure();
  }

  bool isBlock = mmaKind.isBlockIntrinsic();
  if (isBlock && contractionDims.batch.empty()) {
    return failure();
  }

  MLIRContext *context = rewriter.getContext();

  int64_t innerM = contractionDims.m.back();
  int64_t innerN = contractionDims.n.back();
  int64_t innerK = contractionDims.k.back();

  AffineExpr mExpr = rewriter.getAffineDimExpr(innerM);
  AffineExpr nExpr = rewriter.getAffineDimExpr(innerN);
  AffineExpr kExpr = rewriter.getAffineDimExpr(innerK);

  SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();
  AffineMap lhsMap = indexingMaps[0];
  AffineMap rhsMap = indexingMaps[1];
  AffineMap accMap = indexingMaps[2];

  auto getNormalizedPermutation =
      [&](AffineMap map,
          ArrayRef<AffineExpr> expectedDimOrder) -> SmallVector<int64_t> {
    llvm::SmallDenseMap<AffineExpr, int64_t> dimMap;
    for (auto [i, expr] : llvm::enumerate(expectedDimOrder)) {
      dimMap[expr] = i;
    }
    SmallVector<int64_t> permutation;
    for (AffineExpr resExpr : map.getResults()) {
      if (!dimMap.contains(resExpr)) {
        return {};
      }
      permutation.push_back(dimMap[resExpr]);
    }
    return permutation;
  };

  SmallVector<int64_t> lhsInnerPerm, rhsInnerPerm, accInnerPerm;
  SmallVector<int64_t> bounds = linalgOp.getStaticLoopRanges();
  int64_t numDims = lhsMap.getNumDims();
  llvm::SmallDenseSet<int64_t> droppedDims;
  int64_t numInnerDims;

  if (isBlock) {
    int64_t innerB = contractionDims.batch.back();
    AffineExpr bExpr = rewriter.getAffineDimExpr(innerB);
    lhsInnerPerm = getNormalizedPermutation(lhsMap.getMinorSubMap(3),
                                            {bExpr, mExpr, kExpr});
    rhsInnerPerm = getNormalizedPermutation(rhsMap.getMinorSubMap(3),
                                            {bExpr, kExpr, nExpr});
    accInnerPerm = getNormalizedPermutation(accMap.getMinorSubMap(3),
                                            {bExpr, mExpr, nExpr});
    if (lhsInnerPerm.empty() || rhsInnerPerm.empty() || accInnerPerm.empty()) {
      return failure();
    }
    auto [intrinsicB, intrinsicM, intrinsicN, intrinsicK] =
        mmaKind.getBMNKShape();
    if (intrinsicB != bounds[innerB] || intrinsicM != bounds[innerM] ||
        intrinsicN != bounds[innerN] || intrinsicK != bounds[innerK]) {
      return failure();
    }
    droppedDims = {innerB, innerM, innerN, innerK};
    numInnerDims = 4;
  } else {
    lhsInnerPerm =
        getNormalizedPermutation(lhsMap.getMinorSubMap(2), {mExpr, kExpr});
    rhsInnerPerm =
        getNormalizedPermutation(rhsMap.getMinorSubMap(2), {kExpr, nExpr});
    accInnerPerm =
        getNormalizedPermutation(accMap.getMinorSubMap(2), {mExpr, nExpr});
    if (lhsInnerPerm.empty() || rhsInnerPerm.empty() || accInnerPerm.empty()) {
      return failure();
    }
    auto [intrinsicM, intrinsicN, intrinsicK] = mmaKind.getMNKShape();
    if (intrinsicM != bounds[innerM] || intrinsicN != bounds[innerN] ||
        intrinsicK != bounds[innerK]) {
      return failure();
    }
    droppedDims = {innerM, innerN, innerK};
    numInnerDims = 3;
  }

  SmallVector<Value> inputs = linalgOp->getOperands();
  auto [lhsElementType, rhsElementType, accElementType] =
      mmaKind.getABCElementTypes();
  if (cast<RankedTensorType>(inputs[0].getType()).getElementType() !=
          lhsElementType ||
      cast<RankedTensorType>(inputs[1].getType()).getElementType() !=
          rhsElementType ||
      cast<RankedTensorType>(inputs[2].getType()).getElementType() !=
          accElementType) {
    return failure();
  }

  SmallVector<utils::IteratorType> linalgIteratorTypes =
      linalgOp.getIteratorTypesArray();
  llvm::SmallDenseMap<int64_t, int64_t> oldDimsToNewDimsMap;
  int64_t currentDim = 0;
  SmallVector<utils::IteratorType> iteratorTypes;
  for (int64_t dim = 0, e = numDims; dim < e; ++dim) {
    if (droppedDims.contains(dim)) {
      continue;
    }
    iteratorTypes.push_back(linalgIteratorTypes[dim]);
    oldDimsToNewDimsMap[dim] = currentDim++;
  }

  AffineMap outerLhsMap =
      dropDims(context, numDims - numInnerDims, lhsMap, oldDimsToNewDimsMap);
  AffineMap outerRhsMap =
      dropDims(context, numDims - numInnerDims, rhsMap, oldDimsToNewDimsMap);
  AffineMap outerAccMap =
      dropDims(context, numDims - numInnerDims, accMap, oldDimsToNewDimsMap);

  std::optional<SmallVector<SmallVector<int64_t>>> perms =
      SmallVector<SmallVector<int64_t>>{lhsInnerPerm, rhsInnerPerm,
                                        accInnerPerm};
  SmallVector<int64_t> identityPerm =
      isBlock ? SmallVector<int64_t>{0, 1, 2} : SmallVector<int64_t>{0, 1};

  if (lhsInnerPerm == identityPerm && rhsInnerPerm == identityPerm &&
      accInnerPerm == identityPerm) {
    perms = std::nullopt;
  }

  IREE::Codegen::LoweringConfigAttrInterface maybeLoweringConfig =
      getLoweringConfig(linalgOp);

  auto semantics = InnerTiledSemanticsAttr::get(context, /*distributed=*/false,
                                                /*opaque=*/true);
  auto newMmaOp = rewriter.replaceOpWithNewOp<IREE::Codegen::InnerTiledOp>(
      linalgOp, /*inputs=*/ValueRange{inputs}.drop_back(),
      /*inits=*/ValueRange{inputs}.back(),
      ArrayRef<AffineMap>{outerLhsMap, outerRhsMap, outerAccMap}, iteratorTypes,
      mmaKind, semantics, perms);
  if (maybeLoweringConfig) {
    setLoweringConfig(newMmaOp, maybeLoweringConfig);
  }
  return newMmaOp;
}

//===----------------------------------------------------------------------===//
// InnerTiledOp Distribution
//===----------------------------------------------------------------------===//

FailureOr<Operation *>
distributeInnerTiledOp(RewriterBase &rewriter,
                       IREE::Codegen::InnerTiledOp tiledOp) {
  auto semantics =
      dyn_cast<IREE::GPU::InnerTiledSemanticsAttr>(tiledOp.getSemantics());
  if (!semantics) {
    return rewriter.notifyMatchFailure(
        tiledOp, "unexpected tiled op semantics attribute type");
  }
  if (semantics.getDistributed()) {
    return rewriter.notifyMatchFailure(tiledOp, "tiledOp already distributed.");
  }
  if (!tiledOp.hasTensorSemantics()) {
    return rewriter.notifyMatchFailure(
        tiledOp, "tiledOp must have vector semantics for distribution.");
  }

  RewriterBase::InsertionGuard g(rewriter);

  Location loc = tiledOp.getLoc();
  MLIRContext *context = rewriter.getContext();

  OpFoldResult zero = rewriter.getIndexAttr(0);
  OpFoldResult one = rewriter.getIndexAttr(1);

  // Step 1. Create the new scf.forall op with a lane id mapping.
  Attribute mappingType = tiledOp.getKind().getDistributionMappingKind();
  if (!mappingType) {
    return rewriter.notifyMatchFailure(
        tiledOp, "doesn't specify how it's to be distributed");
  }
  OpFoldResult ub =
      tiledOp.getKind().getDistributionWorkerCount(rewriter, loc, tiledOp);
  if (!ub) {
    return tiledOp.emitOpError("failed to specify a worker count for the "
                               "forall it's to be distributed into.");
  }

  auto newForallOp = scf::ForallOp::create(
      rewriter, loc, ArrayRef<OpFoldResult>{zero}, ArrayRef<OpFoldResult>{ub},
      ArrayRef<OpFoldResult>{one}, tiledOp.getOutputs(),
      ArrayAttr::get(context, {mappingType}));

  rewriter.setInsertionPointToStart(newForallOp.getBody());

  // Step 2. Compute the offsets/sizes/strides for each of the operands.
  Value id = newForallOp.getInductionVar(0);
  SmallVector<Value> inputSlices, initSlices;
  SmallVector<tensor::ExtractSliceOp> initSliceOps;
  std::optional<ArrayAttr> maybePerms = tiledOp.getPermutations();
  int64_t firstOutIdx = tiledOp.getNumInputs();

  for (auto [opIndex, operand] : llvm::enumerate(tiledOp.getOperands())) {
    int64_t outerRank = tiledOp.getOperandOuterRank(opIndex);
    SmallVector<OpFoldResult> offsets(outerRank, zero);
    SmallVector<OpFoldResult> sizes;
    for (int64_t i = 0, e = outerRank; i < e; ++i) {
      sizes.push_back(tensor::getMixedSize(rewriter, loc, operand, i));
    }
    ArrayRef<int64_t> innerShape = tiledOp.getOperandInnerShape(opIndex);
    SmallVector<int64_t> permutation;
    if (maybePerms) {
      permutation = llvm::to_vector(
          cast<DenseI64ArrayAttr>((*maybePerms)[opIndex]).asArrayRef());
    } else {
      permutation = llvm::to_vector(llvm::seq(
          static_cast<int64_t>(0), static_cast<int64_t>(innerShape.size())));
    }
    // Slice offsets.
    SmallVector<OpFoldResult> strides(outerRank, one);
    if (failed(tiledOp.getKind().populateOperandOffsetsSizesStrides(
            rewriter, loc, opIndex, id, permutation, offsets, sizes,
            strides))) {
      return tiledOp->emitOpError("failed to populate offsets for operand " +
                                  Twine(opIndex));
    }
    // Extract the rank-reduced slice of the operand based on the expected inner
    // vector shape. If we're slicing the accumulator, extract from the loop
    // variable, since the accumulator operand has been used to create the
    // forall.
    if (opIndex >= firstOutIdx) {
      auto sliceOp = tensor::ExtractSliceOp::create(
          rewriter, loc, newForallOp.getRegionIterArgs()[opIndex - firstOutIdx],
          offsets, sizes, strides);
      initSliceOps.push_back(sliceOp);
      initSlices.push_back(sliceOp);
    } else {
      Value slice = tensor::ExtractSliceOp::create(rewriter, loc, operand,
                                                   offsets, sizes, strides);
      inputSlices.push_back(slice);
    }
  }

  auto distributedSemantics = IREE::GPU::InnerTiledSemanticsAttr::get(
      context, /*distributed=*/true, semantics.getOpaque());

  // Step 3. Create the new inner_tiled op.
  auto newTiledOp = IREE::Codegen::InnerTiledOp::create(
      rewriter, loc, inputSlices, initSlices, tiledOp.getIndexingMaps(),
      tiledOp.getIteratorTypes(), tiledOp.getKind(), distributedSemantics);

  newTiledOp->setDiscardableAttrs(tiledOp->getDiscardableAttrDictionary());

  // Step 4. Insert the result of the inner_tiled using the same offsets/sizes
  // as the accumulator slice.
  scf::InParallelOp terminator = newForallOp.getTerminator();
  rewriter.setInsertionPointToStart(terminator.getBody());
  for (auto [newResult, extractOp] :
       llvm::zip_equal(newTiledOp.getResults(), initSliceOps)) {
    tensor::ParallelInsertSliceOp::create(
        rewriter, loc, newResult, extractOp.getSource(),
        extractOp.getMixedOffsets(), extractOp.getMixedSizes(),
        extractOp.getMixedStrides());
  }

  rewriter.replaceOp(tiledOp, newForallOp);

  return &*newForallOp;
}

//===---------------------------------------------------------------------===//
// Resolving lane mapped forall ops
//===---------------------------------------------------------------------===//

static bool isLaneMappableForall(scf::ForallOp forallOp) {
  if (forallOp.getNumResults() > 0) {
    return false;
  }
  if (forallOp.getRank() != 1) {
    return false;
  }
  if (!forallOp.getMapping().has_value()) {
    return false;
  }
  Attribute mapping = *forallOp.getMapping()->getValue().begin();
  if (mapping != IREE::GPU::LaneIdAttr::get(forallOp.getContext(), 0)) {
    return false;
  }
  return true;
}

static void rewriteForallToLanes(RewriterBase &rewriter, scf::ForallOp forallOp,
                                 bool insertBarrier) {
  Location loc = forallOp->getLoc();
  assert(isLaneMappableForall(forallOp) && "mapping non-lane forall op");

  auto upperBounds = forallOp.getLoopUpperBounds();
  std::optional<IntegerAttr> upperBound;
  if (upperBounds && upperBounds->size() > 0) {
    if (auto upperBoundAttr = (*upperBounds)[0].dyn_cast<Attribute>()) {
      upperBound = dyn_cast<IntegerAttr>(upperBoundAttr);
    }
  }
  Value laneId = gpu::LaneIdOp::create(
      rewriter, loc,
      upperBound ? rewriter.getIndexAttr(upperBound->getInt()) : nullptr);
  rewriter.eraseOp(forallOp.getTerminator());
  rewriter.setInsertionPoint(forallOp);
  rewriter.inlineBlockBefore(forallOp.getBody(), forallOp, {laneId});
  if (insertBarrier) {
    gpu::BarrierOp::create(rewriter, loc, gpu::AddressSpace::Workgroup);
  }
  rewriter.eraseOp(forallOp);
}

void mapLaneForalls(RewriterBase &rewriter, Operation *funcOp,
                    bool insertBarrier) {
  SmallVector<scf::ForallOp> foralls;
  OpBuilder::InsertionGuard g(rewriter);
  funcOp->walk([&](scf::ForallOp forallOp) {
    if (isLaneMappableForall(forallOp)) {
      foralls.push_back(forallOp);
    }
  });
  for (auto forall : foralls) {
    rewriter.setInsertionPoint(forall);
    rewriteForallToLanes(rewriter, forall, insertBarrier);
  }
}

//===---------------------------------------------------------------------===//
// BarrierRegion Lowering
//===---------------------------------------------------------------------===//

namespace {
struct LowerBarrierRegion : OpRewritePattern<IREE::GPU::BarrierRegionOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(IREE::GPU::BarrierRegionOp barrierRegionOp,
                                PatternRewriter &rewriter) const final {
    Location loc = barrierRegionOp.getLoc();

    // Step 1. Synchronize the workers on the shared dest.
    auto writeBarrier = IREE::GPU::ValueBarrierOp::create(
        rewriter, loc, barrierRegionOp.getInputs());

    // Step 2. Inline the barrier op region.
    auto terminator = barrierRegionOp.getBody()->getTerminator();
    rewriter.inlineBlockBefore(barrierRegionOp.getBody(), barrierRegionOp,
                               writeBarrier.getResults());
    rewriter.setInsertionPoint(terminator);

    // Step 3. Synchronize the result value.
    auto barrier = IREE::GPU::ValueBarrierOp::create(rewriter, loc,
                                                     terminator->getOperands());
    rewriter.replaceAllUsesWith(barrierRegionOp.getResults(),
                                barrier.getResults());
    rewriter.eraseOp(terminator);
    return success();
  }
};
} // namespace

void populateIREEGPULowerBarrierRegionPatterns(RewritePatternSet &patterns) {
  patterns.add<LowerBarrierRegion>(patterns.getContext());
}

//===----------------------------------------------------------------------===//
// VectorBarrierOp Lowering
//===----------------------------------------------------------------------===//

namespace {
struct LowerValueBarrierPattern : OpRewritePattern<IREE::GPU::ValueBarrierOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(IREE::GPU::ValueBarrierOp barrier,
                                PatternRewriter &rewriter) const override {
    if (barrier.hasTensorSemantics()) {
      return failure();
    }
    // TODO(kdrewnia): We give this barrier workgroup fence semantics to mach
    // the behavior it had from before we could be more precise about barrier
    // address spaces. Investigate if this could be a sync-only (`memfence []`)
    // barrier
    gpu::BarrierOp::create(rewriter, barrier.getLoc(),
                           gpu::AddressSpace::Workgroup);
    for (auto [result, input] :
         llvm::zip_equal(barrier.getResults(), barrier.getInputs())) {
      rewriter.replaceAllUsesWith(result, input);
    }
    return success();
  }
};
} // namespace

void populateIREEGPULowerValueBarrierPatterns(RewritePatternSet &patterns) {
  patterns.add<LowerValueBarrierPattern>(patterns.getContext());
}

//===----------------------------------------------------------------------===//
// SwizzleHintOp Fold Patterns
//===----------------------------------------------------------------------===//

// The following patterns are adapted from the populateFoldTensorEmptyPatterns
// in upstream llvm-project. The main change is to add support for folding with
// swizzle_hint ops from IREE. Once swizzle_hint ops are more widely used and
// proven stable, we could consider upstreaming this extension.

namespace {
struct FoldSwizzleHintOpWithExtractSliceOp final
    : OpRewritePattern<tensor::ExtractSliceOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
                                PatternRewriter &rewriter) const override {
    // Check for swizzle_hint op source.
    auto swizzleHintOp =
        sliceOp.getSource().getDefiningOp<IREE::Codegen::SwizzleHintOp>();
    if (!swizzleHintOp) {
      return failure();
    }

    // Check for tensor.empty source.
    auto emptyOp = swizzleHintOp.getOperand().getDefiningOp<tensor::EmptyOp>();
    if (!emptyOp) {
      return failure();
    }

    // Check for single use.
    if (!emptyOp->hasOneUse()) {
      return failure();
    }

    // Create new tensor.empty op. tensor.extract_slice may be rank-reducing;
    // its dynamic sizes must be preserved as well as its result type.
    Location loc = sliceOp.getLoc();
    auto sliceType = cast<RankedTensorType>(sliceOp.getType());
    auto tensorType =
        RankedTensorType::get(sliceType.getShape(), sliceType.getElementType(),
                              sliceType.getEncoding());
    auto newEmptyOp =
        tensor::EmptyOp::create(rewriter, loc, tensorType, sliceOp.getSizes());
    rewriter.replaceOpWithNewOp<IREE::Codegen::SwizzleHintOp>(
        sliceOp, newEmptyOp, swizzleHintOp.getSwizzle());
    return success();
  }
};

template <typename ReshapeOp>
struct FoldSwizzleHintOpWithReshapeOp final : OpRewritePattern<ReshapeOp> {
  using OpRewritePattern<ReshapeOp>::OpRewritePattern;
  LogicalResult matchAndRewrite(ReshapeOp reshapeOp,
                                PatternRewriter &rewriter) const override {
    auto swizzleHintOp =
        reshapeOp.getSrc()
            .template getDefiningOp<IREE::Codegen::SwizzleHintOp>();
    if (!swizzleHintOp) {
      return failure();
    }
    auto emptyOp =
        swizzleHintOp.getOperand().template getDefiningOp<tensor::EmptyOp>();
    if (!emptyOp) {
      return failure();
    }

    // Check for single use.
    if (!emptyOp->hasOneUse()) {
      return failure();
    }

    // Reify result shape.
    Location loc = reshapeOp.getLoc();
    ReifiedRankedShapedTypeDims resultShapes;
    if (failed(reifyResultShapes(rewriter, reshapeOp, resultShapes)) ||
        !llvm::hasSingleElement(resultShapes)) {
      return failure();
    }

    // Create new tensor.empty op.
    Value emptyTensor =
        tensor::EmptyOp::create(rewriter, loc, resultShapes[0],
                                reshapeOp.getResultType().getElementType(),
                                reshapeOp.getResultType().getEncoding());
    Value newSwizzleHintOp = IREE::Codegen::SwizzleHintOp::create(
        rewriter, loc, emptyTensor, swizzleHintOp.getSwizzle());
    if (newSwizzleHintOp.getType() != reshapeOp.getResultType()) {
      rewriter.replaceOpWithNewOp<tensor::CastOp>(
          reshapeOp, reshapeOp.getResultType(), newSwizzleHintOp);
    } else {
      rewriter.replaceOp(reshapeOp, newSwizzleHintOp);
    }
    return success();
  }
};

} // namespace

void populateFoldSwizzleHintOpPatterns(RewritePatternSet &patterns) {
  patterns.add<FoldSwizzleHintOpWithReshapeOp<tensor::ExpandShapeOp>,
               FoldSwizzleHintOpWithReshapeOp<tensor::CollapseShapeOp>,
               FoldSwizzleHintOpWithExtractSliceOp>(patterns.getContext());
}

} // namespace mlir::iree_compiler::IREE::GPU
