#pragma once

#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/PatternMatch.h"

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

namespace onnx_mlir {

void populateInitialPatterns(mlir::RewritePatternSet& patterns);
void populateCoreBodyPatterns(mlir::RewritePatternSet& patterns);

void populateTransposeLoweringPatterns(mlir::RewritePatternSet& patterns);
void populateChannelLoweringPatterns(mlir::RewritePatternSet& patterns);
void populateGlobalTensorMaterializationPatterns(mlir::RewritePatternSet& patterns);
void populateTensorPackingPatterns(mlir::RewritePatternSet& patterns);

mlir::RankedTensorType getPackedTensorType(mlir::RankedTensorType elementType, int64_t count);
mlir::Value extractPackedChunk(mlir::Value packedValue,
                               mlir::RankedTensorType chunkType,
                               unsigned index,
                               mlir::OpBuilder& builder,
                               mlir::Location loc);
mlir::Value createPackedExtractRowsSlice(spatial::SpatExtractRowsOp extractRowsOp,
                                         unsigned startIndex,
                                         unsigned count,
                                         mlir::OpBuilder& builder,
                                         mlir::Location loc);
mlir::Value createPackedExtractSliceTensor(mlir::ValueRange values, mlir::OpBuilder& builder, mlir::Location loc);
void eraseUnusedTensorPackingOps(mlir::func::FuncOp funcOp, mlir::IRRewriter& rewriter);

} // namespace onnx_mlir