Raptor/src/PIM/Dialect/Spatial/SpatialOpsCanonicalization.cpp

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Block.h"
#include "mlir/IR/IRMapping.h"

#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/LogicalResult.h"

#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"

using namespace mlir;

namespace onnx_mlir {
namespace spatial {

template <typename ComputeOpTy>
LogicalResult foldComputeLike(ComputeOpTy compute, ::llvm::SmallVectorImpl<::mlir::OpFoldResult>& results) {
  Block& block = compute.getBody().front();
  if (!llvm::hasSingleElement(block))
    return failure();

  auto yieldOp = dyn_cast<SpatYieldOp>(block.front());
  if (!yieldOp)
    return failure();

  for (Value yieldedValue : yieldOp.getOperands()) {
    if (auto blockArg = dyn_cast<BlockArgument>(yieldedValue)) {
      if (blockArg.getOwner() == &block) {
        results.push_back(compute.getOperand(blockArg.getArgNumber()));
        continue;
      }
    }
    results.push_back(yieldedValue);
  }
  return success();
}

LogicalResult SpatGraphCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::mlir::OpFoldResult>& results) {
  return foldComputeLike(*this, results);
}

LogicalResult SpatScheduledCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::mlir::OpFoldResult>& results) {
  return foldComputeLike(*this, results);
}

template <typename ScalarComputeOpTy>
static ScalarComputeOpTy createEmptyScalarCompute(PatternRewriter& rewriter,
                                                  Location loc,
                                                  TypeRange resultTypes,
                                                  ValueRange weights,
                                                  ValueRange inputs) {
  auto computeOp = ScalarComputeOpTy::create(rewriter, loc, resultTypes, weights, inputs);
  SmallVector<Type> blockArgTypes;
  SmallVector<Location> blockArgLocs;
  blockArgTypes.reserve(weights.size() + inputs.size());
  blockArgLocs.reserve(weights.size() + inputs.size());
  for (Value weight : weights) {
    blockArgTypes.push_back(weight.getType());
    blockArgLocs.push_back(weight.getLoc());
  }
  for (Value input : inputs) {
    blockArgTypes.push_back(input.getType());
    blockArgLocs.push_back(input.getLoc());
  }
  rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), blockArgTypes, blockArgLocs);
  rewriter.setInsertionPointToStart(&computeOp.getBody().front());
  return computeOp;
}

static SmallVector<OpFoldResult> remapMixedOffsets(ArrayRef<OpFoldResult> mixedOffsets, IRMapping& mapper) {
  SmallVector<OpFoldResult> remapped;
  remapped.reserve(mixedOffsets.size());
  for (OpFoldResult ofr : mixedOffsets) {
    if (auto value = dyn_cast<Value>(ofr))
      remapped.push_back(cast<Value>(mapper.lookupOrDefault(value)));
    else
      remapped.push_back(cast<Attribute>(ofr));
  }
  return remapped;
}

static SmallVector<Value> createEmptyResults(PatternRewriter& rewriter, Location loc, TypeRange resultTypes) {
  SmallVector<Value> resultValues;
  resultValues.reserve(resultTypes.size());
  for (Type resultType : resultTypes) {
    auto tensorType = dyn_cast<RankedTensorType>(resultType);
    if (!tensorType || !tensorType.hasStaticShape())
      return {};
    resultValues.push_back(tensor::EmptyOp::create(rewriter, loc, tensorType.getShape(), tensorType.getElementType()));
  }
  return resultValues;
}

template <typename ScalarComputeOpTy, typename ComputeBatchOpTy>
static void copyCanonicalizedBatchAttrs(ScalarComputeOpTy compute, ComputeBatchOpTy batch, PatternRewriter& rewriter) {
  for (NamedAttribute attr : batch->getAttrs()) {
    if (attr.getName() == batch.getOperandSegmentSizesAttrName() || attr.getName() == batch.getLaneCountAttrName()
        || attr.getName() == onnx_mlir::kCoreIdsAttrName)
      continue;
    compute->setAttr(attr.getName(), attr.getValue());
  }
  if constexpr (std::is_same_v<ComputeBatchOpTy, SpatScheduledComputeBatch>) {
    if (auto coreIds = batch->template getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName)) {
      assert(coreIds.size() == 1 && "single-lane scheduled compute_batch canonicalization expects exactly one core id");
      compute->setAttr(onnx_mlir::kCoreIdAttrName, rewriter.getI32IntegerAttr(coreIds.asArrayRef().front()));
    }
  }
}

template <typename ComputeBatchOpTy, typename ScalarComputeOpTy>
struct CanonicalizeSingleLaneComputeBatchPattern : OpRewritePattern<ComputeBatchOpTy> {
  using OpRewritePattern<ComputeBatchOpTy>::OpRewritePattern;

  LogicalResult matchAndRewrite(ComputeBatchOpTy compute, PatternRewriter& rewriter) const override {
    if (compute.getLaneCount() != 1)
      return rewriter.notifyMatchFailure(compute, "lane count is not 1");

    Block& oldBlock = compute.getBody().front();
    auto oldLaneArg = compute.getLaneArgument();
    if (!oldLaneArg)
      return rewriter.notifyMatchFailure(compute, "missing compute_batch lane block argument");

    rewriter.setInsertionPointAfter(compute);
    auto newCompute =
      createEmptyScalarCompute<ScalarComputeOpTy>(rewriter, compute.getLoc(), compute.getResultTypes(), compute.getWeights(), compute.getInputs());
    copyCanonicalizedBatchAttrs(newCompute, compute, rewriter);
    auto* newBlock = &newCompute.getBody().front();
    rewriter.setInsertionPointToStart(newBlock);

    IRMapping mapper;
    Value zero = arith::ConstantIndexOp::create(rewriter, compute.getLoc(), 0);
    mapper.map(*oldLaneArg, zero);
    for (auto [index, weight] : llvm::enumerate(compute.getWeights())) {
      auto oldArg = compute.getWeightArgument(index);
      auto newArg = newCompute.getWeightArgument(index);
      if (!oldArg || !newArg)
        return rewriter.notifyMatchFailure(compute, "missing rewritten compute weight block argument");
      mapper.map(*oldArg, *newArg);
    }
    for (auto [index, input] : llvm::enumerate(compute.getInputs())) {
      auto oldArg = compute.getInputArgument(index);
      auto newArg = newCompute.getInputArgument(index);
      if (!oldArg || !newArg)
        return rewriter.notifyMatchFailure(compute, "missing rewritten compute input block argument");
      mapper.map(*oldArg, *newArg);
    }

    SmallVector<Value> resultValues = createEmptyResults(rewriter, compute.getLoc(), compute.getResultTypes());
    if (resultValues.size() != compute.getNumResults())
      return rewriter.notifyMatchFailure(compute, "single-lane compute_batch canonicalization requires static ranked results");
    for (auto [index, resultValue] : llvm::enumerate(resultValues)) {
      auto oldOutputArg = compute.getOutputArgument(index);
      if (!oldOutputArg)
        return rewriter.notifyMatchFailure(compute, "missing compute_batch output block argument");
      mapper.map(*oldOutputArg, resultValue);
    }

    auto oldInParallel = dyn_cast<SpatInParallelOp>(oldBlock.getTerminator());
    auto oldYield = dyn_cast<SpatYieldOp>(oldBlock.getTerminator());
    for (Operation& op : oldBlock.without_terminator())
      rewriter.clone(op, mapper);

    if (oldYield) {
      SpatYieldOp::create(rewriter, oldYield.getLoc(), ValueRange {});
      rewriter.replaceOp(compute, newCompute.getResults());
      return success();
    }
    if (!oldInParallel)
      return rewriter.notifyMatchFailure(compute, "expected spat.in_parallel or empty spat.yield terminator");

    DenseMap<BlockArgument, size_t> outputIndexByArg;
    for (size_t index = 0; index < compute.getNumResults(); ++index) {
      auto oldOutputArg = compute.getOutputArgument(index);
      if (!oldOutputArg)
        return rewriter.notifyMatchFailure(compute, "missing compute_batch output block argument");
      outputIndexByArg[*oldOutputArg] = index;
    }

    for (Operation& op : oldInParallel.getRegion().front()) {
      auto insertSlice = dyn_cast<tensor::ParallelInsertSliceOp>(&op);
      if (!insertSlice)
        return rewriter.notifyMatchFailure(compute, "expected only tensor.parallel_insert_slice in spat.in_parallel");
      auto oldDest = dyn_cast<BlockArgument>(insertSlice.getDest());
      if (!oldDest)
        return rewriter.notifyMatchFailure(compute, "expected tensor.parallel_insert_slice destination to be a block argument");
      auto resultIndexIt = outputIndexByArg.find(oldDest);
      if (resultIndexIt == outputIndexByArg.end())
        return rewriter.notifyMatchFailure(compute, "unexpected tensor.parallel_insert_slice destination");
      size_t resultIndex = resultIndexIt->second;
      Value remappedSource = mapper.lookupOrDefault(insertSlice.getSource());
      auto remappedOffsets = remapMixedOffsets(insertSlice.getMixedOffsets(), mapper);
      auto remappedSizes = remapMixedOffsets(insertSlice.getMixedSizes(), mapper);
      auto remappedStrides = remapMixedOffsets(insertSlice.getMixedStrides(), mapper);
      resultValues[resultIndex] = tensor::InsertSliceOp::create(rewriter,
                                                                insertSlice.getLoc(),
                                                                remappedSource,
                                                                resultValues[resultIndex],
                                                                remappedOffsets,
                                                                remappedSizes,
                                                                remappedStrides)
                                    .getResult();
    }

    SpatYieldOp::create(rewriter, oldInParallel.getLoc(), resultValues);
    rewriter.replaceOp(compute, newCompute.getResults());
    return success();
  }
};

void SpatGraphComputeBatch::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) {
  results.add<CanonicalizeSingleLaneComputeBatchPattern<SpatGraphComputeBatch, SpatGraphCompute>>(context);
}

void SpatScheduledComputeBatch::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) {
  results.add<CanonicalizeSingleLaneComputeBatchPattern<SpatScheduledComputeBatch, SpatScheduledCompute>>(context);
}

} // namespace spatial
} // namespace onnx_mlir