Raptor/src/PIM/Conversion/SpatialToPim/Common.cpp

#include "mlir/IR/ValueRange.h"

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/SCF/IR/SCF.h"

#include "llvm/ADT/STLExtras.h"

#include <cassert>
#include <limits>

#include "Common.hpp"
#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

using namespace llvm;
using namespace mlir;

namespace onnx_mlir {

FailureOr<IntegerAttr> getTensorSizeInBytesAttr(Builder& builder, Operation* anchor, mlir::Value value) {
  auto byteSize = pim::getCheckedShapedTypeSizeInBytes(cast<ShapedType>(value.getType()), anchor, "tensor byte size");
  if (failed(byteSize))
    return failure();
  return pim::getCheckedI32Attr(builder, anchor, *byteSize, "tensor byte size");
}

Operation* getEarliestUserWithinBlock(mlir::Value value) {
  auto users = value.getUsers();

  assert(!users.empty());

  Operation* earliestUser = *users.begin();
  for (auto curUser : users)
    if (curUser->isBeforeInBlock(earliestUser))
      earliestUser = curUser;

  return earliestUser;
}

SmallVector<mlir::Value> getOpOperandsSortedByUses(Operation* operation) {
  auto operandsAndUses =
    map_to_vector(operation->getOperands(), [](mlir::Value operand) -> std::pair<mlir::Value, size_t> {
      return {operand, std::distance(operand.use_begin(), operand.use_end())};
    });
  sort(operandsAndUses, [](auto a, auto b) { return a.second < b.second; });
  return map_to_vector(operandsAndUses, [](auto operandAndUse) { return operandAndUse.first; });
}

bool hasLaterUserInBlock(mlir::Value value, Operation* operation) {
  for (Operation* user : value.getUsers()) {
    if (user->getBlock() != operation->getBlock())
      return true;
    if (operation->isBeforeInBlock(user))
      return true;
  }
  return false;
}

static bool isTensorView(mlir::Value value) {
  return isa_and_nonnull<tensor::CastOp,
                         tensor::CollapseShapeOp,
                         tensor::ExpandShapeOp,
                         tensor::ExtractSliceOp,
                         tensor::ReshapeOp>(value.getDefiningOp());
}

static bool isLoopCarriedOutput(mlir::Value operand, Operation* operation) {
  auto argument = dyn_cast<BlockArgument>(operand);
  if (!argument || argument.getArgNumber() == 0 || operation->getBlock() != argument.getOwner())
    return false;
  auto loop = dyn_cast_or_null<scf::ForOp>(argument.getOwner()->getParentOp());
  return loop && cast<scf::YieldOp>(loop.getBody()->getTerminator())
                   .getOperand(argument.getArgNumber() - 1) == operation->getResult(0);
}

mlir::Value getBestOutputTensorFromOperandsOrAllocate(RewriterBase& rewriter, Operation* operation) {
  assert("Only support operations with a single result" && operation->getNumResults() == 1);
  mlir::Value result = operation->getResult(0);
  auto resultType = result.getType();
  assert("Only support result ShapedType as result type" && isa<ShapedType>(resultType));

  SmallVector<mlir::Value> operands = getOpOperandsSortedByUses(operation);
  auto validOperands = make_filter_range(operands, [operation, resultType](mlir::Value operand) {
    return operand.getType() == resultType
      && (!isa<BlockArgument>(operand) || isLoopCarriedOutput(operand, operation))
      && !operand.getDefiningOp<arith::ConstantOp>()
      && !isTensorView(operand)
      && !hasLaterUserInBlock(operand, operation);
  });
  auto bestOperand = validOperands.begin();

  if (bestOperand != validOperands.end())
    return *bestOperand;

  auto resultShapedType = cast<ShapedType>(resultType);
  rewriter.setInsertionPoint(operation);
  return tensor::EmptyOp::create(
    rewriter, operation->getLoc(), resultShapedType.getShape(), resultShapedType.getElementType());
}

LogicalResult validateFragmentAssemblyMetadata(spatial::SpatBlueprintOp blueprint,
                                               int64_t resultRank,
                                               size_t operandCount,
                                               ArrayRef<int64_t> operandIndices,
                                               ArrayRef<int64_t> sourceOffsets,
                                               ArrayRef<int64_t> flatOffsets,
                                               ArrayRef<int64_t> flatSizes,
                                               ArrayRef<int64_t> flatStrides) {
  if (operandIndices.size() != sourceOffsets.size())
    return blueprint.emitOpError("fragment assembly operand index and source offset counts must match");
  if (flatOffsets.size() != flatSizes.size())
    return blueprint.emitOpError("fragment assembly offset and size arrays must have matching lengths");
  if (flatStrides.size() != flatOffsets.size())
    return blueprint.emitOpError("fragment assembly stride and offset arrays must have matching lengths");
  if (flatOffsets.size() != operandIndices.size() * static_cast<size_t>(resultRank))
    return blueprint.emitOpError("fragment assembly metadata must provide one rank-sized offset/size/stride tuple per fragment");

  for (auto [fragmentIndex, operandIndex] : llvm::enumerate(operandIndices)) {
    if (operandIndex < 0 || operandIndex >= static_cast<int64_t>(operandCount))
      return blueprint.emitOpError("fragment assembly operand index is out of range");
    if (sourceOffsets[fragmentIndex] < 0)
      return blueprint.emitOpError("fragment assembly source offsets must be nonnegative");
  }

  return success();
}

static SmallVector<int64_t, 4> expandFlatElementIndex(int64_t flatIndex, ArrayRef<int64_t> shape) {
  SmallVector<int64_t, 4> indices(shape.size(), 0);
  for (int64_t dim = static_cast<int64_t>(shape.size()) - 1; dim >= 0; --dim) {
    indices[dim] = flatIndex % shape[dim];
    flatIndex /= shape[dim];
  }
  return indices;
}

FailureOr<SmallVector<int64_t, 4>>
getStaticSliceOffsetsForElementOffset(Operation* anchor,
                                      ShapedType sourceType,
                                      ArrayRef<int64_t> fragmentShape,
                                      int64_t sourceElementOffset,
                                      StringRef fieldName) {
  if (!sourceType.hasStaticShape())
    return (anchor->emitOpError() << fieldName << " requires a static source shape"), failure();
  if (sourceElementOffset < 0)
    return (anchor->emitOpError() << fieldName << " requires a nonnegative source element offset"), failure();
  if (sourceType.getRank() != static_cast<int64_t>(fragmentShape.size()))
    return (anchor->emitOpError() << fieldName << " requires fragment rank to match source rank"), failure();

  int64_t sourceElementCount = sourceType.getNumElements();
  int64_t fragmentElementCount = 1;
  for (int64_t dim = 0; dim < sourceType.getRank(); ++dim) {
    if (fragmentShape[dim] < 0)
      return (anchor->emitOpError() << fieldName << " requires nonnegative fragment sizes"), failure();
    fragmentElementCount *= fragmentShape[dim];
  }
  if (sourceElementOffset + fragmentElementCount > sourceElementCount)
    return (anchor->emitOpError() << fieldName << " exceeds the source tensor bounds"), failure();

  SmallVector<int64_t, 4> sliceOffsets = expandFlatElementIndex(sourceElementOffset, sourceType.getShape());
  for (int64_t dim = 0; dim < sourceType.getRank(); ++dim) {
    if (sliceOffsets[dim] + fragmentShape[dim] > sourceType.getDimSize(dim))
      return (anchor->emitOpError() << fieldName << " does not describe a valid unit-stride slice"), failure();
  }
  return sliceOffsets;
}

LogicalResult
forEachContiguousDestinationChunk(ArrayRef<int64_t> destShape,
                                  ArrayRef<int64_t> baseOffsets,
                                  ArrayRef<int64_t> sizes,
                                  llvm::function_ref<LogicalResult(ArrayRef<int64_t>, int64_t, int64_t)> callback) {
  int64_t rank = static_cast<int64_t>(sizes.size());
  int64_t suffixStart = rank - 1;
  while (suffixStart > 0 && sizes[suffixStart] == destShape[suffixStart])
    --suffixStart;
  if (sizes[suffixStart] == destShape[suffixStart] && suffixStart == 0)
    suffixStart = 0;
  else
    ++suffixStart;

  int64_t chunkElements = 1;
  for (int64_t dim = suffixStart; dim < rank; ++dim)
    chunkElements *= sizes[dim];

  SmallVector<int64_t, 4> prefixExtents(sizes.begin(), sizes.begin() + suffixStart);
  SmallVector<int64_t, 4> current(prefixExtents.size(), 0);
  int64_t sourceChunkOrdinal = 0;

  auto visit = [&](auto&& visit, int64_t dim) -> LogicalResult {
    if (dim == static_cast<int64_t>(prefixExtents.size())) {
      SmallVector<int64_t, 4> chunkOffsets(baseOffsets.begin(), baseOffsets.end());
      for (int64_t prefixDim = 0; prefixDim < static_cast<int64_t>(current.size()); ++prefixDim)
        chunkOffsets[prefixDim] += current[prefixDim];
      if (failed(callback(chunkOffsets, sourceChunkOrdinal * chunkElements, chunkElements)))
        return failure();
      ++sourceChunkOrdinal;
      return success();
    }

    for (int64_t index = 0; index < prefixExtents[dim]; ++index) {
      current[dim] = index;
      if (failed(visit(visit, dim + 1)))
        return failure();
    }
    return success();
  };

  if (prefixExtents.empty())
    return callback(baseOffsets, 0, chunkElements);
  return visit(visit, 0);
}

static mlir::Value
createSteppedOffset(OpBuilder& builder, Location loc, mlir::Value start, mlir::Value index,
                    int64_t stepBytes, Operation *constantAnchor) {
  if (stepBytes == 0)
    return start;
  return createOrFoldAffineApply(
    builder, loc, builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1) * stepBytes,
    ValueRange {start, index}, constantAnchor);
}

static mlir::Value createIndexedOffset(OpBuilder& builder,
                                       Location loc,
                                       mlir::Value indexArg,
                                       ArrayRef<int64_t> values,
                                       Operation *constantAnchor) {
  assert(!values.empty() && "expected lane-indexed values");
  if (llvm::all_of(values.drop_front(), [&](int64_t value) { return value == values.front(); }))
    return getOrCreateIndexConstant(builder, constantAnchor, values.front());

  if (values.size() >= 2) {
    int64_t step = values[1] - values[0];
    bool arithmetic = llvm::all_of(llvm::seq<size_t>(2, values.size()), [&](size_t index) {
      return values[index] == values.front() + static_cast<int64_t>(index) * step;
    });
    if (arithmetic) {
      return createOrFoldAffineApply(
        builder, loc, builder.getAffineDimExpr(0) * step + values.front(),
        ValueRange {indexArg}, constantAnchor);
    }
  }

  RankedTensorType tableType = RankedTensorType::get(
    {static_cast<int64_t>(values.size())}, builder.getI64Type());
  DenseElementsAttr tableAttr = DenseElementsAttr::get(tableType, values);
  mlir::Value table = getOrCreateConstant(builder, constantAnchor, tableAttr, tableType);
  mlir::Value selected = tensor::ExtractOp::create(builder, loc, table, ValueRange {indexArg});
  return arith::IndexCastOp::create(builder, loc, builder.getIndexType(), selected).getResult();
}

struct FragmentAssemblyCopyRunFamily {
  FragmentAssemblyCopyRun prototype;
  SmallVector<int64_t, 8> sourceRunStartDeltas;
  SmallVector<int64_t, 8> hostRunStartDeltas;
};

static bool computeUniformRunStartDelta(ArrayRef<int64_t> prototypeStarts,
                                        ArrayRef<int64_t> runStarts,
                                        int64_t& delta) {
  if (prototypeStarts.size() != runStarts.size() || prototypeStarts.empty())
    return false;

  delta = runStarts.front() - prototypeStarts.front();
  return llvm::all_of(llvm::zip_equal(prototypeStarts, runStarts), [&](auto pair) {
    auto [prototypeStart, runStart] = pair;
    return runStart - prototypeStart == delta;
  });
}

static bool canMergeFragmentAssemblyCopyRunIntoFamily(const FragmentAssemblyCopyRunFamily& family,
                                                      const FragmentAssemblyCopyRun& run,
                                                      int64_t& sourceRunStartDelta,
                                                      int64_t& hostRunStartDelta) {
  const FragmentAssemblyCopyRun& prototype = family.prototype;
  if (prototype.source != run.source || prototype.sourceType != run.sourceType
      || prototype.hostTargetIndex != run.hostTargetIndex || prototype.count != run.count
      || prototype.sourceStepBytes != run.sourceStepBytes || prototype.hostStepBytes != run.hostStepBytes
      || prototype.byteSize != run.byteSize)
    return false;

  if (!computeUniformRunStartDelta(prototype.sourceStartBytesByLane, run.sourceStartBytesByLane, sourceRunStartDelta))
    return false;
  return computeUniformRunStartDelta(prototype.hostStartBytesByLane, run.hostStartBytesByLane, hostRunStartDelta);
}

static SmallVector<FragmentAssemblyCopyRunFamily, 8>
groupFragmentAssemblyCopyRunFamilies(ArrayRef<FragmentAssemblyCopyRun> runs) {
  auto compareRunStarts = [](ArrayRef<int64_t> lhs, ArrayRef<int64_t> rhs) {
    return std::lexicographical_compare(lhs.begin(), lhs.end(), rhs.begin(), rhs.end());
  };

  SmallVector<FragmentAssemblyCopyRun, 8> sortedRuns(runs.begin(), runs.end());
  llvm::sort(sortedRuns, [&](const FragmentAssemblyCopyRun& lhs, const FragmentAssemblyCopyRun& rhs) {
    if (lhs.hostTargetIndex != rhs.hostTargetIndex)
      return lhs.hostTargetIndex < rhs.hostTargetIndex;
    if (lhs.source != rhs.source)
      return lhs.source.getAsOpaquePointer() < rhs.source.getAsOpaquePointer();
    if (lhs.byteSize != rhs.byteSize)
      return lhs.byteSize < rhs.byteSize;
    if (lhs.count != rhs.count)
      return lhs.count < rhs.count;
    if (lhs.sourceStepBytes != rhs.sourceStepBytes)
      return lhs.sourceStepBytes < rhs.sourceStepBytes;
    if (lhs.hostStepBytes != rhs.hostStepBytes)
      return lhs.hostStepBytes < rhs.hostStepBytes;
    if (compareRunStarts(lhs.sourceStartBytesByLane, rhs.sourceStartBytesByLane))
      return true;
    if (compareRunStarts(rhs.sourceStartBytesByLane, lhs.sourceStartBytesByLane))
      return false;
    return compareRunStarts(lhs.hostStartBytesByLane, rhs.hostStartBytesByLane);
  });

  SmallVector<FragmentAssemblyCopyRunFamily, 8> families;
  for (const FragmentAssemblyCopyRun& run : sortedRuns) {
    int64_t sourceRunStartDelta = 0;
    int64_t hostRunStartDelta = 0;
    if (!families.empty()
        && canMergeFragmentAssemblyCopyRunIntoFamily(
          families.back(), run, sourceRunStartDelta, hostRunStartDelta)) {
      families.back().sourceRunStartDeltas.push_back(sourceRunStartDelta);
      families.back().hostRunStartDeltas.push_back(hostRunStartDelta);
      continue;
    }

    FragmentAssemblyCopyRunFamily family;
    family.prototype = run;
    family.sourceRunStartDeltas.push_back(0);
    family.hostRunStartDeltas.push_back(0);
    families.push_back(std::move(family));
  }

  return families;
}

FailureOr<SmallVector<FragmentAssemblyCopyRun, 8>>
groupFragmentAssemblyCopyRuns(ArrayRef<FragmentAssemblyCopy> copies, uint32_t laneCount) {
  if (laneCount == 0)
    return failure();

  struct LaneLocalCopyRun {
    FragmentAssemblyCopyRun run;
    int64_t lane = 0;
  };

  SmallVector<FragmentAssemblyCopy, 8> sortedCopies(copies.begin(), copies.end());
  llvm::sort(sortedCopies, [](const FragmentAssemblyCopy& lhs, const FragmentAssemblyCopy& rhs) {
    if (lhs.hostTargetIndex != rhs.hostTargetIndex)
      return lhs.hostTargetIndex < rhs.hostTargetIndex;
    if (lhs.source != rhs.source)
      return lhs.source.getAsOpaquePointer() < rhs.source.getAsOpaquePointer();
    if (lhs.lane != rhs.lane)
      return lhs.lane < rhs.lane;
    if (lhs.byteSize != rhs.byteSize)
      return lhs.byteSize < rhs.byteSize;
    if (lhs.sourceByteOffset != rhs.sourceByteOffset)
      return lhs.sourceByteOffset < rhs.sourceByteOffset;
    return lhs.hostByteOffset < rhs.hostByteOffset;
  });

  SmallVector<LaneLocalCopyRun, 8> laneRuns;
  for (const FragmentAssemblyCopy& copy : sortedCopies) {
    if (copy.lane < 0 || copy.lane >= static_cast<int64_t>(laneCount))
      return failure();

    if (!laneRuns.empty()) {
      LaneLocalCopyRun& laneRun = laneRuns.back();
      FragmentAssemblyCopyRun& run = laneRun.run;
      if (run.source == copy.source && run.sourceType == copy.sourceType
          && run.hostTargetIndex == copy.hostTargetIndex && laneRun.lane == copy.lane && run.byteSize == copy.byteSize
          && run.sourceStartBytesByLane.size() == 1 && run.hostStartBytesByLane.size() == 1) {
        int64_t previousSourceOffset = run.sourceStartBytesByLane.front() + (run.count - 1) * run.sourceStepBytes;
        int64_t previousHostOffset = run.hostStartBytesByLane.front() + (run.count - 1) * run.hostStepBytes;
        int64_t sourceDelta = copy.sourceByteOffset - previousSourceOffset;
        int64_t hostDelta = copy.hostByteOffset - previousHostOffset;
        if (run.count == 1) {
          run.sourceStepBytes = sourceDelta;
          run.hostStepBytes = hostDelta;
          ++run.count;
          continue;
        }
        if (run.sourceStepBytes == sourceDelta && run.hostStepBytes == hostDelta) {
          ++run.count;
          continue;
        }
      }
    }

    LaneLocalCopyRun laneRun;
    laneRun.run.source = copy.source;
    laneRun.run.sourceType = copy.sourceType;
    laneRun.run.hostTargetIndex = copy.hostTargetIndex;
    laneRun.run.count = 1;
    laneRun.run.byteSize = copy.byteSize;
    laneRun.run.sourceStartBytesByLane.push_back(copy.sourceByteOffset);
    laneRun.run.hostStartBytesByLane.push_back(copy.hostByteOffset);
    laneRun.lane = copy.lane;
    laneRuns.push_back(std::move(laneRun));
  }

  SmallVector<FragmentAssemblyCopyRun, 8> mergedRuns;
  for (const LaneLocalCopyRun& laneRun : laneRuns) {
    size_t laneIndex = static_cast<size_t>(laneRun.lane);
    auto mergedIt = llvm::find_if(mergedRuns, [&](const FragmentAssemblyCopyRun& run) {
      return run.source == laneRun.run.source && run.sourceType == laneRun.run.sourceType
             && run.hostTargetIndex == laneRun.run.hostTargetIndex && run.count == laneRun.run.count
             && run.byteSize == laneRun.run.byteSize && run.sourceStepBytes == laneRun.run.sourceStepBytes
             && run.hostStepBytes == laneRun.run.hostStepBytes && laneIndex < run.sourceStartBytesByLane.size()
             && run.sourceStartBytesByLane[laneIndex] == std::numeric_limits<int64_t>::min();
    });

    if (mergedIt == mergedRuns.end()) {
      FragmentAssemblyCopyRun merged = laneRun.run;
      merged.sourceStartBytesByLane.assign(laneCount, std::numeric_limits<int64_t>::min());
      merged.hostStartBytesByLane.assign(laneCount, std::numeric_limits<int64_t>::min());
      merged.sourceStartBytesByLane[laneIndex] = laneRun.run.sourceStartBytesByLane.front();
      merged.hostStartBytesByLane[laneIndex] = laneRun.run.hostStartBytesByLane.front();
      mergedRuns.push_back(std::move(merged));
      continue;
    }

    mergedIt->sourceStartBytesByLane[laneIndex] = laneRun.run.sourceStartBytesByLane.front();
    mergedIt->hostStartBytesByLane[laneIndex] = laneRun.run.hostStartBytesByLane.front();
  }

  for (const FragmentAssemblyCopyRun& run : mergedRuns) {
    if (llvm::any_of(run.sourceStartBytesByLane,
                     [](int64_t value) { return value == std::numeric_limits<int64_t>::min(); }))
      return failure();
    if (llvm::any_of(run.hostStartBytesByLane,
                     [](int64_t value) { return value == std::numeric_limits<int64_t>::min(); }))
      return failure();
  }

  return mergedRuns;
}

static FailureOr<mlir::Value> emitFragmentAssemblyCopyRun(OpBuilder& builder,
                                                          Location loc,
                                                          const FragmentAssemblyCopyRun& run,
                                                          mlir::Value hostTarget,
                                                          Operation* anchor,
                                                          std::optional<mlir::Value> laneArg,
                                                          mlir::Value baseHostOffset,
                                                          mlir::Value sourceRunStartDelta = {},
                                                          mlir::Value hostRunStartDelta = {}) {
  auto sizeAttr = pim::getCheckedI32Attr(builder, anchor, run.byteSize, "fragment assembly host copy byte size");
  if (failed(sizeAttr))
    return failure();

  mlir::Value hostStart;
  mlir::Value sourceStart;
  if (laneArg) {
    hostStart = createIndexedOffset(builder, loc, *laneArg, run.hostStartBytesByLane, anchor);
    sourceStart = createIndexedOffset(builder, loc, *laneArg, run.sourceStartBytesByLane, anchor);
  } else {
    hostStart = getOrCreateIndexConstant(builder, anchor, run.hostStartBytesByLane.front());
    sourceStart = getOrCreateIndexConstant(builder, anchor, run.sourceStartBytesByLane.front());
  }

  if (hostRunStartDelta)
    hostStart = arith::AddIOp::create(builder, loc, hostStart, hostRunStartDelta).getResult();
  if (sourceRunStartDelta)
    sourceStart = arith::AddIOp::create(builder, loc, sourceStart, sourceRunStartDelta).getResult();
  if (baseHostOffset)
    hostStart = arith::AddIOp::create(builder, loc, baseHostOffset, hostStart).getResult();

  if (run.count == 1) {
    return pim::PimMemCopyDevToHostOp::create(builder,
                                              loc,
                                              hostTarget.getType(),
                                              hostStart,
                                              sourceStart,
                                              hostTarget,
                                              run.source,
                                              *sizeAttr)
      .getOutput();
  }

  mlir::Value lowerBound = getOrCreateIndexConstant(builder, anchor, 0);
  mlir::Value upperBound = getOrCreateIndexConstant(builder, anchor, run.count);
  mlir::Value step = getOrCreateIndexConstant(builder, anchor, 1);
  FailureOr<NormalizedLoopResult> loop = buildNormalizedScfFor(
    builder,
    loc,
    lowerBound,
    upperBound,
    step,
    ValueRange {hostTarget},
    [&](OpBuilder& loopBuilder,
        Location bodyLoc,
        mlir::Value flatIndex,
        ValueRange iterArgs,
        SmallVectorImpl<mlir::Value>& yielded) {
      mlir::Value hostOffset = createSteppedOffset(
        loopBuilder, bodyLoc, hostStart, flatIndex, run.hostStepBytes, anchor);
      mlir::Value sourceOffset =
        createSteppedOffset(loopBuilder, bodyLoc, sourceStart, flatIndex, run.sourceStepBytes, anchor);
      mlir::Value copied =
        pim::PimMemCopyDevToHostOp::create(loopBuilder,
                                           bodyLoc,
                                           iterArgs.front().getType(),
                                           hostOffset,
                                           sourceOffset,
                                           iterArgs.front(),
                                           run.source,
                                           *sizeAttr)
          .getOutput();
      yielded.push_back(copied);
      return success();
    });
  if (failed(loop))
    return failure();
  return loop->results.front();
}

static FailureOr<mlir::Value> emitFragmentAssemblyCopyRunFamily(OpBuilder& builder,
                                                                Location loc,
                                                                const FragmentAssemblyCopyRunFamily& family,
                                                                mlir::Value hostTarget,
                                                                Operation* anchor,
                                                                std::optional<mlir::Value> laneArg,
                                                                mlir::Value baseHostOffset) {
  if (family.sourceRunStartDeltas.size() == 1)
    return emitFragmentAssemblyCopyRun(
      builder, loc, family.prototype, hostTarget, anchor, laneArg, baseHostOffset);

  mlir::Value lowerBound = getOrCreateIndexConstant(builder, anchor, 0);
  mlir::Value upperBound = getOrCreateIndexConstant(builder, anchor, family.sourceRunStartDeltas.size());
  mlir::Value step = getOrCreateIndexConstant(builder, anchor, 1);
  FailureOr<NormalizedLoopResult> outerLoop = buildNormalizedScfFor(
    builder,
    loc,
    lowerBound,
    upperBound,
    step,
    ValueRange {hostTarget},
    [&](OpBuilder& loopBuilder,
        Location bodyLoc,
        mlir::Value runIndex,
        ValueRange iterArgs,
        SmallVectorImpl<mlir::Value>& yielded) {
      mlir::Value sourceRunStartDelta =
        createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.sourceRunStartDeltas, anchor);
      mlir::Value hostRunStartDelta =
        createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.hostRunStartDeltas, anchor);
      FailureOr<mlir::Value> copied = emitFragmentAssemblyCopyRun(loopBuilder,
                                                                  bodyLoc,
                                                                  family.prototype,
                                                                  iterArgs.front(),
                                                                  anchor,
                                                                  laneArg,
                                                                  baseHostOffset,
                                                                  sourceRunStartDelta,
                                                                  hostRunStartDelta);
      if (failed(copied))
        return failure();
      yielded.push_back(*copied);
      return success();
    });
  if (failed(outerLoop))
    return failure();
  return outerLoop->results.front();
}

FailureOr<mlir::Value> emitFragmentAssemblyCopyRuns(IRRewriter& rewriter,
                                                    Location loc,
                                                    ArrayRef<FragmentAssemblyCopyRun> runs,
                                                    mlir::Value hostTarget,
                                                    Operation* anchor,
                                                    std::optional<mlir::Value> laneArg,
                                                    mlir::Value baseHostOffset) {
  for (const FragmentAssemblyCopyRunFamily& family : groupFragmentAssemblyCopyRunFamilies(runs)) {
    FailureOr<mlir::Value> updatedHostTarget =
      emitFragmentAssemblyCopyRunFamily(rewriter, loc, family, hostTarget, anchor, laneArg, baseHostOffset);
    if (failed(updatedHostTarget))
      return failure();
    hostTarget = *updatedHostTarget;
  }

  return hostTarget;
}

} // namespace onnx_mlir