Raptor/src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp

#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Transforms/DialectConversion.h"

#include "llvm/ADT/SmallVector.h"

#include <algorithm>
#include <numeric>
#include <optional>
#include <type_traits>

#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"

using namespace mlir;

namespace onnx_mlir {
namespace {

struct ReduceMeanSemantics {
  SmallVector<int64_t> axes;
  int64_t keepdims = 1;
  bool isIdentity = false;
};

static bool isNoneValueLike(Value value) { return isa_and_nonnull<ONNXNoneOp>(value.getDefiningOp()); }

static FailureOr<SmallVector<int64_t>> getConstantIntValues(Value value) {
  auto denseAttr = dyn_cast_or_null<DenseIntElementsAttr>(getHostConstDenseElementsAttr(value));
  if (!denseAttr)
    return failure();
  return SmallVector<int64_t>(denseAttr.getValues<int64_t>().begin(), denseAttr.getValues<int64_t>().end());
}

static FailureOr<SmallVector<int64_t>> normalizeAxesChecked(ArrayRef<int64_t> axes, int64_t rank) {
  SmallVector<int64_t> normalizedAxes;
  normalizedAxes.reserve(axes.size());
  for (int64_t axis : axes) {
    auto normalizedAxis = normalizeAxisChecked(axis, rank);
    if (failed(normalizedAxis))
      return failure();
    normalizedAxes.push_back(*normalizedAxis);
  }
  llvm::sort(normalizedAxes);
  normalizedAxes.erase(std::unique(normalizedAxes.begin(), normalizedAxes.end()), normalizedAxes.end());
  return normalizedAxes;
}

template <typename ReduceMeanOp, typename ReduceMeanOpAdaptor>
static FailureOr<ReduceMeanSemantics>
getReduceMeanSemantics(ReduceMeanOp reduceMeanOp, ReduceMeanOpAdaptor adaptor, int64_t inputRank) {
  ReduceMeanSemantics semantics;
  semantics.keepdims = reduceMeanOp.getKeepdims();

  if constexpr (std::is_same_v<ReduceMeanOp, ONNXReduceMeanV13Op>) {
    auto axes = onnx_mlir::normalizeAxesChecked(std::optional<ArrayAttr>(reduceMeanOp.getAxesAttr()), inputRank);
    if (failed(axes))
      return failure();
    semantics.axes = std::move(*axes);
    return semantics;
  }
  else {
    if (isNoneValueLike(adaptor.getAxes())) {
      if (reduceMeanOp.getNoopWithEmptyAxes() != 0) {
        semantics.isIdentity = true;
        return semantics;
      }

      semantics.axes.reserve(inputRank);
      for (int64_t axis = 0; axis < inputRank; ++axis)
        semantics.axes.push_back(axis);
      return semantics;
    }

    auto axes = getConstantIntValues(adaptor.getAxes());
    if (failed(axes))
      return failure();

    if (axes->empty()) {
      if (reduceMeanOp.getNoopWithEmptyAxes() != 0) {
        semantics.isIdentity = true;
        return semantics;
      }

      semantics.axes.reserve(inputRank);
      for (int64_t axis = 0; axis < inputRank; ++axis)
        semantics.axes.push_back(axis);
      return semantics;
    }

    auto normalizedAxes = normalizeAxesChecked(*axes, inputRank);
    if (failed(normalizedAxes))
      return failure();
    semantics.axes = std::move(*normalizedAxes);
    return semantics;
  }
}

static SmallVector<bool> buildReducedAxesMask(ArrayRef<int64_t> axes, int64_t rank) {
  SmallVector<bool> reducedAxes(rank, false);
  for (int64_t axis : axes) {
    if (axis < 0 || axis >= rank)
      return {};
    reducedAxes[axis] = true;
  }
  return reducedAxes;
}

static RankedTensorType getAllOnesType(RankedTensorType inputType, Type elementType) {
  return RankedTensorType::get(SmallVector<int64_t>(inputType.getRank(), 1), elementType);
}

static RankedTensorType getKeepdimsType(RankedTensorType inputType, Type elementType, ArrayRef<bool> reducedAxes) {
  SmallVector<int64_t> shape;
  shape.reserve(inputType.getRank());
  for (auto [dim, isReduced] : llvm::zip_equal(inputType.getShape(), reducedAxes))
    shape.push_back(isReduced ? 1 : dim);
  return RankedTensorType::get(shape, elementType, inputType.getEncoding());
}

static RankedTensorType getCompactKeptType(RankedTensorType inputType, Type elementType, ArrayRef<bool> reducedAxes) {
  SmallVector<int64_t> shape;
  for (auto [dim, isReduced] : llvm::zip_equal(inputType.getShape(), reducedAxes))
    if (!isReduced)
      shape.push_back(dim);
  return RankedTensorType::get(shape, elementType, inputType.getEncoding());
}

static RankedTensorType getReducedSliceType(RankedTensorType inputType, ArrayRef<bool> reducedAxes) {
  SmallVector<int64_t> shape;
  shape.reserve(inputType.getRank());
  for (auto [dim, isReduced] : llvm::zip_equal(inputType.getShape(), reducedAxes))
    shape.push_back(isReduced ? dim : 1);
  return RankedTensorType::get(shape, inputType.getElementType(), inputType.getEncoding());
}

static RankedTensorType getLanePackedKeepdimsType(int64_t laneCount, RankedTensorType leafType) {
  SmallVector<int64_t> shape(leafType.getShape().begin(), leafType.getShape().end());
  shape.front() = laneCount;
  return RankedTensorType::get(shape, leafType.getElementType(), leafType.getEncoding());
}

static SmallVector<int64_t> getKeptAxes(ArrayRef<bool> reducedAxes) {
  SmallVector<int64_t> keptAxes;
  for (auto [axis, isReduced] : llvm::enumerate(reducedAxes))
    if (!isReduced)
      keptAxes.push_back(static_cast<int64_t>(axis));
  return keptAxes;
}

static Value
computeLaneIndex(Value lane, int64_t stride, int64_t dimSize, ConversionPatternRewriter& rewriter, Location loc) {
  if (dimSize == 1)
    return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);

  MLIRContext* context = rewriter.getContext();
  AffineExpr d0 = getAffineDimExpr(0, context);
  AffineExpr expr = d0;
  if (stride != 1)
    expr = expr.floorDiv(stride);
  if (dimSize != 1)
    expr = expr % dimSize;
  return createOrFoldAffineApply(rewriter, loc, expr, ValueRange {lane}, rewriter.getInsertionBlock()->getParentOp());
}

static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
                                                     ArrayRef<bool> reducedAxes,
                                                     RankedTensorType batchType,
                                                     RankedTensorType leafType,
                                                     ConversionPatternRewriter& rewriter,
                                                     Location loc) {
  auto inputType = cast<RankedTensorType>(input.getType());
  auto sliceType = getReducedSliceType(inputType, reducedAxes);
  SmallVector<int64_t> keptAxes = getKeptAxes(reducedAxes);

  int64_t laneCount = 1;
  SmallVector<int64_t> keptAxisStrides(keptAxes.size(), 1);
  for (int64_t index = static_cast<int64_t>(keptAxes.size()) - 1; index >= 0; --index) {
    keptAxisStrides[index] = laneCount;
    int64_t dimSize = inputType.getDimSize(keptAxes[index]);
    if (dimSize <= 0)
      return failure();
    if (laneCount > std::numeric_limits<int32_t>::max() / dimSize)
      return failure();
    laneCount *= dimSize;
  }

  SmallVector<OpFoldResult> sliceOffsets;
  SmallVector<OpFoldResult> sliceSizes;
  SmallVector<OpFoldResult> insertOffsets;
  SmallVector<OpFoldResult> insertSizes(inputType.getRank(), rewriter.getIndexAttr(1));
  SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, inputType.getRank());
  sliceOffsets.reserve(inputType.getRank());
  sliceSizes.reserve(inputType.getRank());
  insertOffsets.reserve(inputType.getRank());

  auto batchOp =
    createSpatComputeBatch(rewriter,
                           loc,
                           TypeRange {batchType},
                           laneCount,
                           {},
                           ValueRange {input},
                           [&](detail::SpatComputeBatchBodyArgs args) {
                             size_t keptAxisIndex = 0;
                             sliceOffsets.clear();
                             sliceSizes.clear();
                             insertOffsets.clear();
                             for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
                               if (isReduced) {
                                 sliceOffsets.push_back(rewriter.getIndexAttr(0));
                                 sliceSizes.push_back(rewriter.getIndexAttr(inputType.getDimSize(axis)));
                                 continue;
                               }

                               Value axisIndex = computeLaneIndex(
                                 args.lane, keptAxisStrides[keptAxisIndex], inputType.getDimSize(axis), rewriter, loc);
                               ++keptAxisIndex;
                               sliceOffsets.push_back(axisIndex);
                               sliceSizes.push_back(rewriter.getIndexAttr(1));
                             }

                             insertOffsets.push_back(args.lane);
                             insertOffsets.append(inputType.getRank() - 1, rewriter.getIndexAttr(0));

                             Value slice = tensor::ExtractSliceOp::create(
                               rewriter, loc, sliceType, args.inputs.front(), sliceOffsets, sliceSizes, unitStrides);
                             Value reduced = spatial::SpatVAvgOp::create(rewriter, loc, leafType, slice).getResult();
                             createParallelInsertSliceIntoBatchOutput(
                               rewriter, loc, reduced, args.outputs.front(), insertOffsets, insertSizes, unitStrides);
                           });
  if (failed(batchOp))
    return failure();
  return (*batchOp).getResult(0);
}

static Value buildKeepdimsFromLanePackedBatch(Value batchValue,
                                              RankedTensorType keepdimsType,
                                              RankedTensorType compactKeptType,
                                              ArrayRef<bool> reducedAxes,
                                              ConversionPatternRewriter& rewriter,
                                              Location loc) {
  auto batchType = cast<RankedTensorType>(batchValue.getType());
  if (batchType == keepdimsType)
    return batchValue;

  SmallVector<ReassociationIndices> collapseToFlat {{}};
  for (int64_t axis = 0; axis < batchType.getRank(); ++axis)
    collapseToFlat.front().push_back(axis);

  SmallVector<ReassociationIndices> expandFlatToCompact(1);
  for (int64_t axis = 0; axis < compactKeptType.getRank(); ++axis)
    expandFlatToCompact.front().push_back(axis);

  SmallVector<ReassociationIndices> expandCompactToKeepdims;
  ReassociationIndices pendingLeadingReducedAxes;
  for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
    if (isReduced) {
      if (expandCompactToKeepdims.empty())
        pendingLeadingReducedAxes.push_back(axis);
      else
        expandCompactToKeepdims.back().push_back(axis);
      continue;
    }

    expandCompactToKeepdims.emplace_back();
    auto& group = expandCompactToKeepdims.back();
    group.append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());
    pendingLeadingReducedAxes.clear();
    group.push_back(axis);
  }
  if (!pendingLeadingReducedAxes.empty())
    expandCompactToKeepdims.back().append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());

  auto reshapeCompute =
    createSpatCompute<1>(rewriter, loc, TypeRange {keepdimsType}, {}, ValueRange {batchValue}, [&](Value input) {
      auto flatType =
        RankedTensorType::get({batchType.getDimSize(0)}, batchType.getElementType(), batchType.getEncoding());
      Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, input, collapseToFlat);
      Value compact = flat;
      if (compactKeptType != flatType)
        compact = tensor::ExpandShapeOp::create(rewriter, loc, compactKeptType, flat, expandFlatToCompact);
      Value keepdims = compact;
      if (keepdimsType != compactKeptType)
        keepdims = tensor::ExpandShapeOp::create(rewriter, loc, keepdimsType, compact, expandCompactToKeepdims);
      spatial::SpatYieldOp::create(rewriter, loc, keepdims);
    });
  return reshapeCompute.getResult(0);
}

static SmallVector<ReassociationIndices> buildCollapseReassociation(ArrayRef<bool> reducedAxes) {
  SmallVector<ReassociationIndices> reassociation;
  ReassociationIndices currentGroup;

  for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
    currentGroup.push_back(axis);
    if (!isReduced) {
      reassociation.push_back(currentGroup);
      currentGroup.clear();
    }
  }

  if (!currentGroup.empty()) {
    if (reassociation.empty())
      reassociation.push_back(std::move(currentGroup));
    else
      reassociation.back().append(currentGroup.begin(), currentGroup.end());
  }

  return reassociation;
}

static Value squeezeReducedAxes(Value keepdimsValue,
                                RankedTensorType resultType,
                                ArrayRef<bool> reducedAxes,
                                ConversionPatternRewriter& rewriter,
                                Location loc) {
  SmallVector<ReassociationIndices> reassociation =
    resultType.getRank() == 0 ? SmallVector<ReassociationIndices> {} : buildCollapseReassociation(reducedAxes);
  if (isCompileTimeComputable(keepdimsValue))
    return tensor::CollapseShapeOp::create(rewriter, loc, resultType, keepdimsValue, reassociation).getResult();

  auto squeezeCompute =
    createSpatCompute<1>(rewriter, loc, TypeRange {resultType}, {}, ValueRange {keepdimsValue}, [&](Value input) {
      Value collapsed = tensor::CollapseShapeOp::create(rewriter, loc, resultType, input, reassociation);
      spatial::SpatYieldOp::create(rewriter, loc, collapsed);
    });
  return squeezeCompute.getResult(0);
}

template <typename ReduceMeanOp>
struct ReduceMeanToSpatialCompute : OpConversionPattern<ReduceMeanOp> {
  using OpConversionPattern<ReduceMeanOp>::OpConversionPattern;
  using Adaptor = typename ReduceMeanOp::Adaptor;

  LogicalResult matchAndRewrite(ReduceMeanOp reduceMeanOp,
                                Adaptor adaptor,
                                ConversionPatternRewriter& rewriter) const override {
    auto inputType = dyn_cast<RankedTensorType>(adaptor.getData().getType());
    auto resultType = dyn_cast<RankedTensorType>(reduceMeanOp.getReduced().getType());
    if (!inputType || !resultType || !inputType.hasStaticShape() || !resultType.hasStaticShape())
      return failure();
    if (inputType.getRank() == 0) {
      rewriter.replaceOp(reduceMeanOp, adaptor.getData());
      return success();
    }

    auto semantics = getReduceMeanSemantics(reduceMeanOp, adaptor, inputType.getRank());
    if (failed(semantics))
      return rewriter.notifyMatchFailure(reduceMeanOp, "requires compile-time constant, in-range ReduceMean axes");
    if (semantics->isIdentity) {
      if (inputType != resultType)
        return rewriter.notifyMatchFailure(
          reduceMeanOp, "noop_with_empty_axes identity requires the result type to match the input type");
      rewriter.replaceOp(reduceMeanOp, adaptor.getData());
      return success();
    }

    SmallVector<bool> reducedAxes = buildReducedAxesMask(semantics->axes, inputType.getRank());
    if (reducedAxes.empty() && inputType.getRank() != 0)
      return failure();

    Location loc = reduceMeanOp.getLoc();
    RankedTensorType leafType = getAllOnesType(inputType, resultType.getElementType());
    RankedTensorType compactKeptType = getCompactKeptType(inputType, resultType.getElementType(), reducedAxes);
    RankedTensorType keepdimsType = getKeepdimsType(inputType, resultType.getElementType(), reducedAxes);
    int64_t laneCount = 1;
    for (int64_t dim : compactKeptType.getShape())
      laneCount *= dim;
    RankedTensorType batchType = getLanePackedKeepdimsType(laneCount, leafType);

    auto lanePackedKeepdims =
      buildReduceMeanKeepdimsBatch(adaptor.getData(), reducedAxes, batchType, leafType, rewriter, loc);
    if (failed(lanePackedKeepdims))
      return failure();
    Value reducedKeepdims =
      buildKeepdimsFromLanePackedBatch(*lanePackedKeepdims, keepdimsType, compactKeptType, reducedAxes, rewriter, loc);

    if (semantics->keepdims != 0) {
      rewriter.replaceOp(reduceMeanOp, reducedKeepdims);
      return success();
    }

    Value reduced = squeezeReducedAxes(reducedKeepdims, resultType, reducedAxes, rewriter, loc);
    rewriter.replaceOp(reduceMeanOp, reduced);
    return success();
  }
};

} // namespace

void populateReduceMeanPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
  patterns.add<ReduceMeanToSpatialCompute<ONNXReduceMeanV13Op>, ReduceMeanToSpatialCompute<ONNXReduceMeanOp>>(ctx);
}

} // namespace onnx_mlir