// Copyright 2020 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/Dialect/Flow/Transforms/RegionOpUtils.h"

#include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/Support/CommandLine.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Analysis/TopologicalSortUtils.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tensor/Utils/Utils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/Block.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/Value.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/RegionUtils.h"

#define DEBUG_TYPE "iree-flow-region-op-utils"

// NOTE: These flags are added for experimental purposes only
// for developer control. These should be treated as internal
// compiler implementation details.
static llvm::cl::opt<int> clInlineConstantByteLength(
    "iree-flow-inline-constants-max-byte-length",
    llvm::cl::desc("Maximum byte-length of tensor constant that can be inlined "
                   "into a dispatch region or 0 to disable inlining."),
    llvm::cl::init(256));

// TODO(#18457, #18447): Remove once backends support gather fusion.
static llvm::cl::opt<bool>
    clEnableGatherFusion("iree-flow-enable-gather-fusion",
                         llvm::cl::desc("Fuse gather-like ops with consumer."),
                         llvm::cl::init(false));

namespace mlir::iree_compiler::IREE::Flow {

//===----------------------------------------------------------------------===//
// Methods for getting the workload information for dispatch region creation.
//===----------------------------------------------------------------------===//

static SmallVector<Range> getLoopRangesImpl(TilingInterface tilableOp,
                                            Location loc, OpBuilder &builder) {
  SmallVector<Range> loopRanges = tilableOp.getIterationDomain(builder);
  Value one = builder.create<arith::ConstantIndexOp>(loc, 1);
  for (auto iteratorType : llvm::enumerate(tilableOp.getLoopIteratorTypes())) {
    if (iteratorType.value() == utils::IteratorType::reduction) {
      loopRanges[iteratorType.index()].size = one;
    }
  }
  return loopRanges;
}

static SmallVector<Range> getLoopRangesFromValue(Value source, Location loc,
                                                 OpBuilder &builder) {
  SmallVector<OpFoldResult> dimValues =
      tensor::getMixedSizes(builder, loc, source);
  OpFoldResult zero = builder.getIndexAttr(0);
  OpFoldResult one = builder.getIndexAttr(1);
  return llvm::map_to_vector(dimValues, [&](OpFoldResult dimValue) {
    return Range{zero, dimValue, one};
  });
}

static SmallVector<Range>
getLoopRangesImpl(ReifyRankedShapedTypeOpInterface shapedOp, Location loc,
                  OpBuilder &builder) {
  Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
  Value one = builder.create<arith::ConstantIndexOp>(loc, 1);
  ReifiedRankedShapedTypeDims resultDims;
  LogicalResult status = shapedOp.reifyResultShapes(builder, resultDims);
  (void)status;
  assert(succeeded(status) && "reifyResultShapes failed");
  return llvm::map_to_vector(
      resultDims[0], [&](OpFoldResult v) { return Range{zero, v, one}; });
}

/// For a given operation returns the loop ranges needed to compute the op.
SmallVector<Range> getLoopRanges(Operation *op, Location loc,
                                 OpBuilder &builder) {
  return llvm::TypeSwitch<Operation *, SmallVector<Range>>(op)
      .Case<Encoding::SetEncodingOp, Encoding::UnsetEncodingOp,
            tensor::InsertSliceOp>([&](auto op) {
        return getLoopRangesFromValue(op.getSource(), loc, builder);
      })
      .Case<TilingInterface>([&](TilingInterface op) {
        return getLoopRangesImpl(op, loc, builder);
      })
      .Case<ReifyRankedShapedTypeOpInterface>([&](auto shapedOp) {
        return getLoopRangesImpl(shapedOp, loc, builder);
      })
      .Default([](Operation *op) -> SmallVector<Range> {
        llvm_unreachable("op not supported");
      });
}

/// Return `true` if an operation is within a `flow.dispatch.region` or
/// `flow.dispatch.workgroups` op.
bool isNonNullAndOutsideDispatch(Operation *op) {
  if (!op)
    return false;
  Operation *parentOp = op->getParentOp();
  while (parentOp) {
    if (isa<IREE::Flow::DispatchRegionOp, IREE::Flow::DispatchWorkgroupsOp>(
            parentOp)) {
      return false;
    }
    parentOp = parentOp->getParentOp();
  }
  return true;
}
bool isNonNullAndOutsideDispatch(ArrayRef<Operation *> operations) {
  return llvm::all_of(operations, [](Operation *op) {
    return isNonNullAndOutsideDispatch(op);
  });
}

/// Compute the workload to use for the workgroup based on the root op.
static SmallVector<Value> getWorkloadForRootOp(OpBuilder &builder,
                                               Operation *rootOp) {
  // Compute workgroup count to use for the dispatch op. These are the ranges
  // of the outermost parallel loops that can be distributed.
  Location loc = rootOp->getLoc();
  SmallVector<Range> loopRanges =
      IREE::Flow::getLoopRanges(rootOp, loc, builder);
  AffineExpr s0, s1, s2;
  bindSymbols(builder.getContext(), s0, s1, s2);
  AffineMap workload = AffineMap::get(0, 3, (s1 - s0).ceilDiv(s2));
  return llvm::map_to_vector(loopRanges, [&](Range r) -> Value {
    Value offset = getValueOrCreateConstantIndexOp(builder, loc, r.offset);
    Value size = getValueOrCreateConstantIndexOp(builder, loc, r.size);
    Value stride = getValueOrCreateConstantIndexOp(builder, loc, r.stride);
    return builder.create<affine::AffineApplyOp>(
        rootOp->getLoc(), workload, ValueRange{offset, size, stride});
  });
}

/// For very specific cases where the shape is data dependent, we cannot
/// use the normal way of workgroup count calculation through use of
/// program slices within the body of the op. This is because
/// the program slice will also include other tensors.. The general solution
/// here requires making the workgroup count calculation also run on the device.
/// This isnt plumbed through yet. For these, just use
/// the workgroup count calculation based on the root of the DAG.
static bool checkShapeIsDataDependant(Operation *op) {
  if (auto linalgOp = dyn_cast<linalg::GenericOp>(op)) {
    /// Currently checked case is
    /// - linalg.generic op
    /// - all parallel iterators
    /// - output shape depends on tensor.extract;
    if (!llvm::all_of(linalgOp.getIteratorTypesArray(),
                      linalg::isParallelIterator)) {
      return false;
    }
    BackwardSliceOptions options;
    options.inclusive = true;
    options.filter = [](Operation *op) {
      // Only look for slices with a few ops to not blow up the slice
      // computation.
      if (!isa<arith::IndexCastOp, tensor::EmptyOp, tensor::ExtractOp>(op)) {
        return false;
      }
      // The slice computation method has an assert that the parent op
      // needs to have a single region with a single block. This seems
      // unnecessary for IREEs use case. For now avoid this assert by bailing if
      // any operands are block arguments.
      if (llvm::any_of(op->getOperands(), llvm::IsaPred<BlockArgument>)) {
        auto parentOp = op->getParentOp();
        if (parentOp->getNumRegions() != 1 ||
            parentOp->getRegion(0).getBlocks().size() != 1) {
          return false;
        }
      }
      return true;
    };
    llvm::SetVector<Operation *> slice;
    for (Value initOperand : linalgOp.getDpsInits()) {
      mlir::getBackwardSlice(initOperand, &slice, options);
    }
    return llvm::any_of(slice, llvm::IsaPred<tensor::ExtractOp>);
  }
  return false;
}

/// Populate the workgroup count calculation to be based on root of the DAG op.
/// These is the legacy path for workgroup count calculation that
/// arent handled by the current default.
static void createWorkgroupCountFromDagRootRegion(
    RewriterBase &rewriter, IREE::Flow::DispatchRegionOp &regionOp,
    TypeRange workloadTypes, ArrayRef<Location> workloadLocs) {
  Region &countRegion = regionOp.getWorkgroupCount();
  if (!countRegion.empty())
    return;
  Block *body = rewriter.createBlock(&countRegion, countRegion.begin(),
                                     workloadTypes, workloadLocs);
  auto args = body->getArguments();
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPointToStart(body);
  Location loc = regionOp.getLoc();
  auto countOp =
      rewriter.create<IREE::Flow::DispatchWorkgroupCountFromDagRootOp>(loc,
                                                                       args);
  rewriter.create<IREE::Flow::ReturnOp>(loc, countOp->getResults());
}

/// Return `true` if the given type is a ShapedType and has at least one
/// dynamic dimension.
static bool hasDynamicShape(Type t) {
  auto shapedType = llvm::dyn_cast<ShapedType>(t);
  if (!shapedType)
    return false;
  return !shapedType.hasStaticShape();
}

/// Reify the dynamic dimensions of the given value.
static LogicalResult
reifyDynamicResultDimsImpl(OpBuilder &b, Value value,
                           SmallVectorImpl<Value> &dynamicDims,
                           bool createTensorDimOps) {
  OpBuilder::InsertionGuard guard(b);

  // Case 1: No dynamic result dims.
  if (!hasDynamicShape(value.getType()))
    return success();

  // There is at least one dynamic dimension, continue...
  ShapedType shapedType = llvm::cast<ShapedType>(value.getType());

  // Helper function that generates tensor.dim ops.
  auto emitTensorDimOps = [&]() {
    for (int64_t i = 0; i < shapedType.getRank(); ++i) {
      if (shapedType.isDynamicDim(i)) {
        Value dim = b.create<tensor::DimOp>(value.getLoc(), value, i);
        dynamicDims.push_back(dim);
      }
    }
  };

  // Case 2: Value is a block argument.
  if (auto bbArg = llvm::dyn_cast<BlockArgument>(value)) {
    if (!createTensorDimOps)
      return failure();

    b.setInsertionPointToStart(bbArg.getOwner());
    emitTensorDimOps();
    return success();
  }

  // Value is an OpResult.
  Operation *op = value.getDefiningOp();
  OpResult opResult = llvm::cast<OpResult>(value);
  b.setInsertionPoint(op);

  // Case 3: Value is tied. Reify the dimensions of the tied operand.
  auto tiedOp = dyn_cast<IREE::Util::TiedOpInterface>(op);
  if (tiedOp) {
    Value tiedOperand = tiedOp.getTiedResultOperand(value);
    if (tiedOperand && tiedOperand.getType() == value.getType())
      return reifyDynamicResultDimsImpl(b, tiedOperand, dynamicDims,
                                        createTensorDimOps);
  }

  // Case 4: Query ShapeAwareOpInterface.
  auto shapeAwareOp = dyn_cast<IREE::Util::ShapeAwareOpInterface>(op);
  if (shapeAwareOp) {
    ValueRange dims =
        shapeAwareOp.getResultDynamicDims(opResult.getResultNumber());
    dynamicDims.append(dims.begin(), dims.end());
    return success();
  }

  // Case 5: Query ReifyRankedShapedTypeOpInterface.
  auto reifyShapeOp = dyn_cast<ReifyRankedShapedTypeOpInterface>(op);
  if (reifyShapeOp) {
    ReifiedRankedShapedTypeDims dims;
    if (failed(reifyShapeOp.reifyResultShapes(b, dims)))
      return failure();
    for (int64_t i = 0; i < shapedType.getRank(); ++i)
      if (shapedType.isDynamicDim(i))
        dynamicDims.push_back(dims[opResult.getResultNumber()][i].get<Value>());
    return success();
  }

  if (!createTensorDimOps)
    return failure();

  // None of the above. Insert tensor.dim ops.
  b.setInsertionPointAfter(op);
  emitTensorDimOps();
  return success();
}

/// Reify the dynamic dimensions of the given value.
LogicalResult reifyDynamicResultDims(OpBuilder &b, Value value,
                                     SmallVectorImpl<Value> &dynamicDims) {
  return reifyDynamicResultDimsImpl(b, value, dynamicDims,
                                    /*createTensorDimOps=*/true);
}

LogicalResult
getOptimizedDynamicResultDims(OpBuilder &b, Value value,
                              SmallVectorImpl<Value> &dynamicDims) {
  return reifyDynamicResultDimsImpl(b, value, dynamicDims,
                                    /*createTensorDimOps=*/false);
}

// Append a result to the given DispatchRegionOp. The newly created
// DispatchRegionOp is returned.
FailureOr<IREE::Flow::DispatchRegionOp> appendDispatchRegionResults(
    RewriterBase &rewriter, IREE::Flow::DispatchRegionOp regionOp,
    ArrayRef<Value> results, ArrayRef<SmallVector<Value>> dynamicDims) {
  if (results.empty()) {
    return regionOp;
  }
  assert(results.size() == dynamicDims.size() &&
         "expected as many dynamic dims list as the number of results");

  // Collect the current region dynamic dims, and result types.
  SmallVector<Value> regionDynamicDims(regionOp.getResultDims().begin(),
                                       regionOp.getResultDims().end());
  SmallVector<Type> resultTypes(regionOp.getResultTypes().begin(),
                                regionOp.getResultTypes().end());

  // Collect the current dispatch yielded values.
  auto returnOp =
      cast<IREE::Flow::ReturnOp>(regionOp.getBody().front().getTerminator());
  SmallVector<Value> returnedValues(returnOp.getOperands().begin(),
                                    returnOp.getOperands().end());

  for (auto [index, result] : llvm::enumerate(results)) {
#ifndef NDEBUG
    auto tensorType = cast<RankedTensorType>(result.getType());
    assert(tensorType.getNumDynamicDims() == dynamicDims[index].size() &&
           "incorrect number of dynamicDims provided");
#endif // NDEBUG
    resultTypes.push_back(result.getType());
    regionDynamicDims.append(dynamicDims[index]);
    returnedValues.push_back(result);
  }

  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPoint(regionOp);

  // Create new DispatchRegionOp and move over the body.
  auto newRegionOp = rewriter.create<IREE::Flow::DispatchRegionOp>(
      regionOp->getLoc(), resultTypes, regionDynamicDims,
      regionOp.getWorkload());
  rewriter.inlineRegionBefore(regionOp.getBody(), newRegionOp.getBody(),
                              newRegionOp.getBody().begin());
  rewriter.replaceOp(
      regionOp, newRegionOp.getResults().take_front(regionOp->getNumResults()));

  // Update terminator.
  auto newRegionReturnOp =
      cast<IREE::Flow::ReturnOp>(newRegionOp.getBody().front().getTerminator());
  rewriter.setInsertionPoint(newRegionReturnOp);
  rewriter.replaceOpWithNewOp<IREE::Flow::ReturnOp>(newRegionReturnOp,
                                                    returnedValues);

  return newRegionOp;
}

IREE::Flow::DispatchRegionOp
makeEmptyDispatchRegion(OpBuilder &builder, Location loc, ValueRange workload) {
  OpBuilder::InsertionGuard guard(builder);

  // Create RegionOp.
  auto regionOp = builder.create<IREE::Flow::DispatchRegionOp>(
      loc, /*resultTypes=*/TypeRange(), /*dynamicDims=*/ValueRange(), workload);
  Block &body = regionOp.getBody().emplaceBlock();
  builder.setInsertionPointToStart(&body);
  builder.create<IREE::Flow::ReturnOp>(loc, ValueRange());
  return regionOp;
}

// Clone a `target` op that is preceding the given dispatch region op into the
// dispatch region.
FailureOr<Operation *>
clonePrecedingOpIntoDispatchRegion(RewriterBase &rewriter, Operation *target,
                                   IREE::Flow::DispatchRegionOp regionOp) {
  Block &body = regionOp.getBody().front();

  // Gather all uses of `target`.
  SmallVector<OpOperand *> usesInsideOfRegion;
  for (OpOperand &use : target->getUses()) {
    Operation *parentOperation = use.getOwner();
    Region *parentRegion = parentOperation->getParentRegion();

    while ((parentOperation = parentOperation->getParentOp())) {
      if (regionOp.getOperation() == parentOperation)
        break;
      parentRegion = parentOperation->getParentRegion();
    }

    if (parentOperation && &parentRegion->front() == &body) {
      usesInsideOfRegion.push_back(&use);
    }
  }

  // Clone op into dispatch region.
  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPointToStart(&body);
  Operation *newTargetOp = rewriter.clone(*target);

  // Replace all uses in the dispatch region.
  for (OpOperand *use : usesInsideOfRegion) {
    rewriter.modifyOpInPlace(use->getOwner(), [&]() {
      use->set(newTargetOp->getResult(
          llvm::cast<OpResult>(use->get()).getResultNumber()));
    });
  }

  return newTargetOp;
}

// Move a `target` op that is preceding the given dispatch region op into the
// dispatch region.
FailureOr<IREE::Flow::DispatchRegionOp>
movePrecedingOpsIntoDispatchRegion(RewriterBase &rewriter,
                                   ArrayRef<Operation *> targets,
                                   IREE::Flow::DispatchRegionOp regionOp) {
  // Values replaced by moving the `targets` into the dispatch region.
  SmallVector<Value> replacedValues;

  // List of dynamic dimensions for each new results added to the dispatch
  // region.
  SmallVector<SmallVector<Value>> dispatchOpNewResultsDynamicDims;

  // New values that are yielded from dispatch.
  SmallVector<Value> yieldedResults;

  llvm::SetVector<Operation *> targetSet;
  targetSet.insert(targets.begin(), targets.end());

  Block &body = regionOp.getBody().front();
  for (Operation *target : llvm::reverse(targets)) {
    // Clone op into dispatch region.
    OpBuilder::InsertionGuard g(rewriter);
    rewriter.setInsertionPointToStart(&body);
    Operation *clonedTarget = rewriter.clone(*target);

    // Gather all uses of `target`.
    for (auto [index, result] : llvm::enumerate(target->getResults())) {
      bool hasUsesOutsideOfRegion =
          llvm::any_of(result.getUses(), [&](OpOperand &use) {
            Operation *user = use.getOwner();
            // The use is not in
            // 1. the current dispatch
            // 2. Not in one of the targets.
            return !regionOp->isProperAncestor(user) && !targetSet.count(user);
          });
      if (hasUsesOutsideOfRegion) {
        replacedValues.push_back(result);
        yieldedResults.push_back(clonedTarget->getResult(index));
        OpBuilder::InsertionGuard guard(rewriter);
        rewriter.setInsertionPoint(target);
        SmallVector<Value> &dims =
            dispatchOpNewResultsDynamicDims.emplace_back();
        if (failed(reifyDynamicResultDims(rewriter, result, dims))) {
          return target->emitOpError(
              "failed to reify dynamic dims of result to be yielded from "
              "dispatch region");
        }
      }
    }

    rewriter.replaceOpUsesWithinBlock(target, clonedTarget->getResults(),
                                      &body);
  }

  FailureOr<IREE::Flow::DispatchRegionOp> newRegionOp =
      appendDispatchRegionResults(rewriter, regionOp, yieldedResults,
                                  dispatchOpNewResultsDynamicDims);

  if (failed(newRegionOp)) {
    return regionOp->emitOpError("failed to append results to op");
  }

  ValueRange replacements =
      newRegionOp->getResults().take_back(replacedValues.size());
  for (auto [index, replacedVal] : llvm::enumerate(replacedValues)) {
    rewriter.replaceAllUsesWith(replacedVal, replacements[index]);
  }
  for (auto target : llvm::reverse(targets)) {
    rewriter.eraseOp(target);
  }
  return newRegionOp.value();
}

// Move a `target` op that is following the given dispatch region op into the
// dispatch region.
FailureOr<IREE::Flow::DispatchRegionOp>
moveFollowingOpIntoDispatchRegion(RewriterBase &rewriter, Operation *target,
                                  IREE::Flow::DispatchRegionOp regionOp) {
  // Fail if any of the `target` operands do not dominate the dispatch region.
  mlir::DominanceInfo dominanceInfo(regionOp);
  for (Value operand : target->getOperands()) {
    Operation *definingOp = operand.getDefiningOp();
    if (definingOp && !dominanceInfo.dominates(definingOp, regionOp)) {
      return rewriter.notifyMatchFailure(
          target, "target operands do not dominate the dispatch region op.");
    }
  }

  // Values replaced by moving the `target` into the dispatch region.
  SmallVector<Value> replacedValues;

  // List of dynamic dimensions for each new results added to the dispatch
  // region.
  SmallVector<SmallVector<Value>> dispatchOpNewResultsDynamicDims;

  // New values that are yielded from dispatch.
  SmallVector<Value> yieldedResults;

  Block &body = regionOp.getBody().front();
  // Clone op into dispatch region.
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(body.getTerminator());
  Operation *clonedTarget = rewriter.clone(*target);

  // Replace any operands returned by the `regionOp` with the results yielded
  // inside of the `regionOp`.
  for (OpOperand &operand : clonedTarget->getOpOperands()) {
    if (operand.get().getDefiningOp() != regionOp) {
      continue;
    }
    auto returnOp =
        cast<IREE::Flow::ReturnOp>(regionOp.getBody().front().getTerminator());
    auto opResult = cast<OpResult>(operand.get());
    Value yieldedValue = returnOp->getOperand(opResult.getResultNumber());
    rewriter.modifyOpInPlace(clonedTarget, [&]() {
      clonedTarget->setOperand(operand.getOperandNumber(), yieldedValue);
    });
  }

  // Gather all uses of `target`.
  for (auto [index, result] : llvm::enumerate(target->getResults())) {
    replacedValues.push_back(result);
    yieldedResults.push_back(clonedTarget->getResult(index));
    rewriter.setInsertionPoint(target);
    SmallVector<Value> &dims = dispatchOpNewResultsDynamicDims.emplace_back();
    if (failed(reifyDynamicResultDims(rewriter, result, dims))) {
      return target->emitOpError(
          "failed to reify dynamic dims of result to be yielded from "
          "dispatch region");
    }
  }

  FailureOr<IREE::Flow::DispatchRegionOp> newRegionOp =
      appendDispatchRegionResults(rewriter, regionOp, yieldedResults,
                                  dispatchOpNewResultsDynamicDims);

  if (failed(newRegionOp)) {
    return regionOp->emitOpError("failed to append results to op");
  }

  ValueRange replacements =
      newRegionOp->getResults().take_back(replacedValues.size());
  for (auto [index, replacedVal] : llvm::enumerate(replacedValues)) {
    rewriter.replaceAllUsesWith(replacedVal, replacements[index]);
  }
  rewriter.eraseOp(target);
  return newRegionOp.value();
}

FailureOr<IREE::Flow::DispatchRegionOp>
wrapOpInDispatchRegion(RewriterBase &rewriter, Operation *op) {
  OpBuilder::InsertionGuard g(rewriter);

  SmallVector<Value> workload;
  bool useWorkloadCountFromDagRootMode = checkShapeIsDataDependant(op);
  if (useWorkloadCountFromDagRootMode) {
    rewriter.setInsertionPoint(op);
    workload = getWorkloadForRootOp(rewriter, op);
  }

  rewriter.setInsertionPointAfter(op);
  IREE::Flow::DispatchRegionOp regionOp =
      IREE::Flow::makeEmptyDispatchRegion(rewriter, op->getLoc(), workload);

  // Move the op into the dispatch region.
  auto newRegionOp = movePrecedingOpsIntoDispatchRegion(rewriter, op, regionOp);

  if (succeeded(newRegionOp) && useWorkloadCountFromDagRootMode) {
    auto newWorkload = newRegionOp->getWorkload();
    auto workloadTypes =
        llvm::map_to_vector(newWorkload, [](Value v) { return v.getType(); });
    auto workloadLocs =
        llvm::map_to_vector(newWorkload, [](Value v) { return v.getLoc(); });
    createWorkgroupCountFromDagRootRegion(rewriter, *newRegionOp, workloadTypes,
                                          workloadLocs);
  }

  return newRegionOp;
}

FailureOr<Operation *> hoistOutOfDispatch(RewriterBase &rewriter,
                                          Operation *op) {
  assert(op && !isNonNullAndOutsideDispatch(op) &&
         "op expected to be in a dispatch");

  // Step 1: Clone the op outside of the dispatch region.

  OpBuilder::InsertionGuard g(rewriter);
  auto dispatchRegionOp = op->getParentOfType<DispatchRegionOp>();

  // If all operands of the `op` come from outside the dispatch, then the op can
  // be hoisted out before the dispatch region. Otherwise, the op can be hoisted
  // out below the dispatch if the only users of the op are the dispatch return.
  if (llvm::none_of(op->getOperands(), [&](Value operand) {
        Operation *producer = operand.getDefiningOp();
        return producer && producer->getParentOfType<DispatchRegionOp>();
      })) {
    rewriter.setInsertionPoint(dispatchRegionOp);
  } else if (llvm::all_of(op->getUsers(), [&](Operation *user) {
               return isa<IREE::Flow::ReturnOp>(user);
             })) {
    rewriter.setInsertionPointAfter(dispatchRegionOp);
  } else {
    return rewriter.notifyMatchFailure(
        op, "op has both operands and users insided of its dispatch");
  }
  Operation *hoistedOp = rewriter.clone(*op);

  // Step 2: Replace op uses inside and outside of the dispatch region with the
  //         hoisted results.

  auto getMatchingDispatchResult =
      [&](Value result) -> std::optional<OpResult> {
    for (OpOperand &use : result.getUses()) {
      if (isa<IREE::Flow::ReturnOp>(use.getOwner())) {
        return dispatchRegionOp.getResults()[use.getOperandNumber()];
      }
    }
    return std::nullopt;
  };
  bool yieldsResults = false;
  for (OpResult result : op->getResults()) {
    Value hoistedResult = hoistedOp->getResult(result.getResultNumber());
    // Replace all results yielded by the dispatch region with the hoisted
    // op results.
    std::optional<OpResult> dispResult = getMatchingDispatchResult(result);
    if (dispResult.has_value()) {
      yieldsResults = true;
      rewriter.replaceAllUsesWith(dispResult.value(), hoistedResult);
    }
    // Replace uses inside the dispatch region.
    rewriter.replaceAllUsesWith(result, hoistedResult);
  }
  // If no results were yielded from `op`, then nothing more to do.
  if (!yieldsResults) {
    return hoistedOp;
  }

  // Step 3: Collect the new set of dispatch results and dynamic dims, and
  //         create a new dispatch region to replace the old one. The new
  //         dispatch may have duplicated results,

  // Get the new dispatch region return values and dynamic dims, excluding the
  // ones coming from the `hoistedOp`.
  auto dispatchReturnOp = cast<IREE::Flow::ReturnOp>(
      dispatchRegionOp.getBody().front().getTerminator());
  SmallVector<Value, 2> newDispatchReturnOperands;
  SmallVector<Value, 4> newDispatchResultDynamicDims;
  // Keep track of which results in the original dispatch region correspond to
  // which results in the new dispatch region with `oldDispatchResultInds`.
  SmallVector<int64_t, 2> oldDispatchResultInds;
  for (OpOperand &operand : dispatchReturnOp->getOpOperands()) {
    if (operand.get().getDefiningOp() == hoistedOp) {
      continue;
    }
    oldDispatchResultInds.push_back(operand.getOperandNumber());
    newDispatchReturnOperands.push_back(operand.get());
    auto dims =
        dispatchRegionOp.getResultDynamicDims(operand.getOperandNumber());
    newDispatchResultDynamicDims.append(dims.begin(), dims.end());
  }

  // Add the operands of the `op` to the new return values of the dispatch, and
  // add their result dynamic dims to the new result dynamic dims.
  // Save the result index in the new dispatch corresponding to each hoisted op
  // operand in `resultIndsForHoistedOperands`, so uses can be replaced later.
  SmallVector<int64_t> resultIndsForHoistedOperands;
  for (OpOperand &operand : op->getOpOperands()) {
    // Only need to yield operands defined in the dispatch region.
    if (operand.get().getParentRegion() != &dispatchRegionOp.getBody()) {
      continue;
    }

    // If the operand is already yielded by the dispatch, don't yield it again,
    // and save the result index.
    bool resultAlreadyYielded = false;
    for (auto [idx, returnOperand] :
         llvm::enumerate(newDispatchReturnOperands)) {
      if (returnOperand == operand.get()) {
        resultAlreadyYielded = true;
        resultIndsForHoistedOperands.push_back(idx);
        break;
      }
    }
    if (resultAlreadyYielded) {
      break;
    }
    resultIndsForHoistedOperands.push_back(newDispatchReturnOperands.size());
    newDispatchReturnOperands.push_back(operand.get());

    // Save operand and dynamic dims to add to the dispatch region.
    SmallVector<Value> dims;
    if (failed(reifyDynamicResultDims(rewriter, operand.get(), dims))) {
      return op->emitOpError(
          "failed to reify dynamic dims of result to be yielded from "
          "dispatch region");
    }
    newDispatchResultDynamicDims.append(dims.begin(), dims.end());
  }

  // Create the new dispatch region op. `newDispatchReturnOperands` now has all
  // the original return operands, excluding the hoisted op's results, and
  // including any new results coming from the hoisted op's old operands. The
  // `newDispatchResultDynamicDims` contains the corresponding result dynamic
  // dims for `newDispatchReturnOperands`.
  SmallVector<Type> newResultTypes =
      llvm::map_to_vector(newDispatchReturnOperands,
                          [](Value operand) { return operand.getType(); });
  rewriter.setInsertionPoint(dispatchRegionOp);
  auto newDispatchRegionOp = rewriter.create<IREE::Flow::DispatchRegionOp>(
      dispatchRegionOp->getLoc(), newResultTypes, newDispatchResultDynamicDims,
      dispatchRegionOp.getWorkload());
  rewriter.inlineRegionBefore(dispatchRegionOp.getBody(),
                              newDispatchRegionOp.getBody(),
                              newDispatchRegionOp.getBody().begin());
  // Move the workgroup count region over.
  if (!dispatchRegionOp.getWorkgroupCount().empty()) {
    Region &newWorkgroupCountRegion = newDispatchRegionOp.getWorkgroupCount();
    rewriter.inlineRegionBefore(dispatchRegionOp.getWorkgroupCount(),
                                newWorkgroupCountRegion,
                                newWorkgroupCountRegion.begin());
  }
  // Need to make a new flow.return op, since the body was copied from the
  // old dispatch region.
  auto newDispatchReturnOp = cast<IREE::Flow::ReturnOp>(
      newDispatchRegionOp.getBody().front().getTerminator());
  rewriter.setInsertionPoint(newDispatchReturnOp);
  rewriter.replaceOpWithNewOp<IREE::Flow::ReturnOp>(newDispatchReturnOp,
                                                    newDispatchReturnOperands);

  // Replace operands of the `hoistedOp` with dispatch region results. They are
  // currently using values from inside the dispatch region.
  for (auto [idx, operand] : llvm::enumerate(hoistedOp->getOperands())) {
    auto newResultIdx = resultIndsForHoistedOperands[idx];
    Value newDispatchResult = newDispatchRegionOp->getResults()[newResultIdx];
    rewriter.replaceUsesWithIf(operand, newDispatchResult,
                               [&](OpOperand &opOperand) {
                                 return opOperand.getOwner() == hoistedOp;
                               });
  }

  // Step 4: Fixup all uses. Still need to replace the operands of the hoisted
  //         op, and replace the remaining uses of the old dispatch region with
  //         the new dispatch region results.

  // Replace the uses of the original dispatch region results with the final
  // dispatch region results.
  for (auto [oldIdx, newIdx] : llvm::enumerate(oldDispatchResultInds)) {
    Value newDispatchResult = newDispatchRegionOp->getResults()[newIdx];
    Value dispatchResult = dispatchRegionOp->getResults()[oldIdx];
    rewriter.replaceAllUsesWith(dispatchResult, newDispatchResult);
  }

  return hoistedOp;
}

//===---------------------------------------------------------------------===//
// Utilities to make a dispatch region isolated from above
//===---------------------------------------------------------------------===//

/// Operations that are cloned into dispatch regions formed with other
/// operations as roots.
bool isClonableIntoDispatchOp(Operation *op) {
  // TODO(#8637): `tensor.collapse_shape` and `tensor.expand_shape` are
  // trivially clonable too, but they cause problems
  // with bufferization. Make them clonable when fixed.
  if (isa<affine::AffineApplyOp, arith::IndexCastOp, linalg::FillOp,
          tensor::EmptyOp, tensor::ExtractOp, tensor::ExtractSliceOp,
          complex::CreateOp>(op)) {
    return true;
  }
  if (LinalgExt::isBitExtendOp(op)) {
    return true;
  }
  if (clEnableGatherFusion && LinalgExt::isGatherlikeOp(op)) {
    return true;
  }
  if (isa<arith::ConstantOp>(op) || isa<complex::ConstantOp>(op)) {
    if (clInlineConstantByteLength == 0)
      return false;
    Attribute constantValueAttr;
    if (!matchPattern(op->getResult(0), m_Constant(&constantValueAttr))) {
      return false;
    }

    auto constantType = op->getResult(0).getType();
    if (llvm::isa<SplatElementsAttr>(constantValueAttr)) {
      return true;
    } else if (auto attr = llvm::dyn_cast<ElementsAttr>(constantValueAttr)) {
      auto shapedType = llvm::cast<ShapedType>(constantType);
      uint64_t estimatedByteLength =
          (shapedType.getNumElements() *
           IREE::Util::getTypeBitWidth(shapedType.getElementType())) /
          8;
      return attr.isSplat() ||
             estimatedByteLength <= clInlineConstantByteLength;
    } else if (constantType.isIntOrIndexOrFloat() ||
               isa<ComplexType>(constantType)) {
      return true;
    }
  }
  if (op->getDialect() ==
          op->getContext()->getLoadedDialect<arith::ArithDialect>() &&
      llvm::all_of(op->getOperands(),
                   [&](Value v) { return v.getType().isIntOrFloat(); }) &&
      llvm::all_of(op->getResults(),
                   [&](Value v) { return v.getType().isIntOrFloat(); })) {
    return true;
  }
  return false;
}

/// Checks if the `Value` has a use within the dispatch that is unfusable.
static bool hasUnfusableUseInDispatch(Value v, Operation *dispatchOp) {
  for (OpOperand &use : v.getUses()) {
    Operation *user = use.getOwner();

    // Do not fuse operations if they are already an operand of the
    // owner and have an index return type as that means its a shape
    // computation that needs to happen on the host.
    if (user == dispatchOp && v.getType().isIndex() &&
        isa<IREE::Flow::DispatchRegionOp>(dispatchOp)) {
      return true;
    }

    Operation *ownerWorkgroupsOp =
        user->getParentOfType<IREE::Flow::DispatchWorkgroupsOp>();
    Operation *ownerRegionOp =
        user->getParentOfType<IREE::Flow::DispatchRegionOp>();
    Operation *owner = ownerWorkgroupsOp ? ownerWorkgroupsOp : ownerRegionOp;

    // Ignore uses outside of dispatch workgroups op.
    if (owner != dispatchOp)
      continue;

    // Cannot fuse producer of `dest` with `tensor.insert_slice`.
    if (auto insertSliceUser = dyn_cast<tensor::InsertSliceOp>(user)) {
      if (insertSliceUser.getDest() == v)
        return true;
    }
  }
  return false;
}

SmallVector<Operation *>
getCloneableOps(IREE::Flow::DispatchRegionOp regionOp) {
  // Find values that are used inside of the dispatch region but defined outside
  // of the dispatch region.
  llvm::SetVector<Value> valuesDefinedAbove;
  mlir::getUsedValuesDefinedAbove(regionOp.getBody(), valuesDefinedAbove);
  if (valuesDefinedAbove.empty())
    return {};

  // Traverse the defining ops of these values (and the ops on their reverse
  // SSA use-def chain).
  SmallVector<Operation *> result;
  llvm::SetVector<Value> visited;
  SmallVector<Value> worklist;
  worklist.assign(valuesDefinedAbove.begin(), valuesDefinedAbove.end());
  while (!worklist.empty()) {
    Value outsideValue = worklist.pop_back_val();
    // Skip values that were already visited.
    if (visited.count(outsideValue))
      continue;
    visited.insert(outsideValue);

    Operation *definingOp = outsideValue.getDefiningOp();
    if (!definingOp || !IREE::Flow::isClonableIntoDispatchOp(definingOp) ||
        hasUnfusableUseInDispatch(outsideValue, regionOp)) {
      valuesDefinedAbove.insert(outsideValue);
      continue;
    }
    result.push_back(definingOp);
    worklist.append(definingOp->operand_begin(), definingOp->operand_end());
    llvm::SetVector<Value> nestedValues;
    mlir::getUsedValuesDefinedAbove(definingOp->getRegions(), nestedValues);
    worklist.append(nestedValues.begin(), nestedValues.end());
  }

  return result;
}

/// Clone producers into the dispatch region.
LogicalResult cloneProducersToRegion(RewriterBase &rewriter,
                                     IREE::Flow::DispatchRegionOp regionOp) {
  SmallVector<Operation *> cloneableOps = getCloneableOps(regionOp);
  bool sortResult = mlir::computeTopologicalSorting(cloneableOps);
  (void)sortResult;
  assert(sortResult && "could not compute topological sorting");

  for (Operation *producer : llvm::reverse(cloneableOps)) {
    if (failed(
            clonePrecedingOpIntoDispatchRegion(rewriter, producer, regionOp))) {
      return failure();
    }
  }

  return success();
}

} // namespace mlir::iree_compiler::IREE::Flow
