// Copyright 2025 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/PCF/IR/PCF.h"
#include "iree/compiler/Codegen/Dialect/PCF/IR/PCFOps.h"
#include "iree/compiler/Codegen/Dialect/PCF/Transforms/Passes.h"
#include "iree/compiler/Codegen/Dialect/PCF/Transforms/Transforms.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"
#include "mlir/Interfaces/TilingInterface.h"
#include "mlir/Support/WalkResult.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

#define DEBUG_TYPE "iree-pcf-fuse-consumers"

namespace mlir::iree_compiler::IREE::PCF {

#define GEN_PASS_DEF_FUSECONSUMERSPASS
#include "iree/compiler/Codegen/Dialect/PCF/Transforms/Passes.h.inc"

namespace {

struct FuseConsumersPass final
    : impl::FuseConsumersPassBase<FuseConsumersPass> {
  void runOnOperation() override;
};

struct FuseIntoGenericOp final : OpRewritePattern<IREE::PCF::GenericOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(IREE::PCF::GenericOp genericOp,
                                PatternRewriter &rewriter) const override {
    ConsumerFusionParams params;
    TilingInterface fusionTarget;
    for (Operation *user : genericOp->getUsers()) {
      fusionTarget = dyn_cast<TilingInterface>(user);
      ConsumerFusionParams tempParams;
      if (fusionTarget && succeeded(matchTilableConsumer(
                              rewriter, genericOp, fusionTarget, tempParams))) {
        std::swap(params, tempParams);
        break;
      }
      fusionTarget = TilingInterface();
    }
    if (!fusionTarget) {
      // Match failure reason message generated by `matchTilableConsumer`.
      return failure();
    }
    fuseTilableConsumer(rewriter, genericOp, fusionTarget, params);
    return success();
  }
};

struct FuseIntoLoopOp final : OpRewritePattern<IREE::PCF::LoopOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(IREE::PCF::LoopOp loopOp,
                                PatternRewriter &rewriter) const override {
    ConsumerFusionParams params;
    TilingInterface fusionTarget;
    for (Operation *user : loopOp->getUsers()) {
      fusionTarget = dyn_cast<TilingInterface>(user);
      ConsumerFusionParams tempParams;
      if (fusionTarget && succeeded(matchTilableConsumer(
                              rewriter, loopOp, fusionTarget, tempParams))) {
        std::swap(params, tempParams);
        break;
      }
      fusionTarget = TilingInterface();
    }
    if (!fusionTarget) {
      // Match failure reason message generated by `matchTilableConsumer`.
      return failure();
    }
    fuseTilableConsumer(rewriter, loopOp, fusionTarget, params);
    return success();
  }
};

struct FuseExtractSliceIntoLoopOp final
    : OpRewritePattern<tensor::ExtractSliceOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(tensor::ExtractSliceOp extractSliceOp,
                                PatternRewriter &rewriter) const override {
    auto loopOp = extractSliceOp.getSource().getDefiningOp<IREE::PCF::LoopOp>();
    if (!loopOp) {
      return rewriter.notifyMatchFailure(extractSliceOp, "No loop op producer");
    }

    return fuseExtractSliceIntoProducerLoop(rewriter, loopOp, extractSliceOp);
  }
};

struct FuseExtractSliceIntoGenericOp final
    : OpRewritePattern<tensor::ExtractSliceOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(tensor::ExtractSliceOp extractSliceOp,
                                PatternRewriter &rewriter) const override {
    auto genericOp =
        extractSliceOp.getSource().getDefiningOp<IREE::PCF::GenericOp>();
    if (!genericOp) {
      return rewriter.notifyMatchFailure(extractSliceOp,
                                         "No generic op producer");
    }

    return fuseExtractSliceIntoProducerGeneric(rewriter, genericOp,
                                               extractSliceOp);
  }
};

struct FuseCollapseShapeIntoGenericOp final
    : OpRewritePattern<tensor::CollapseShapeOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(tensor::CollapseShapeOp collapseOp,
                                PatternRewriter &rewriter) const override {
    auto genericOp = collapseOp.getSrc().getDefiningOp<IREE::PCF::GenericOp>();
    if (!genericOp) {
      return rewriter.notifyMatchFailure(collapseOp, "No generic op producer");
    }
    return fuseCollapseShapeIntoProducerGeneric(rewriter, genericOp,
                                                collapseOp);
  }
};

struct FuseCollapseShapeIntoLoopOp final
    : OpRewritePattern<tensor::CollapseShapeOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(tensor::CollapseShapeOp collapseOp,
                                PatternRewriter &rewriter) const override {
    auto loopOp = collapseOp.getSrc().getDefiningOp<IREE::PCF::LoopOp>();
    if (!loopOp) {
      return rewriter.notifyMatchFailure(collapseOp, "No loop op producer");
    }
    return fuseCollapseShapeIntoProducerLoop(rewriter, loopOp, collapseOp);
  }
};

WalkResult verifyOperationLegality(Operation *op) {
  if (isa<UnrealizedConversionCastOp>(op)) {
    return WalkResult::interrupt();
  }
  return WalkResult::advance();
}

void FuseConsumersPass::runOnOperation() {
  RewritePatternSet patterns(&getContext());
  patterns.add<FuseIntoGenericOp, FuseIntoLoopOp>(&getContext());
  patterns.add<FuseExtractSliceIntoLoopOp, FuseExtractSliceIntoGenericOp>(
      &getContext());
  patterns.add<FuseCollapseShapeIntoGenericOp, FuseCollapseShapeIntoLoopOp>(
      &getContext());
  populatePCFDropUnusedResultPatterns(patterns);
  if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
    return signalPassFailure();
  }

  // Verify that no unrealized conversion casts remain.
  if (getOperation()->walk(verifyOperationLegality).wasInterrupted()) {
    return signalPassFailure();
  }
}

//===---------------------------------------------------------------------===//
// Consumer fusion impls
//===---------------------------------------------------------------------===//

template <typename OpTy>
static LogicalResult
lookupProducerSlices(OpResult result,
                     SmallVectorImpl<PCF::WriteSliceOp> &slices) {
  OpTy owner = cast<OpTy>(result.getOwner());
  Value tiedArg = owner.getRegionRefArgs()[result.getResultNumber()];
  auto srefType = cast<PCF::ShapedRefType>(tiedArg.getType());
  if (!srefType.isReturnOnlySync()) {
    return failure();
  }
  for (Operation *user : tiedArg.getUsers()) {
    // We can ignore memory reads.
    if (isa<PCF::ReadSliceOp>(user)) {
      continue;
    }
    auto sliceOp = dyn_cast<PCF::WriteSliceOp>(user);
    // TODO: Support vector operands.
    if (!sliceOp || !isa<RankedTensorType>(sliceOp.getSourceType()) ||
        !sliceOp.hasUnitStride()) {
      return failure();
    }
    slices.push_back(sliceOp);
  }
  return success();
}

// Two cases: one operand with multiple producer slices or multiple operands
// with a single producer slice per operand.
// Currently multiple operands <-> multiple producers is unsupported.
template <typename OpTy>
static LogicalResult
matchTilableConsumerImpl(RewriterBase &rewriter, OpTy producerOp,
                         TilingInterface target, ConsumerFusionParams &params) {
  // To create a loop result we need either an initializer or a shape. This
  // can come from either ReifyRankedShapedTypeOpInterface or DPS. If the
  // operand being fused along is itself a destination we get the shape via
  // passthrough with the producer's init/shape.
  if (!isa<ReifyRankedShapedTypeOpInterface, DestinationStyleOpInterface>(
          *target)) {
    return rewriter.notifyMatchFailure(
        target, "unsupported non-reify result shapes or dps op");
  }

  SmallVector<unsigned> &targetOperands = params.operands;
  SetVector<unsigned> &targetResults = params.results;
  SmallVector<PCF::WriteSliceOp> &slices = params.slices;
  assert(targetOperands.empty() && "unexpected non-empty operand list");
  assert(targetResults.empty() && "unexpected non-empty result set");
  assert(slices.empty() && "unexpected non-empty slice list");
  DominanceInfo dominanceInfo(producerOp->getParentOp());
  // First collect the set of operands/results fused along. Additionally verify
  // dominance for other operands.
  for (OpOperand &operand : target->getOpOperands()) {
    auto opResult = dyn_cast<OpResult>(operand.get());
    if (opResult && opResult.getOwner() == producerOp) {
      targetOperands.push_back(operand.getOperandNumber());
      targetResults.insert(opResult.getResultNumber());
    } else {
      if (!dominanceInfo.dominates(operand.get(), producerOp)) {
        return rewriter.notifyMatchFailure(
            target, "unable to fuse due to operand dominance");
      }
    }
  }

  if (targetOperands.empty()) {
    return rewriter.notifyMatchFailure(target, "no operands to fuse along");
  }

  // This dominance check can be expensive in the most general case, however
  // the majority of tilable ops have no or small regions so in practice this
  // isn't so bad.
  WalkResult res = target->walk([&](Operation *containedOp) -> WalkResult {
    if (containedOp == target) {
      return WalkResult::advance();
    }
    bool dominates = llvm::all_of(containedOp->getOperands(), [&](Value v) {
      auto bbArg = dyn_cast<BlockArgument>(v);
      // Check if the tilable op owns the producer of this operand or if the
      // producer dominates the loop we're fusing into.
      Operation *owner = bbArg ? bbArg.getParentRegion()->getParentOp()
                               : cast<OpResult>(v).getOwner();
      return target->isAncestor(owner) ||
             dominanceInfo.dominates(v, producerOp);
    });
    return dominates ? WalkResult::advance() : WalkResult::interrupt();
  });
  if (res.wasInterrupted()) {
    return rewriter.notifyMatchFailure(
        target, "target region users don't dominate producer");
  }

  // Case 1: Single result to fuse along.
  if (targetResults.size() == 1) {
    unsigned resultIndex = *targetResults.begin();
    if (failed(lookupProducerSlices<OpTy>(producerOp->getOpResult(resultIndex),
                                          slices))) {
      return rewriter.notifyMatchFailure(producerOp,
                                         "non write slice producer");
    }

    for (PCF::WriteSliceOp writeSlice : slices) {
      SmallVector<SmallVector<OpFoldResult>> allOffsets(
          targetOperands.size(), writeSlice.getMixedOffsets());
      SmallVector<SmallVector<OpFoldResult>> allSizes(
          targetOperands.size(), writeSlice.getMixedSizes());
      if (!target.isOpFusableWithProducerSlices(targetOperands, allOffsets,
                                                allSizes)) {
        return rewriter.notifyMatchFailure(
            target, "unsupported fusion for single producer");
      }
    }
  } else {
    // Case 2: Multiple results. We must find the most dominated slice to use
    // as the insertion point in this case.
    int64_t leader = -1;
    for (unsigned operandIndex : targetOperands) {
      int64_t currNumSlices = slices.size();
      auto opResult = cast<OpResult>(target->getOperand(operandIndex));
      if (failed(lookupProducerSlices<OpTy>(opResult, slices))) {
        return rewriter.notifyMatchFailure(producerOp,
                                           "non write slice producer");
      }
      if (slices.size() - currNumSlices != 1) {
        return rewriter.notifyMatchFailure(
            target,
            "multiple operand fusion with multiple producers unsupported");
      }

      if (leader < 0) {
        leader = currNumSlices;
      } else {
        PCF::WriteSliceOp currLeader = slices[leader];
        PCF::WriteSliceOp next = slices.back();
        if (next != currLeader) {
          // Match for all writes within the same block to guarantee all
          // required result values are produced "together" (i.e. all required
          // operand slices are written at the same time by the same thread per
          // the *control flow* rather than just the slice offsets/shape).
          if (next->getBlock() != currLeader->getBlock()) {
            return rewriter.notifyMatchFailure(
                target, "unsupported different block insertion points");
          }
          if (dominanceInfo.dominates(currLeader, next)) {
            leader = currNumSlices;
          } else if (!dominanceInfo.dominates(next, currLeader)) {
            return rewriter.notifyMatchFailure(
                target, "could not find single insertion point for multiple "
                        "producer slices");
          }
        }
      }
    }
    SmallVector<SmallVector<OpFoldResult>> allOffsets = llvm::map_to_vector(
        slices, [](PCF::WriteSliceOp op) { return op.getMixedOffsets(); });
    SmallVector<SmallVector<OpFoldResult>> allSizes = llvm::map_to_vector(
        slices, [](PCF::WriteSliceOp op) { return op.getMixedSizes(); });
    if (!target.isOpFusableWithProducerSlices(targetOperands, allOffsets,
                                              allSizes)) {
      return rewriter.notifyMatchFailure(
          target, "unsupported fusion for multiple producer slices");
    }
    if (leader != 0) {
      // Swap the most dominant slice to the beginning.
      std::swap(slices[leader], slices[0]);
      std::swap(targetOperands[leader], targetOperands[0]);
    }
  }
  return success();
}

/// Fuses tilable op |target| into the list of |slices|, one per operand.
/// For example, if |operands| was [0, 2, 4], then the 3 entries in |slices|
/// correspond to the inputs for operands 0, 2, and 4 respectively.
/// The most dominant slice (i.e. the insertion point for the tiled + fused op)
/// is always assumed to be slices[0]. |newResultDests| is the list of new
/// DPS destinations for the tiled op to write to.
static void fuseIntoWriteSlices(RewriterBase &rewriter, TilingInterface target,
                                ArrayRef<unsigned> operands,
                                MutableArrayRef<PCF::WriteSliceOp> slices,
                                ValueRange newResultDests) {
  assert(operands.size() == slices.size() &&
         "expected same number of operands and slices to fuse into");
  OpBuilder::InsertionGuard g(rewriter);
  Location loc = target.getLoc();

  // The contract with the matcher is that the first slice in the list is the
  // most dominant and thus the insertion point for the fused op.
  rewriter.setInsertionPoint(slices.front());

  // Clone the op and replace all operands being fused along with unrealized
  // conversion casts from the distributed producer tile to the undistributed
  // tile. We will forward the input to the unrealized conversion cast directly
  // to the tiled op once finished.
  auto clonedOp = cast<TilingInterface>(rewriter.clone(*target));
  SmallVector<UnrealizedConversionCastOp> unrealizedConversions;
  for (auto [operand, slice] : llvm::zip_equal(operands, slices)) {
    OpOperand &currOperand = clonedOp->getOpOperand(operand);
    Type undistributedType = currOperand.get().getType();
    auto conversion = UnrealizedConversionCastOp::create(
        rewriter, loc, undistributedType, slice.getSource());
    currOperand.assign(conversion.getResult(0));
    unrealizedConversions.push_back(conversion);
  }

  // Get the iteration domain in terms of the operand tiles. This is required
  // to fetch the result tile positions. This and all subsequent interface
  // queries must succeed per the matcher check.
  SmallVector<SmallVector<OpFoldResult>> allOffsets = llvm::map_to_vector(
      slices, [](PCF::WriteSliceOp op) { return op.getMixedOffsets(); });
  SmallVector<SmallVector<OpFoldResult>> allSizes = llvm::map_to_vector(
      slices, [](PCF::WriteSliceOp op) { return op.getMixedSizes(); });
  SmallVector<OpFoldResult> iterDomainOffsets, iterDomainSizes;
  [[maybe_unused]] LogicalResult res =
      clonedOp.getIterationDomainTileFromOperandTiles(
          rewriter, operands, allOffsets, allSizes, iterDomainOffsets,
          iterDomainSizes);
  assert(succeeded(res) && "unexpected iteration domain fetch failed");

  unsigned numResults = clonedOp->getNumResults();
  SmallVector<SmallVector<OpFoldResult>> resultOffsets(numResults);
  SmallVector<SmallVector<OpFoldResult>> resultSizes(numResults);
  for (auto [idx, v] : llvm::enumerate(clonedOp->getResults())) {
    [[maybe_unused]] LogicalResult res = clonedOp.getResultTilePosition(
        rewriter, idx, iterDomainOffsets, iterDomainSizes, resultOffsets[idx],
        resultSizes[idx]);
    assert(succeeded(res) && "Unexpected failure to get result tile position");
  }

  // Tile the cloned op based on the slice shapes.
  FailureOr<TilingResult> tiledResult =
      clonedOp.getTiledImplementationFromOperandTiles(rewriter, operands,
                                                      allOffsets, allSizes);
  assert(succeeded(tiledResult) && "unexpected tiling failure");

  // Create write_slice ops updating the destination for each result.
  OpFoldResult one = rewriter.getIndexAttr(1);
  for (auto [offsets, sizes, replacement, dest] :
       llvm::zip_equal(resultOffsets, resultSizes, tiledResult->tiledValues,
                       newResultDests)) {
    SmallVector<OpFoldResult> strides(offsets.size(), one);
    PCF::WriteSliceOp::create(rewriter, loc, replacement, dest, offsets, sizes,
                              strides);
  }

  // Finally forward the sources of the unrealized conversion casts past each
  // `tensor.extract_slice` consumer. If this fails for any reason we leave the
  // unrealized cast in and fail later for better diagnostics as it is
  // unrecoverable.
  for (UnrealizedConversionCastOp unrealizedCast : unrealizedConversions) {
    SmallVector<Operation *> users(unrealizedCast->getUsers());
    for (Operation *user : users) {
      if (auto extract = dyn_cast<tensor::ExtractSliceOp>(user)) {
        if (extract.getResultType() ==
            unrealizedCast->getOperandTypes().front()) {
          rewriter.replaceOp(extract, unrealizedCast.getOperand(0));
        }
      }
    }
    if (unrealizedCast->use_empty()) {
      rewriter.eraseOp(unrealizedCast);
    }
  }
}

static void addSrefArguments(MLIRContext *context, Location loc,
                             int64_t newArgIndex, Block *entryBlock,
                             TypeRange resultTypes,
                             PCF::ScopeAttrInterface scope) {
  // Add the new region arguments with parent sync scope.
  Attribute syncScope = PCF::SyncOnReturnAttr::get(context);
  for (Type resultType : resultTypes) {
    auto tensorType = cast<RankedTensorType>(resultType);
    auto newSrefType =
        PCF::ShapedRefType::get(context, tensorType.getShape(),
                                tensorType.getElementType(), scope, syncScope);
    entryBlock->insertArgument(newArgIndex, newSrefType, loc);
    ++newArgIndex;
  }
}

static PCF::LoopOp addResults(RewriterBase &rewriter, PCF::LoopOp loopOp,
                              ArrayRef<bool> isTied, ArrayRef<Value> tiedArgs,
                              ArrayRef<Value> dynamicSizes,
                              TypeRange resultTypes) {

  // Append the parameters for the new results to the existing lists.
  SmallVector<Type> newResultTypes(loopOp->getResultTypes());
  llvm::append_range(newResultTypes, resultTypes);
  SmallVector<bool> newIsTied(loopOp.getIsTied());
  llvm::append_range(newIsTied, isTied);
  SmallVector<Value> newDynamicSizes(loopOp.getDynamicSizes());
  llvm::append_range(newDynamicSizes, dynamicSizes);
  SmallVector<Value> newTiedArgs(loopOp.getInits());
  llvm::append_range(newTiedArgs, tiedArgs);

  int64_t numOriginalResults = loopOp->getNumResults();

  // Get the index of the last region ref arg before moving the body over.
  // + 1 because we want the new args to go at the end.
  int64_t newArgIndex = loopOp.getRegionRefArgs().back().getArgNumber() + 1;

  auto newLoopOp =
      PCF::LoopOp::create(rewriter, loopOp.getLoc(), newResultTypes,
                          loopOp.getScope(), loopOp.getCount(), newTiedArgs,
                          newDynamicSizes, newIsTied, loopOp.getSyncOnReturn());
  newLoopOp.getRegion().takeBody(loopOp.getRegion());

  // Add the new region arguments with parent sync scope.
  addSrefArguments(rewriter.getContext(), loopOp.getLoc(), newArgIndex,
                   newLoopOp.getBody(), resultTypes, loopOp.getScope());

  rewriter.replaceOp(loopOp,
                     newLoopOp->getResults().take_front(numOriginalResults));
  return newLoopOp;
}

static PCF::GenericOp
addResults(RewriterBase &rewriter, PCF::GenericOp genericOp,
           ArrayRef<bool> isTied, ArrayRef<Value> tiedArgs,
           ArrayRef<Value> dynamicSizes, TypeRange resultTypes) {

  // Append the parameters for the new results to the existing lists.
  SmallVector<Type> newResultTypes(genericOp->getResultTypes());
  llvm::append_range(newResultTypes, resultTypes);
  SmallVector<bool> newIsTied(genericOp.getIsTied());
  llvm::append_range(newIsTied, isTied);
  SmallVector<Value> newDynamicSizes(genericOp.getDynamicSizes());
  llvm::append_range(newDynamicSizes, dynamicSizes);
  SmallVector<Value> newTiedArgs(genericOp.getInits());
  llvm::append_range(newTiedArgs, tiedArgs);

  int64_t numOriginalResults = genericOp->getNumResults();

  // Get the index of the last region ref arg before moving the body over.
  // + 1 because we want the new args to go at the end.
  int64_t newArgIndex = genericOp.getRegionRefArgs().back().getArgNumber() + 1;

  auto newGenericOp = PCF::GenericOp::create(
      rewriter, genericOp.getLoc(), newResultTypes, genericOp.getScope(),
      newTiedArgs, newDynamicSizes, newIsTied, genericOp.getNumIterators(),
      genericOp.getSyncOnReturn());
  newGenericOp.getRegion().takeBody(genericOp.getRegion());
  newGenericOp.getInitializer().takeBody(genericOp.getInitializer());
  newGenericOp.setNumLeadingArgs(genericOp.getNumLeadingArgs());

  // Add the new region arguments with parent sync scope.
  addSrefArguments(rewriter.getContext(), genericOp.getLoc(), newArgIndex,
                   &newGenericOp.getRegion().front(), resultTypes,
                   genericOp.getScope());

  rewriter.replaceOp(genericOp,
                     newGenericOp->getResults().take_front(numOriginalResults));
  return newGenericOp;
}

template <typename OpTy>
static void fuseTilableConsumerImpl(RewriterBase &rewriter, OpTy producerOp,
                                    TilingInterface target,
                                    ConsumerFusionParams &params) {
  assert(!params.results.empty() && "unexpected empty number of results");

  Location loc = target.getLoc();

  // Step 1. To compute the set of new result shapes, we need to either reify
  // result shapes or get it from a destination a la DPS.
  SmallVector<bool> isTied;
  SmallVector<Value> tiedArgs;
  SmallVector<Value> dynamicSizes;
  SmallVector<Type> resultTypes(target->getResultTypes());

  auto getInitOrCreateEmpty = [&](int64_t resultNumber) -> Value {
    if (OpOperand *tiedInit = producerOp.getTiedInit(resultNumber)) {
      return tiedInit->get();
    }
    return tensor::EmptyOp::create(rewriter, loc,
                                   producerOp.getResultType(resultNumber),
                                   producerOp.getResultDims(resultNumber));
  };

  if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(*target)) {
    for (Value init : dpsOp.getDpsInits()) {
      if (init.getDefiningOp() == producerOp) {
        // There are two options if we are fusing along an init operand:
        //  1. Create a new empty init with the same shape.
        //  2. Use the init of the producer op.
        //
        // This opts for the latter because when fusing we'll replace the actual
        // operand with the thread-local version, so the code will still be
        // correct and it's the closest to the original intent of the
        // ops destination.
        auto result = cast<OpResult>(init);
        if (OpOperand *operand =
                producerOp.getTiedInit(result.getResultNumber())) {
          isTied.push_back(true);
          tiedArgs.push_back(operand->get());
        } else {
          // If there is no tied init, copy the dimensions over.
          isTied.push_back(false);
          ValueRange resultDims =
              producerOp.getResultDims(result.getResultNumber());
          llvm::append_range(dynamicSizes, resultDims);
        }
      } else {
        // Otherwise we just use the init as the tied operand directly.
        isTied.push_back(true);
        tiedArgs.push_back(init);
      }
    }
  } else {
    // For reification ops, we need to construct the result dims in terms of
    // the producer's operands. To do this we replace all operands of |target|
    // coming from the producer with equivalently shaped inits/tensor.empty ops
    // and call reification on that. This is ostensibly a hack as there is no
    // formal guarantee that swapping out operands like this is valid per the
    // interface, however in practice this is valid for all known operations
    // that implement the interface today.

    SmallVector<Value> originalOperands;
    rewriter.setInsertionPoint(producerOp);
    for (unsigned operandIndex : params.operands) {
      Value operand = target->getOperand(operandIndex);
      originalOperands.push_back(operand);
      target->getOpOperand(operandIndex)
          .assign(
              getInitOrCreateEmpty(cast<OpResult>(operand).getResultNumber()));
    }

    SmallVector<SmallVector<OpFoldResult>> outputShapes;
    Operation *nextNode = target->getNextNode();
    Block *currBlock = target->getBlock();

    // Move the op immediately before the producer to get the SSA dominance
    // needed for the result shape dims.
    rewriter.moveOpBefore(target, producerOp);

    // If a fusable operation cannot reify its result shapes under any
    // circumstance, then it was not fusable and should not have been marked as
    // such.
    [[maybe_unused]] auto reifyOp =
        cast<ReifyRankedShapedTypeOpInterface>(*target);
    assert(succeeded(reifyOp.reifyResultShapes(rewriter, outputShapes)) &&
           "unexpected reify result shapes failed");
    if (nextNode) {
      rewriter.moveOpBefore(target, nextNode);
    } else {
      rewriter.moveOpAfter(target, &currBlock->back());
    }

    for (ArrayRef<OpFoldResult> outputShape : outputShapes) {
      dynamicSizes.append(
          llvm::map_to_vector(outputShape, [&](OpFoldResult ofr) {
            return getValueOrCreateConstantIndexOp(rewriter, loc, ofr);
          }));
    }
    isTied.append(outputShapes.size(), false);
  }

  OpTy newRegionOp = addResults(rewriter, producerOp, isTied, tiedArgs,
                                dynamicSizes, resultTypes);
  ValueRange newResultDests =
      newRegionOp.getRegionRefArgs().take_back(resultTypes.size());
  ValueRange replacements =
      newRegionOp.getResults().take_back(resultTypes.size());

  if (params.results.size() == 1) {
    // Single fusion vector means each slice is a different insertion point.
    for (PCF::WriteSliceOp slice : params.slices) {
      // Replicate the slice for each operand. Per the matcher the only
      // supported multi-operand case is if each operand refers to the same
      // value.
      SmallVector<PCF::WriteSliceOp> slices(params.operands.size(), slice);
      fuseIntoWriteSlices(rewriter, target, params.operands, slices,
                          newResultDests);
    }
  } else {
    fuseIntoWriteSlices(rewriter, target, params.operands, params.slices,
                        newResultDests);
  }

  // Finally replace the original fusion target with the newly created loop
  // results.
  rewriter.replaceOp(target, replacements);
}

//===---------------------------------------------------------------------===//
// Extract slice consumer fusion
//===---------------------------------------------------------------------===//

// Compute clamped offsets and sizes of a write_slice to fit within extract
// bounds. This creates affine.min ops to clamp the sizes.
static void computeClampedOffsetsAndSizes(
    RewriterBase &rewriter, Location loc, ArrayRef<OpFoldResult> sliceOffsets,
    ArrayRef<OpFoldResult> sliceSizes, ArrayRef<OpFoldResult> extractSizes,
    SmallVectorImpl<OpFoldResult> &clampedOffsets,
    SmallVectorImpl<OpFoldResult> &clampedSizes) {
  // Clamp sizes to fit within extract bounds.
  for (auto [sliceOffset, sliceSize, extractSize] :
       llvm::zip_equal(sliceOffsets, sliceSizes, extractSizes)) {
    // Compute min(sliceOffset + sliceSize, extractSize) - sliceOffset
    // = min(sliceSize, extractSize - sliceOffset).
    AffineExpr d0, d1, d2;
    bindDims(rewriter.getContext(), d0, d1, d2);
    // d0 = sliceSize, d1 = extractSize, d2 = sliceOffset.
    // clampedSize = min(d0, d1 - d2).
    AffineMap minMap =
        AffineMap::get(3, 0, {d0, d1 - d2}, rewriter.getContext());
    OpFoldResult clampedSize = affine::makeComposedFoldedAffineMin(
        rewriter, loc, minMap, {sliceSize, extractSize, sliceOffset});
    clampedOffsets.push_back(sliceOffset);
    clampedSizes.push_back(clampedSize);
  }
}

template <typename OpTy>
static LogicalResult
fuseExtractSliceIntoProducerImpl(RewriterBase &rewriter, OpTy producerOp,
                                 tensor::ExtractSliceOp extractSliceOp) {
  OpResult producerResult = cast<OpResult>(extractSliceOp.getSource());
  if (!producerResult.hasOneUse()) {
    return rewriter.notifyMatchFailure(producerOp,
                                       "producer result has multiple uses");
  }

  // Only zero offset extract_slice ops are supported.
  if (!llvm::all_of(extractSliceOp.getMixedOffsets(), [](OpFoldResult ofr) {
        return isConstantIntValue(ofr, 0);
      })) {
    return rewriter.notifyMatchFailure(extractSliceOp,
                                       "extract_slice has non-zero offsets");
  }

  // Rank-reducing extract_slice is not yet supported.
  RankedTensorType extractedType = extractSliceOp.getType();
  auto producerResultType = cast<RankedTensorType>(producerResult.getType());
  if (extractedType.getRank() != producerResultType.getRank()) {
    return rewriter.notifyMatchFailure(
        extractSliceOp, "rank-reducing extract_slice not yet supported");
  }

  // Get the write_slice ops for this result.
  unsigned resultIdx = producerResult.getResultNumber();
  SmallVector<PCF::WriteSliceOp> slices;
  if (failed(lookupProducerSlices<OpTy>(producerResult, slices))) {
    return rewriter.notifyMatchFailure(producerOp,
                                       "failed to lookup producer slices");
  }

  if (slices.empty()) {
    return rewriter.notifyMatchFailure(producerOp, "no write_slice producers");
  }

  // Verify all write_slices have unit stride.
  // Only zero-offset extract_slice is supported (already checked above).
  SmallVector<OpFoldResult> extractSizes = extractSliceOp.getMixedSizes();
  for (PCF::WriteSliceOp slice : slices) {
    if (!slice.hasUnitStride()) {
      return rewriter.notifyMatchFailure(slice,
                                         "write_slice has non-unit stride");
    }
  }

  // Get the tied init for this result if it exists.
  OpOperand *tiedInit = producerOp.getTiedInit(resultIdx);
  Value initValue;
  if (tiedInit) {
    // Extract from the tied init.
    rewriter.setInsertionPoint(producerOp);
    initValue = tensor::ExtractSliceOp::create(
        rewriter, producerOp.getLoc(), tiedInit->get(),
        extractSliceOp.getMixedOffsets(), extractSliceOp.getMixedSizes(),
        extractSliceOp.getMixedStrides());
  }

  // Compute new dynamic sizes for the result.
  SmallVector<Value> newDynamicSizes;
  int64_t dynamicDimIdx = 0;

  // First, copy dynamic sizes for results before this one.
  for (unsigned i = 0; i < resultIdx; ++i) {
    auto prevResultType =
        cast<RankedTensorType>(producerOp->getResult(i).getType());
    for (int64_t j = 0, rank = prevResultType.getRank(); j < rank; ++j) {
      if (prevResultType.isDynamicDim(j)) {
        newDynamicSizes.push_back(
            producerOp.getDynamicSizes()[dynamicDimIdx++]);
      }
    }
  }

  // Skip dynamic sizes for the current result (we'll add new ones).
  for (int64_t j = 0, rank = producerResultType.getRank(); j < rank; ++j) {
    if (producerResultType.isDynamicDim(j)) {
      dynamicDimIdx++;
    }
  }

  // Add new dynamic sizes from the extract_slice.
  rewriter.setInsertionPoint(producerOp);
  for (int64_t j = 0, rank = extractedType.getRank(); j < rank; ++j) {
    if (extractedType.isDynamicDim(j)) {
      OpFoldResult size = extractSliceOp.getMixedSizes()[j];
      newDynamicSizes.push_back(
          getValueOrCreateConstantIndexOp(rewriter, producerOp.getLoc(), size));
    }
  }

  // Copy remaining dynamic sizes.
  while (dynamicDimIdx <
         static_cast<int64_t>(producerOp.getDynamicSizes().size())) {
    newDynamicSizes.push_back(producerOp.getDynamicSizes()[dynamicDimIdx++]);
  }

  // Update tied init if present.
  SmallVector<Value> newInits(producerOp.getInits());
  if (tiedInit) {
    // Find the init index using the same logic as getTiedInit.
    int64_t initIdx =
        llvm::count(producerOp.getIsTied().take_front(resultIdx), true);
    newInits[initIdx] = initValue;
  }

  // Create the new result types.
  SmallVector<Type> newResultTypes;
  for (unsigned i = 0, e = producerOp->getNumResults(); i < e; ++i) {
    if (i == resultIdx) {
      newResultTypes.push_back(extractedType);
    } else {
      newResultTypes.push_back(producerOp->getResult(i).getType());
    }
  }

  // Clone the producer op with updated result types and dynamic sizes.
  OpTy newOp;
  if constexpr (std::is_same_v<OpTy, PCF::LoopOp>) {
    newOp = PCF::LoopOp::create(
        rewriter, producerOp.getLoc(), newResultTypes, producerOp.getScope(),
        producerOp.getCount(), newInits, newDynamicSizes,
        producerOp.getIsTied(), producerOp.getSyncOnReturn());
    newOp.getRegion().takeBody(producerOp.getRegion());
  } else {
    newOp = PCF::GenericOp::create(
        rewriter, producerOp.getLoc(), newResultTypes, producerOp.getScope(),
        newInits, newDynamicSizes, producerOp.getIsTied(),
        producerOp.getNumIterators(), producerOp.getSyncOnReturn());
    newOp.getRegion().takeBody(producerOp.getRegion());
    newOp.getInitializer().takeBody(producerOp.getInitializer());
    newOp.setNumLeadingArgs(producerOp.getNumLeadingArgs());
  }

  // Update the region ref arg type to match the new result size.
  Value newRefArg = newOp.getRegionRefArgs()[resultIdx];
  auto oldSrefType = cast<PCF::ShapedRefType>(newRefArg.getType());
  auto newSrefType = PCF::ShapedRefType::get(
      rewriter.getContext(), extractedType.getShape(),
      extractedType.getElementType(), producerOp.getScope(),
      oldSrefType.getSyncScope());
  newRefArg.setType(newSrefType);

  // Get the write_slices in the new op's region (they were moved with the
  // body).
  SmallVector<PCF::WriteSliceOp> newSlices;
  for (Operation *user : newRefArg.getUsers()) {
    if (auto writeSlice = dyn_cast<PCF::WriteSliceOp>(user)) {
      newSlices.push_back(writeSlice);
    }
  }

  // For each write_slice, clamp it to fit within the extracted bounds.
  for (PCF::WriteSliceOp slice : newSlices) {
    OpBuilder::InsertionGuard g(rewriter);
    rewriter.setInsertionPoint(slice);
    Location loc = slice.getLoc();

    // Compute clamped offsets and sizes.
    SmallVector<OpFoldResult> clampedOffsets, clampedSizes;
    computeClampedOffsetsAndSizes(rewriter, loc, slice.getMixedOffsets(),
                                  slice.getMixedSizes(), extractSizes,
                                  clampedOffsets, clampedSizes);

    Value source = slice.getSource();
    auto sourceType = cast<RankedTensorType>(source.getType());

    // Create extract_slice of source to get the clamped portion.
    SmallVector<OpFoldResult> sourceOffsets(sourceType.getRank(),
                                            rewriter.getIndexAttr(0));
    SmallVector<OpFoldResult> sourceStrides(sourceType.getRank(),
                                            rewriter.getIndexAttr(1));
    auto clampedSource = tensor::ExtractSliceOp::create(
        rewriter, loc, source, sourceOffsets, clampedSizes, sourceStrides);

    // Create the new write_slice with clamped offsets/sizes.
    SmallVector<OpFoldResult> strides(clampedOffsets.size(),
                                      rewriter.getIndexAttr(1));
    PCF::WriteSliceOp::create(rewriter, loc, clampedSource, slice.getDest(),
                              clampedOffsets, clampedSizes, strides);

    rewriter.eraseOp(slice);
  }

  // Replace the producer and extract_slice.
  SmallVector<Value> replacements(newOp->getResults());
  rewriter.replaceOp(producerOp, replacements);
  rewriter.replaceOp(extractSliceOp, newOp->getResult(resultIdx));

  return success();
}

//===---------------------------------------------------------------------===//
// Collapse shape consumer fusion
//===---------------------------------------------------------------------===//

/// For a reassociation group, determines the index within the group where
/// "retention" begins. Dimensions from retainedStart to the end of the group
/// form a contiguous chunk in the collapsed layout. Dimensions before
/// retainedStart become loop iterations via scf.forall.
///
/// The innermost dimension is always retained. Each subsequent outer dimension
/// can be folded into the contiguous chunk if the dimension just inside it is
/// fully covered (offset == 0, size == dim_size). Dynamic dimensions stop the
/// extension since we cannot statically verify full coverage.
static int64_t computeRetainedStart(const ReassociationIndices &group,
                                    ArrayRef<OpFoldResult> sliceOffsets,
                                    ArrayRef<OpFoldResult> sliceSizes,
                                    ArrayRef<int64_t> producerShape) {
  if (group.size() <= 1) {
    return 0;
  }

  int64_t retainedStart = group.size() - 1; // Innermost always retained.

  // Extend retained range from innermost outward. To include group position j,
  // the dim just inside (j+1) must be fully covered: offset == 0 and
  // size == dim_size. Dynamic dimensions stop the extension since we cannot
  // statically verify full coverage.
  for (int64_t j = static_cast<int64_t>(group.size()) - 2; j >= 0; --j) {
    int64_t innerDim = group[j + 1];
    if (ShapedType::isDynamic(producerShape[innerDim])) {
      break;
    }
    if (isConstantIntValue(sliceOffsets[innerDim], 0) &&
        isConstantIntValue(sliceSizes[innerDim], producerShape[innerDim])) {
      retainedStart = j;
    } else {
      break;
    }
  }

  return retainedStart;
}

/// Gets OpFoldResult dim sizes for the producer result. Static dims become
/// IntegerAttr, dynamic dims become Values from the tied init or the op's
/// dynamic_sizes.
template <typename OpTy>
static SmallVector<OpFoldResult> getProducerResultDimSizes(
    RewriterBase &rewriter, OpTy producerOp, unsigned resultIdx,
    RankedTensorType producerResultType, OpOperand *tiedInit) {
  if (tiedInit) {
    return tensor::getMixedSizes(rewriter, producerOp.getLoc(),
                                 tiedInit->get());
  }
  return getMixedValues(producerResultType.getShape(),
                        producerOp.getResultDims(resultIdx), rewriter);
}

/// Computes the linearized offset and size for a single reassociation group.
///
/// For a group [d0, d1, ..., dN] with retainedStart = k:
///   Loop dims: d0, ..., d_{k-1} (iterate with scf.forall)
///   Retained dims: d_k, ..., d_N (form contiguous chunk)
///
/// The linearized offset is computed via affine.linearize_index with
/// multi-index [adj_0, ..., adj_k, 0, ..., 0] and basis
/// [producerDimSizes[group[0..N]]]:
///   adj_j = O_j + i_j (loop dim, j < k)
///         = O_j       (first retained dim, j == k)
///         = 0         (inner retained dims, j > k)
///
/// Linearized size = sizes[group[k]] * product(producerDimSizes[group[k+1..N]])
static void computeGroupLinearizedOffsetAndSize(
    RewriterBase &rewriter, Location loc, const ReassociationIndices &group,
    int64_t retainedStart, ArrayRef<OpFoldResult> sliceOffsets,
    ArrayRef<OpFoldResult> sliceSizes, ArrayRef<OpFoldResult> producerDimSizes,
    ArrayRef<Value> groupLoopIndices, OpFoldResult &collapsedOffset,
    OpFoldResult &collapsedSize) {
  MLIRContext *ctx = rewriter.getContext();

  // Build multi-index and basis for affine.linearize_index.
  SmallVector<Value> multiIndex;
  SmallVector<OpFoldResult> basis;
  // Pre-create zero constant if needed for inner retained dims.
  Value zero;
  if (retainedStart + 1 < static_cast<int64_t>(group.size())) {
    zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
  }
  int64_t loopIdx = 0;
  for (int64_t j = 0, e = static_cast<int64_t>(group.size()); j < e; ++j) {
    basis.push_back(producerDimSizes[group[j]]);
    if (j < retainedStart) {
      // Loop dim: adjusted offset = sliceOffset + loopIndex.
      AffineExpr e0, e1;
      bindDims(ctx, e0, e1);
      OpFoldResult idx = affine::makeComposedFoldedAffineApply(
          rewriter, loc, e0 + e1,
          {sliceOffsets[group[j]], groupLoopIndices[loopIdx++]});
      multiIndex.push_back(getValueOrCreateConstantIndexOp(rewriter, loc, idx));
    } else if (j == retainedStart) {
      // First retained dim: offset passes through.
      multiIndex.push_back(getValueOrCreateConstantIndexOp(
          rewriter, loc, sliceOffsets[group[j]]));
    } else {
      // Inner retained dim: zero offset (fully covered).
      multiIndex.push_back(zero);
    }
  }

  // Linearize the offset using the group's producer dimension sizes as basis.
  collapsedOffset = affine::AffineLinearizeIndexOp::create(
                        rewriter, loc, multiIndex, basis, /*disjoint=*/true)
                        .getResult();

  // Compute the linearized size: sliceSizes[retainedStart] multiplied by the
  // product of producer dim sizes for all dims after retainedStart.
  collapsedSize = sliceSizes[group[retainedStart]];
  for (int64_t j = retainedStart + 1, e = static_cast<int64_t>(group.size());
       j < e; ++j) {
    AffineExpr d0, d1;
    bindDims(ctx, d0, d1);
    collapsedSize = affine::makeComposedFoldedAffineApply(
        rewriter, loc, d0 * d1, {collapsedSize, producerDimSizes[group[j]]});
  }
}

/// Computes collapsed offsets and sizes for a write_slice given reassociation
/// indices and per-group retained start analysis. For groups with loop dims,
/// the provided loopIndices are consumed in order (group by group, inner loop
/// dims first).
static void computeCollapsedOffsetsAndSizes(
    RewriterBase &rewriter, Location loc, ArrayRef<OpFoldResult> sliceOffsets,
    ArrayRef<OpFoldResult> sliceSizes, ArrayRef<OpFoldResult> producerDimSizes,
    ArrayRef<ReassociationIndices> reassociation,
    ArrayRef<int64_t> groupRetainedStarts, ArrayRef<Value> loopIndices,
    SmallVectorImpl<OpFoldResult> &collapsedOffsets,
    SmallVectorImpl<OpFoldResult> &collapsedSizes) {
  int64_t loopIdxOffset = 0;
  for (auto [groupIdx, group] : llvm::enumerate(reassociation)) {
    if (group.size() == 1) {
      // Singleton group: offset and size pass through.
      collapsedOffsets.push_back(sliceOffsets[group[0]]);
      collapsedSizes.push_back(sliceSizes[group[0]]);
      continue;
    }

    int64_t retainedStart = groupRetainedStarts[groupIdx];
    int64_t numLoopDims = retainedStart;
    ArrayRef<Value> groupLoopIndices =
        loopIndices.slice(loopIdxOffset, numLoopDims);
    loopIdxOffset += numLoopDims;

    OpFoldResult offset, size;
    computeGroupLinearizedOffsetAndSize(
        rewriter, loc, group, retainedStart, sliceOffsets, sliceSizes,
        producerDimSizes, groupLoopIndices, offset, size);
    collapsedOffsets.push_back(offset);
    collapsedSizes.push_back(size);
  }
}

template <typename OpTy>
LogicalResult
fuseCollapseShapeIntoProducerImpl(RewriterBase &rewriter, OpTy producerOp,
                                  tensor::CollapseShapeOp collapseOp) {
  OpResult producerResult = cast<OpResult>(collapseOp.getSrc());
  if (!producerResult.hasOneUse()) {
    return rewriter.notifyMatchFailure(producerOp,
                                       "producer result has multiple uses");
  }

  unsigned resultIdx = producerResult.getResultNumber();
  auto producerResultType = cast<RankedTensorType>(producerResult.getType());
  RankedTensorType collapsedType = collapseOp.getResultType();
  ArrayRef<int64_t> producerShape = producerResultType.getShape();

  SmallVector<ReassociationIndices> reassociation =
      collapseOp.getReassociationIndices();

  // Get the write_slice ops for this result.
  SmallVector<PCF::WriteSliceOp> slices;
  if (failed(lookupProducerSlices<OpTy>(producerResult, slices))) {
    return rewriter.notifyMatchFailure(producerOp,
                                       "failed to lookup producer slices");
  }

  if (slices.empty()) {
    return rewriter.notifyMatchFailure(producerOp, "no write_slice producers");
  }

  // Verify all write_slices have unit stride.
  for (PCF::WriteSliceOp slice : slices) {
    if (!slice.hasUnitStride()) {
      return rewriter.notifyMatchFailure(slice,
                                         "write_slice has non-unit stride");
    }
  }

  // Get the tied init and setup the sizes in terms of values available to the
  // producer if present.
  OpOperand *tiedInit = producerOp.getTiedInit(resultIdx);

  rewriter.setInsertionPoint(producerOp);
  SmallVector<OpFoldResult> producerDimSizes = getProducerResultDimSizes(
      rewriter, producerOp, resultIdx, producerResultType, tiedInit);

  // Collapse the init if it exists.
  Value initValue;
  if (tiedInit) {
    rewriter.setInsertionPoint(producerOp);
    initValue = tensor::CollapseShapeOp::create(rewriter, producerOp.getLoc(),
                                                collapsedType, tiedInit->get(),
                                                reassociation);
  }

  // Compute new dynamic sizes. getDynamicSizes() only contains entries for
  // untied results; tied results get their dynamic sizes from the init.
  bool isResultTied = (tiedInit != nullptr);
  SmallVector<Value> newDynamicSizes;
  int64_t dynamicDimIdx = 0;

  // Copy dynamic sizes for untied results before this one.
  for (unsigned i = 0; i < resultIdx; ++i) {
    if (producerOp.getTiedInit(i)) {
      continue; // Tied results don't have entries in getDynamicSizes().
    }
    auto prevResultType =
        cast<RankedTensorType>(producerOp->getResult(i).getType());
    for (int64_t j = 0, rank = prevResultType.getRank(); j < rank; ++j) {
      if (prevResultType.isDynamicDim(j)) {
        newDynamicSizes.push_back(
            producerOp.getDynamicSizes()[dynamicDimIdx++]);
      }
    }
  }

  // Skip dynamic sizes for the current result (only if untied).
  if (!isResultTied) {
    for (int64_t j = 0, rank = producerResultType.getRank(); j < rank; ++j) {
      if (producerResultType.isDynamicDim(j)) {
        dynamicDimIdx++;
      }
    }
  }

  // Add dynamic sizes for the collapsed type (only if untied). For dynamic
  // collapsed dims, compute the product of producer dim sizes in the group.
  Location producerLoc = producerOp.getLoc();
  if (!isResultTied) {
    rewriter.setInsertionPoint(producerOp);
    for (int64_t j = 0, rank = collapsedType.getRank(); j < rank; ++j) {
      if (collapsedType.isDynamicDim(j)) {
        const ReassociationIndices &group = reassociation[j];
        OpFoldResult product = rewriter.getIndexAttr(1);
        for (int64_t dim : group) {
          AffineExpr d0, d1;
          bindDims(rewriter.getContext(), d0, d1);
          product = affine::makeComposedFoldedAffineApply(
              rewriter, producerLoc, d0 * d1, {product, producerDimSizes[dim]});
        }
        newDynamicSizes.push_back(
            getValueOrCreateConstantIndexOp(rewriter, producerLoc, product));
      }
    }
  }

  // Copy remaining dynamic sizes.
  while (dynamicDimIdx <
         static_cast<int64_t>(producerOp.getDynamicSizes().size())) {
    newDynamicSizes.push_back(producerOp.getDynamicSizes()[dynamicDimIdx++]);
  }

  // Update tied init if present.
  SmallVector<Value> newInits(producerOp.getInits());
  if (tiedInit) {
    int64_t initIdx =
        llvm::count(producerOp.getIsTied().take_front(resultIdx), true);
    newInits[initIdx] = initValue;
  }

  // Create the new result types with collapsed type for this result.
  SmallVector<Type> newResultTypes;
  for (unsigned i = 0, e = producerOp->getNumResults(); i < e; ++i) {
    if (i == resultIdx) {
      newResultTypes.push_back(collapsedType);
    } else {
      newResultTypes.push_back(producerOp->getResult(i).getType());
    }
  }

  // Clone the producer op with updated result type.
  OpTy newOp;
  if constexpr (std::is_same_v<OpTy, PCF::LoopOp>) {
    newOp = PCF::LoopOp::create(
        rewriter, producerLoc, newResultTypes, producerOp.getScope(),
        producerOp.getCount(), newInits, newDynamicSizes,
        producerOp.getIsTied(), producerOp.getSyncOnReturn());
    newOp.getRegion().takeBody(producerOp.getRegion());
  } else {
    newOp = PCF::GenericOp::create(
        rewriter, producerLoc, newResultTypes, producerOp.getScope(), newInits,
        newDynamicSizes, producerOp.getIsTied(), producerOp.getNumIterators(),
        producerOp.getSyncOnReturn());
    newOp.getRegion().takeBody(producerOp.getRegion());
    newOp.getInitializer().takeBody(producerOp.getInitializer());
    newOp.setNumLeadingArgs(producerOp.getNumLeadingArgs());
  }

  // Update the region ref arg type to match the collapsed shape.
  Value newRefArg = newOp.getRegionRefArgs()[resultIdx];
  auto oldSrefType = cast<PCF::ShapedRefType>(newRefArg.getType());
  auto newSrefType = PCF::ShapedRefType::get(
      rewriter.getContext(), collapsedType.getShape(),
      collapsedType.getElementType(), producerOp.getScope(),
      oldSrefType.getSyncScope());
  newRefArg.setType(newSrefType);

  // Get the write_slices in the new op's region.
  SmallVector<PCF::WriteSliceOp> newSlices;
  for (Operation *user : newRefArg.getUsers()) {
    if (auto writeSlice = dyn_cast<PCF::WriteSliceOp>(user)) {
      newSlices.push_back(writeSlice);
    }
  }

  // Transform each write_slice to use collapsed offsets/sizes.
  for (PCF::WriteSliceOp slice : newSlices) {
    OpBuilder::InsertionGuard g(rewriter);
    rewriter.setInsertionPoint(slice);
    Location loc = slice.getLoc();

    SmallVector<OpFoldResult> sliceOffsets = slice.getMixedOffsets();
    SmallVector<OpFoldResult> sliceSizes = slice.getMixedSizes();
    Value source = slice.getSource();
    auto sourceType = cast<RankedTensorType>(source.getType());

    // Per-group analysis: determine retained start and collect loop dims.
    SmallVector<int64_t> groupRetainedStarts;
    // Each entry is (global dim index, loop bound).
    SmallVector<std::pair<int64_t, OpFoldResult>> loopDims;
    for (const ReassociationIndices &group : reassociation) {
      int64_t retainedStart =
          computeRetainedStart(group, sliceOffsets, sliceSizes, producerShape);
      groupRetainedStarts.push_back(retainedStart);
      for (int64_t j = 0; j < retainedStart; ++j) {
        loopDims.push_back({group[j], sliceSizes[group[j]]});
      }
    }

    if (loopDims.empty()) {
      // No loops needed. Directly collapse source and linearize offsets.
      SmallVector<OpFoldResult> collapsedOffsets, collapsedSizes;
      computeCollapsedOffsetsAndSizes(
          rewriter, loc, sliceOffsets, sliceSizes, producerDimSizes,
          reassociation, groupRetainedStarts,
          /*loopIndices=*/{}, collapsedOffsets, collapsedSizes);

      Value collapsedSource =
          tensor::CollapseShapeOp::create(rewriter, loc, source, reassociation);

      SmallVector<OpFoldResult> strides(collapsedOffsets.size(),
                                        rewriter.getIndexAttr(1));
      PCF::WriteSliceOp::create(rewriter, loc, collapsedSource, slice.getDest(),
                                collapsedOffsets, collapsedSizes, strides);
    } else {
      // Emit scf.forall loop over the non-retained (loop) dimensions.
      SmallVector<OpFoldResult> lbs(loopDims.size(), rewriter.getIndexAttr(0));
      SmallVector<OpFoldResult> ubs;
      SmallVector<OpFoldResult> steps(loopDims.size(),
                                      rewriter.getIndexAttr(1));
      for (auto &[dim, bound] : loopDims) {
        ubs.push_back(bound);
      }

      auto forallOp = scf::ForallOp::create(rewriter, loc, lbs, ubs, steps,
                                            /*outputs=*/ValueRange(),
                                            /*mapping=*/std::nullopt);

      // Build the forall body.
      rewriter.setInsertionPointToStart(forallOp.getBody());
      SmallVector<Value> inductionVars(forallOp.getInductionVars());

      // Build a map from global dim index to forall induction variable index.
      DenseMap<int64_t, int64_t> dimToLoopIdx;
      for (auto [idx, dimAndBound] : llvm::enumerate(loopDims)) {
        dimToLoopIdx[dimAndBound.first] = idx;
      }

      // Build extract_slice: loop dims get induction var offset with size 1,
      // retained dims get offset 0 with original size.
      SmallVector<OpFoldResult> extractOffsets(sourceType.getRank(),
                                               rewriter.getIndexAttr(0));
      SmallVector<OpFoldResult> extractSizes;
      SmallVector<OpFoldResult> extractStrides(sourceType.getRank(),
                                               rewriter.getIndexAttr(1));
      for (int64_t dim = 0, rank = sourceType.getRank(); dim < rank; ++dim) {
        auto it = dimToLoopIdx.find(dim);
        if (it != dimToLoopIdx.end()) {
          extractOffsets[dim] = inductionVars[it->second];
          extractSizes.push_back(rewriter.getIndexAttr(1));
        } else {
          extractSizes.push_back(sliceSizes[dim]);
        }
      }

      // Extract the sub-tensor for this loop iteration.
      Value extractedSource = tensor::ExtractSliceOp::create(
          rewriter, loc, source, extractOffsets, extractSizes, extractStrides);

      // Collapse the extracted source according to the reassociation.
      Value collapsedSource = tensor::CollapseShapeOp::create(
          rewriter, loc, extractedSource, reassociation);

      // Compute collapsed offsets and sizes.
      SmallVector<Value> loopIndicesVec(inductionVars.begin(),
                                        inductionVars.end());
      SmallVector<OpFoldResult> collapsedOffsets, collapsedSizes;
      computeCollapsedOffsetsAndSizes(rewriter, loc, sliceOffsets, sliceSizes,
                                      producerDimSizes, reassociation,
                                      groupRetainedStarts, loopIndicesVec,
                                      collapsedOffsets, collapsedSizes);

      SmallVector<OpFoldResult> strides(collapsedOffsets.size(),
                                        rewriter.getIndexAttr(1));
      PCF::WriteSliceOp::create(rewriter, loc, collapsedSource, slice.getDest(),
                                collapsedOffsets, collapsedSizes, strides);
    }

    rewriter.eraseOp(slice);
  }

  // Replace the original producer and collapse_shape.
  SmallVector<Value> replacements(newOp->getResults());
  rewriter.replaceOp(producerOp, replacements);
  rewriter.replaceOp(collapseOp, newOp->getResult(resultIdx));

  return success();
}

} // namespace

//===---------------------------------------------------------------------===//
// Public API Specializations
//===---------------------------------------------------------------------===//

LogicalResult
fuseExtractSliceIntoProducerLoop(RewriterBase &rewriter, PCF::LoopOp loopOp,
                                 tensor::ExtractSliceOp extractSliceOp) {
  return fuseExtractSliceIntoProducerImpl(rewriter, loopOp, extractSliceOp);
}

LogicalResult
fuseExtractSliceIntoProducerGeneric(RewriterBase &rewriter,
                                    PCF::GenericOp genericOp,
                                    tensor::ExtractSliceOp extractSliceOp) {
  return fuseExtractSliceIntoProducerImpl(rewriter, genericOp, extractSliceOp);
}

LogicalResult
fuseCollapseShapeIntoProducerLoop(RewriterBase &rewriter, PCF::LoopOp loopOp,
                                  tensor::CollapseShapeOp collapseOp) {
  return fuseCollapseShapeIntoProducerImpl(rewriter, loopOp, collapseOp);
}

LogicalResult
fuseCollapseShapeIntoProducerGeneric(RewriterBase &rewriter,
                                     PCF::GenericOp genericOp,
                                     tensor::CollapseShapeOp collapseOp) {
  return fuseCollapseShapeIntoProducerImpl(rewriter, genericOp, collapseOp);
}

LogicalResult matchTilableConsumer(RewriterBase &rewriter,
                                   PCF::GenericOp producerOp,
                                   TilingInterface target,
                                   ConsumerFusionParams &params) {
  return matchTilableConsumerImpl(rewriter, producerOp, target, params);
}

LogicalResult matchTilableConsumer(RewriterBase &rewriter,
                                   PCF::LoopOp producerOp,
                                   TilingInterface target,
                                   ConsumerFusionParams &params) {
  return matchTilableConsumerImpl(rewriter, producerOp, target, params);
}

void fuseTilableConsumer(RewriterBase &rewriter, PCF::GenericOp producerOp,
                         TilingInterface target, ConsumerFusionParams &params) {
  return fuseTilableConsumerImpl(rewriter, producerOp, target, params);
}

void fuseTilableConsumer(RewriterBase &rewriter, PCF::LoopOp producerOp,
                         TilingInterface target, ConsumerFusionParams &params) {
  return fuseTilableConsumerImpl(rewriter, producerOp, target, params);
}

} // namespace mlir::iree_compiler::IREE::PCF
