#pragma once

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

#include "llvm/ADT/SmallVector.h"

#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"

namespace onnx_mlir {

/// Slices a statically shaped tensor along one axis into contiguous pieces of
/// at most `sliceSize` elements.
llvm::SmallVector<mlir::Value> sliceTensor(const mlir::Value& tensorToSlice,
                                           size_t axis,
                                           int64_t sliceSize,
                                           mlir::PatternRewriter& rewriter,
                                           mlir::Location loc);

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

/// Partitions one logical vector into per-core crossbar-sized slices using the
/// current PIM target geometry.
llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>> sliceVectorPerCrossbarPerCore(
  const mlir::Value& vectorToSlice, mlir::PatternRewriter& rewriter, mlir::Location loc);

} // namespace onnx_mlir