Raptor/src/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.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/Transforms/DialectConversion.h"

#include <cassert>
#include <cstddef>
#include <limits>
#include <type_traits>
#include <utility>

#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"

namespace onnx_mlir {

struct FragmentAssemblyEntry {
  int64_t sourceSlot;
  int64_t sourceOffset;
  llvm::SmallVector<int64_t, 4> destinationOffsets;
  llvm::SmallVector<int64_t, 4> sizes;
};

namespace detail {

inline mlir::ValueRange getBlockArgs(mlir::Block* block) { return mlir::ValueRange(block->getArguments()); }

inline mlir::ValueRange getInputBlockArgs(mlir::Block* block, size_t weightCount) {
  return mlir::ValueRange(block->getArguments()).drop_front(weightCount);
}

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

template <typename Fn, size_t... Is>
decltype(auto) invokeWithValues(Fn&& fn, mlir::ValueRange values, std::index_sequence<Is...>) {
  return std::forward<Fn>(fn)(values[Is]...);
}

template <size_t>
using ValueArg = mlir::Value;

template <typename Fn, typename Seq>
struct InvokeWithBlockArgsResult;

template <typename Fn, size_t... Is>
struct InvokeWithBlockArgsResult<Fn, std::index_sequence<Is...>> {
  using type = std::invoke_result_t<Fn, ValueArg<Is>...>;
};

template <typename Fn, typename Seq>
using InvokeWithBlockArgsResultT = typename InvokeWithBlockArgsResult<Fn, Seq>::type;

template <typename Fn>
using InvokeWithValueRangeResultT = std::invoke_result_t<Fn, mlir::ValueRange>;

struct SpatComputeBatchBodyArgs {
  mlir::Value lane;
  mlir::ValueRange weights;
  mlir::ValueRange inputs;
  mlir::ValueRange outputs;
};

inline mlir::SmallVector<mlir::Type> getGraphComputeBlockArgTypes(mlir::ValueRange weights, mlir::ValueRange inputs) {
  mlir::SmallVector<mlir::Type> blockArgTypes;
  blockArgTypes.reserve(weights.size() + inputs.size());
  for (mlir::Value weight : weights)
    blockArgTypes.push_back(weight.getType());
  for (mlir::Value input : inputs)
    blockArgTypes.push_back(input.getType());
  return blockArgTypes;
}

inline mlir::SmallVector<mlir::Location> getGraphComputeBlockArgLocs(mlir::Location defaultLoc,
                                                                     mlir::ValueRange weights,
                                                                     mlir::ValueRange inputs) {
  mlir::SmallVector<mlir::Location> blockArgLocs;
  blockArgLocs.reserve(weights.size() + inputs.size());
  for (mlir::Value weight : weights)
    blockArgLocs.push_back(weight.getLoc());
  for (mlir::Value input : inputs)
    blockArgLocs.push_back(input.getLoc());
  return blockArgLocs;
}

inline mlir::SmallVector<mlir::Type> getGraphComputeBatchBlockArgTypes(mlir::OpBuilder& builder,
                                                                       mlir::TypeRange resultTypes,
                                                                       mlir::ValueRange weights,
                                                                       mlir::ValueRange inputs) {
  mlir::SmallVector<mlir::Type> blockArgTypes {builder.getIndexType()};
  blockArgTypes.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
  for (mlir::Value weight : weights)
    blockArgTypes.push_back(weight.getType());
  for (mlir::Value input : inputs)
    blockArgTypes.push_back(input.getType());
  llvm::append_range(blockArgTypes, resultTypes);
  return blockArgTypes;
}

inline mlir::SmallVector<mlir::Location> getGraphComputeBatchBlockArgLocs(mlir::Location defaultLoc,
                                                                          mlir::TypeRange resultTypes,
                                                                          mlir::ValueRange weights,
                                                                          mlir::ValueRange inputs) {
  mlir::SmallVector<mlir::Location> blockArgLocs {defaultLoc};
  blockArgLocs.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
  for (mlir::Value weight : weights)
    blockArgLocs.push_back(weight.getLoc());
  for (mlir::Value input : inputs)
    blockArgLocs.push_back(input.getLoc());
  blockArgLocs.append(resultTypes.size(), defaultLoc);
  return blockArgLocs;
}

} // namespace detail

template <typename RewriterT>
inline mlir::Value createSpatConcat(RewriterT& rewriter, mlir::Location loc, int64_t axis, mlir::ValueRange inputs) {
  assert(!inputs.empty() && "spat.concat requires at least one input");
  if (inputs.size() == 1)
    return inputs.front();

  auto firstType = mlir::cast<mlir::RankedTensorType>(inputs.front().getType());
  auto outputShape = llvm::to_vector(firstType.getShape());
  int64_t concatDimSize = 0;
  bool concatDimDynamic = false;

  for (mlir::Value input : inputs) {
    auto inputType = mlir::cast<mlir::RankedTensorType>(input.getType());
    assert(inputType.getRank() == firstType.getRank() && "spat.concat expects same-rank inputs");
    if (mlir::ShapedType::isDynamic(inputType.getDimSize(axis)))
      concatDimDynamic = true;
    else
      concatDimSize += inputType.getDimSize(axis);
  }

  outputShape[axis] = concatDimDynamic ? mlir::ShapedType::kDynamic : concatDimSize;
  auto outputType = mlir::RankedTensorType::get(outputShape, firstType.getElementType(), firstType.getEncoding());
  return spatial::SpatConcatOp::create(rewriter, loc, outputType, rewriter.getI64IntegerAttr(axis), inputs).getOutput();
}

template <typename RewriterT>
spatial::SpatGraphCompute createEmptySpatGraphCompute(RewriterT& rewriter,
                                                      mlir::Location loc,
                                                      mlir::TypeRange resultTypes,
                                                      mlir::ValueRange weights,
                                                      mlir::ValueRange inputs,
                                                      mlir::TypeRange blockArgTypes,
                                                      llvm::ArrayRef<mlir::Location> blockArgLocs) {
  auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs);
  rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), blockArgTypes, blockArgLocs);
  rewriter.setInsertionPointToStart(&computeOp.getBody().front());
  return computeOp;
}

template <typename RewriterT>
spatial::SpatGraphCompute createEmptySpatGraphCompute(RewriterT& rewriter,
                                                      mlir::Location loc,
                                                      mlir::TypeRange resultTypes,
                                                      mlir::ValueRange weights,
                                                      mlir::ValueRange inputs) {
  auto blockArgTypes = detail::getGraphComputeBlockArgTypes(weights, inputs);
  auto blockArgLocs = detail::getGraphComputeBlockArgLocs(loc, weights, inputs);
  return createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs, blockArgTypes, blockArgLocs);
}

/// Builds a `spat.graph_compute` with a fixed number of SSA inputs and erases it if
/// the body callback reports failure.
template <size_t NumInputs, typename RewriterT, typename BodyFn>
auto createSpatGraphCompute(RewriterT& 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 = createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs);
  auto* block = &computeOp.getBody().front();

  using BodyResult = detail::InvokeWithBlockArgsResultT<std::decay_t<BodyFn>, std::make_index_sequence<NumInputs>>;
  if constexpr (std::is_same_v<BodyResult, void>) {
    detail::invokeWithValues(std::forward<BodyFn>(body),
                             detail::getInputBlockArgs(block, weights.size()),
                             std::make_index_sequence<NumInputs> {});

    rewriter.setInsertionPointAfter(computeOp);
    return computeOp;
  }
  else {
    auto bodyResult = detail::invokeWithValues(std::forward<BodyFn>(body),
                                               detail::getInputBlockArgs(block, weights.size()),
                                               std::make_index_sequence<NumInputs> {});
    if (mlir::failed(bodyResult)) {
      rewriter.setInsertionPointAfter(computeOp);
      rewriter.eraseOp(computeOp);
      return mlir::FailureOr<spatial::SpatGraphCompute>(mlir::failure());
    }
    rewriter.setInsertionPointAfter(computeOp);
    return mlir::FailureOr<spatial::SpatGraphCompute>(computeOp);
  }
}

/// Builds a `spat.graph_compute` whose body consumes the block arguments as a single
/// `ValueRange`, which is convenient for variadic reductions/concats.
template <typename RewriterT, typename BodyFn>
auto createSpatGraphCompute(RewriterT& rewriter,
                            mlir::Location loc,
                            mlir::TypeRange resultTypes,
                            mlir::ValueRange weights,
                            mlir::ValueRange inputs,
                            BodyFn&& body) {
  auto computeOp = createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs);
  auto* block = &computeOp.getBody().front();

  using BodyResult = detail::InvokeWithValueRangeResultT<std::decay_t<BodyFn>>;
  if constexpr (std::is_same_v<BodyResult, void>) {
    std::forward<BodyFn>(body)(detail::getInputBlockArgs(block, weights.size()));

    rewriter.setInsertionPointAfter(computeOp);
    return computeOp;
  }
  else {
    auto bodyResult = std::forward<BodyFn>(body)(detail::getInputBlockArgs(block, weights.size()));
    if (mlir::failed(bodyResult)) {
      rewriter.setInsertionPointAfter(computeOp);
      rewriter.eraseOp(computeOp);
      return mlir::FailureOr<spatial::SpatGraphCompute>(mlir::failure());
    }
    rewriter.setInsertionPointAfter(computeOp);
    return mlir::FailureOr<spatial::SpatGraphCompute>(computeOp);
  }
}

template <typename RewriterT>
auto createEmptySpatGraphComputeBatch(RewriterT& rewriter,
                                      mlir::Location loc,
                                      mlir::TypeRange resultTypes,
                                      int64_t laneCount,
                                      mlir::ValueRange weights,
                                      mlir::ValueRange inputs,
                                      mlir::TypeRange blockArgTypes,
                                      llvm::ArrayRef<mlir::Location> blockArgLocs) {
  if (laneCount <= 0 || laneCount > std::numeric_limits<int32_t>::max())
    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());

  auto laneCountAttr = pim::getCheckedI32Attr(rewriter, loc, laneCount, "spatial compute_batch lane count");
  if (mlir::failed(laneCountAttr))
    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());

  auto batchOp = spatial::SpatGraphComputeBatch::create(rewriter, loc, resultTypes, *laneCountAttr, weights, inputs);
  rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), blockArgTypes, blockArgLocs);
  rewriter.setInsertionPointToStart(&batchOp.getBody().front());
  return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
}

template <typename RewriterT>
auto createEmptySpatGraphComputeBatch(RewriterT& rewriter,
                                      mlir::Location loc,
                                      mlir::TypeRange resultTypes,
                                      int64_t laneCount,
                                      mlir::ValueRange weights,
                                      mlir::ValueRange inputs) {
  auto blockArgTypes = detail::getGraphComputeBatchBlockArgTypes(rewriter, resultTypes, weights, inputs);
  auto blockArgLocs = detail::getGraphComputeBatchBlockArgLocs(loc, resultTypes, weights, inputs);
  return createEmptySpatGraphComputeBatch(
    rewriter, loc, resultTypes, laneCount, weights, inputs, blockArgTypes, blockArgLocs);
}

template <typename RewriterT, typename BodyFn>
auto createSpatGraphComputeBatch(RewriterT& rewriter,
                                 mlir::Location loc,
                                 mlir::TypeRange resultTypes,
                                 int64_t laneCount,
                                 mlir::ValueRange weights,
                                 mlir::ValueRange inputs,
                                 BodyFn&& body) {
  auto batchOp = createEmptySpatGraphComputeBatch(rewriter, loc, resultTypes, laneCount, weights, inputs);
  if (failed(batchOp))
    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
  auto* block = &(*batchOp).getBody().front();

  detail::SpatComputeBatchBodyArgs args {
    block->getArgument(0),
    mlir::ValueRange(block->getArguments()).slice(1, weights.size()),
    mlir::ValueRange(block->getArguments()).slice(1 + weights.size(), inputs.size()),
    mlir::ValueRange(block->getArguments()).drop_front(1 + weights.size() + inputs.size())};

  using BodyResult = std::invoke_result_t<BodyFn, detail::SpatComputeBatchBodyArgs>;
  if constexpr (std::is_same_v<BodyResult, void>) {
    std::forward<BodyFn>(body)(args);
    rewriter.setInsertionPointAfter(*batchOp);
    return batchOp;
  }
  else {
    auto bodyResult = std::forward<BodyFn>(body)(args);
    if (mlir::failed(bodyResult)) {
      rewriter.setInsertionPointAfter(*batchOp);
      rewriter.eraseOp(*batchOp);
      return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
    }
    rewriter.setInsertionPointAfter(*batchOp);
    return batchOp;
  }
}

template <size_t NumInputs, typename RewriterT, typename BodyFn>
auto createSpatCompute(RewriterT& rewriter,
                       mlir::Location loc,
                       mlir::TypeRange resultTypes,
                       mlir::ValueRange weights,
                       mlir::ValueRange inputs,
                       BodyFn&& body) {
  return createSpatGraphCompute<NumInputs>(
    rewriter, loc, resultTypes, weights, inputs, std::forward<BodyFn>(body));
}

template <typename RewriterT, typename BodyFn>
auto createSpatCompute(RewriterT& rewriter,
                       mlir::Location loc,
                       mlir::TypeRange resultTypes,
                       mlir::ValueRange weights,
                       mlir::ValueRange inputs,
                       BodyFn&& body) {
  return createSpatGraphCompute(rewriter, loc, resultTypes, weights, inputs, std::forward<BodyFn>(body));
}

template <typename RewriterT, typename BodyFn>
auto createSpatComputeBatch(RewriterT& rewriter,
                            mlir::Location loc,
                            mlir::TypeRange resultTypes,
                            int64_t laneCount,
                            mlir::ValueRange weights,
                            mlir::ValueRange inputs,
                            BodyFn&& body) {
  return createSpatGraphComputeBatch(
    rewriter, loc, resultTypes, laneCount, weights, inputs, std::forward<BodyFn>(body));
}

inline void createParallelInsertSliceIntoBatchOutput(mlir::PatternRewriter& rewriter,
                                                     mlir::Location loc,
                                                     mlir::Value source,
                                                     mlir::Value dest,
                                                     mlir::ArrayRef<mlir::OpFoldResult> offsets,
                                                     mlir::ArrayRef<mlir::OpFoldResult> sizes,
                                                     mlir::ArrayRef<mlir::OpFoldResult> strides) {
  auto inParallelOp = spatial::SpatInParallelOp::create(rewriter, loc);
  rewriter.setInsertionPointToStart(&inParallelOp.getRegion().front());
  mlir::tensor::ParallelInsertSliceOp::create(rewriter, loc, source, dest, offsets, sizes, strides);
}

inline void publishGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
                                              mlir::Location loc,
                                              mlir::Value fragment,
                                              mlir::Value output,
                                              mlir::Value physicalSlot) {
  auto fragmentType = mlir::cast<mlir::RankedTensorType>(fragment.getType());
  mlir::SmallVector<mlir::OpFoldResult> offsets {physicalSlot};
  mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
  mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
  for (int64_t dim : fragmentType.getShape()) {
    offsets.push_back(rewriter.getIndexAttr(0));
    sizes.push_back(rewriter.getIndexAttr(dim));
    strides.push_back(rewriter.getIndexAttr(1));
  }
  createParallelInsertSliceIntoBatchOutput(rewriter, loc, fragment, output, offsets, sizes, strides);
}

inline mlir::FailureOr<mlir::Value>
extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
                                  mlir::Location loc,
                                  mlir::Value physicalBatch,
                                  mlir::OpFoldResult slot,
                                  mlir::RankedTensorType fragmentType) {
  if (fragmentType.getRank() == 0)
    return mlir::failure();
  auto physicalType = mlir::dyn_cast<mlir::RankedTensorType>(physicalBatch.getType());
  if (!physicalType || physicalType.getRank() != fragmentType.getRank() + 1)
    return mlir::failure();
  mlir::SmallVector<mlir::OpFoldResult> offsets {slot};
  mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
  mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
  for (int64_t dim : fragmentType.getShape()) {
    offsets.push_back(rewriter.getIndexAttr(0));
    sizes.push_back(rewriter.getIndexAttr(dim));
    strides.push_back(rewriter.getIndexAttr(1));
  }
  return extractMixedSliceOrIdentity(
    rewriter, loc, physicalBatch, fragmentType, {offsets, sizes, strides});
}

template <typename BodyFn>
mlir::Value materializeOrComputeUnary(mlir::Value input,
                                      mlir::RankedTensorType resultType,
                                      mlir::PatternRewriter& rewriter,
                                      mlir::Location loc,
                                      BodyFn&& build) {
  auto&& buildFn = build;
  if (isCompileTimeComputable(input))
    return buildFn(input);

  auto computeOp = createSpatCompute<1>(
    rewriter, loc, mlir::TypeRange {resultType}, {}, mlir::ValueRange {input}, [&](mlir::Value computeInput) {
      mlir::Value result = buildFn(computeInput);
      spatial::SpatYieldOp::create(rewriter, loc, result);
    });
  return computeOp.getResult(0);
}

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

mlir::FailureOr<mlir::Value> createFragmentAssemblyBlueprint(mlir::Value physicalBatch,
                                                             mlir::RankedTensorType logicalType,
                                                             llvm::ArrayRef<FragmentAssemblyEntry> entries,
                                                             llvm::StringRef physicalLayout,
                                                             llvm::StringRef indexMap,
                                                             mlir::PatternRewriter& rewriter,
                                                             mlir::Location loc);

} // namespace onnx_mlir