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

#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"

using namespace mlir;

namespace onnx_mlir {

struct Concat : public OpConversionPattern<ONNXConcatOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult matchAndRewrite(ONNXConcatOp maxpoolOp,
                                ONNXConcatOpAdaptor adaptor,
                                ConversionPatternRewriter& rewriter) const final {
    auto inputs = adaptor.getInputs();
    int64_t axis = adaptor.getAxis();

    rewriter.replaceOpWithNewOp<tensor::ConcatOp>(maxpoolOp, axis, inputs);

    return success();
  }
};

void populateConcatPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
  patterns.insert<Concat>(ctx);
}

} // namespace onnx_mlir
