Raptor/src/PIM/Conversion/ONNXToSpatial/Common.hpp

#pragma once

#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Block.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/DialectConversion.h"

#include <cassert>
#include <type_traits>
#include <utility>

#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"

#define DEFINE_MAP_OP(opname) opname,

namespace onnx_mlir {

template <class ShapedType>
inline auto getImageWidth(const ShapedType& shapedType) {
  return shapedType.getDimSize(2);
}

template <class ShapedType>
inline auto getImageHeight(const ShapedType& shapedType) {
  return shapedType.getDimSize(3);
}

template <class ShapedType>
inline auto getImageChannel(const ShapedType& shapedType) {
  return shapedType.getDimSize(1);
}

template <class ShapedType>
inline auto getImageN(const ShapedType& shapedType) {
  return shapedType.getDimSize(0);
}

template <class ShapedType>
inline auto getKernelWidth(const ShapedType& shapedType) {
  return shapedType.getDimSize(2);
}

template <class ShapedType>
inline auto getKernelHeight(const ShapedType& shapedType) {
  return shapedType.getDimSize(3);
}

template <class ShapedType>
inline auto getFilterCount(const ShapedType& shapedType) {
  return shapedType.getDimSize(0);
}

using HSliceId = size_t;
using CoreId = size_t;

template <class A, class B, class C = std::common_type_t<A, B>>
constexpr C ceilIntegerDivide(A a, B b) {
  static_assert(std::is_integral_v<A>, "A must be an integer type");
  static_assert(std::is_integral_v<B>, "B must be an integer type");
  C ac = static_cast<C>(a);
  C bc = static_cast<C>(b);
  return 1 + (ac - 1) / bc;
}

template <class A, class B, class C = std::common_type_t<A, B>>
constexpr std::pair<C, C> ceilIntegerDivideWithRemainder(A a, B b) {
  static_assert(std::is_integral_v<A>, "A must be an integer type");
  static_assert(std::is_integral_v<B>, "B must be an integer type");
  C ac = static_cast<C>(a);
  C bc = static_cast<C>(b);
  return {ceilIntegerDivide(ac, bc), ac % bc};
}

template <class T>
bool isVectorShape(mlir::ArrayRef<T> shape) {
  return shape.size() == 2 && (shape[0] == 1 || shape[1] == 1);
}

template <class T>
bool isMatrixShape(mlir::ArrayRef<T> shape) {
  return shape.size() == 2;
}

template <class T>
bool isHVectorShape(mlir::ArrayRef<T> shape) {
  return shape.size() == 2 && shape[0] == 1;
}

template <class T>
bool isVVectorShape(mlir::ArrayRef<T> shape) {
  return shape.size() == 2 && shape[1] == 1;
}

template <class T>
T getVectorLength(mlir::ArrayRef<T> shape) {
  assert(isVectorShape(shape));
  return shape[0] != 1 ? shape[0] : shape[1];
}

inline auto getTensorShape(mlir::Value tensor) {
  return mlir::cast<mlir::RankedTensorType>(tensor.getType()).getShape();
}

namespace detail {

template <typename Fn, size_t... Is>
void invokeWithBlockArgs(Fn&& fn, mlir::Block* block, std::index_sequence<Is...>) {
  std::forward<Fn>(fn)(block->getArgument(Is)...);
}

} // namespace detail

template <size_t NumInputs, typename BodyFn>
spatial::SpatWeightedCompute createSpatCompute(mlir::ConversionPatternRewriter& rewriter,
                                               mlir::Location loc,
                                               mlir::TypeRange resultTypes,
                                               mlir::ValueRange weights,
                                               mlir::ValueRange inputs,
                                               BodyFn&& body) {
  assert(inputs.size() == NumInputs && "NumInputs must match the number of input values");
  auto computeOp = spatial::SpatWeightedCompute::create(rewriter, loc, resultTypes, weights, inputs);

  auto* block = new mlir::Block();
  for (mlir::Value input : inputs)
    block->addArgument(input.getType(), loc);

  computeOp.getBody().push_back(block);
  rewriter.setInsertionPointToStart(block);

  detail::invokeWithBlockArgs(std::forward<BodyFn>(body), block, std::make_index_sequence<NumInputs> {});

  rewriter.setInsertionPointAfter(computeOp);
  return computeOp;
}

llvm::SmallVector<mlir::Value> sliceTensor(const mlir::Value& tensorToSlice,
                                           size_t axis,
                                           int64_t sliceSize,
                                           mlir::ConversionPatternRewriter& rewriter,
                                           mlir::Location loc);

llvm::SmallVector<mlir::Value> sliceVector(const mlir::Value& vectorToSlice,
                                           int64_t sliceSize,
                                           mlir::ConversionPatternRewriter& rewriter,
                                           mlir::Location loc);

llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>> sliceVectorPerCrossbarPerCore(
  const mlir::Value& vectorToSlice, mlir::ConversionPatternRewriter& rewriter, mlir::Location loc);

llvm::DenseMap<HSliceId, llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>>>
tileMatrix(mlir::Value& matrixToTile,
           int64_t hSliceSize,
           int64_t vSliceSize,
           mlir::ConversionPatternRewriter& rewriter,
           mlir::Location& loc);

mlir::tensor::SplatOp broadcastToVector(mlir::Value scalarToBroadcast,
                                        int64_t length,
                                        mlir::ConversionPatternRewriter& rewriter,
                                        mlir::Location loc);

mlir::Value sumTensors(mlir::ArrayRef<mlir::Value> tensors, mlir::ConversionPatternRewriter& rewriter);

}; // namespace onnx_mlir