#pragma once

#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"

namespace onnx_mlir {

inline constexpr llvm::StringLiteral kRowStripIndexMap = "nchw_row_strip_fragments";

struct RowStripPhysicalValue {
  mlir::Value storage;
  mlir::RankedTensorType logicalType;
  llvm::SmallVector<int64_t, 16> fragmentOffsets;
  llvm::SmallVector<int64_t, 16> fragmentSizes;
};

std::pair<llvm::SmallVector<int64_t>, llvm::SmallVector<int64_t>>
buildRowStripMetadata(mlir::RankedTensorType type);

mlir::RankedTensorType getRowStripFragmentType(mlir::RankedTensorType logicalType);

mlir::RankedTensorType getRowStripStorageType(mlir::RankedTensorType logicalType);

llvm::SmallVector<mlir::OpFoldResult> buildRowStripFragmentOffsets(mlir::PatternRewriter& rewriter,
                                                                   mlir::OpFoldResult row);

llvm::SmallVector<mlir::OpFoldResult> buildRowStripFragmentSizes(mlir::PatternRewriter& rewriter,
                                                                 mlir::RankedTensorType logicalType);

mlir::Value extractRowStripFragment(mlir::Value storage,
                                    mlir::RankedTensorType logicalType,
                                    mlir::OpFoldResult row,
                                    mlir::PatternRewriter& rewriter,
                                    mlir::Location loc);

void insertRowStripFragment(mlir::Value fragment,
                            mlir::Value output,
                            mlir::RankedTensorType logicalType,
                            mlir::OpFoldResult row,
                            mlir::PatternRewriter& rewriter,
                            mlir::Location loc);

mlir::FailureOr<mlir::Value> createPerChannelConstantFragment(mlir::DenseElementsAttr denseAttr,
                                                              mlir::RankedTensorType fragmentType,
                                                              mlir::PatternRewriter& rewriter);

mlir::FailureOr<mlir::Value> createRowStripStorageFromRows(mlir::Value rows,
                                                           mlir::RankedTensorType logicalType,
                                                           mlir::PatternRewriter& rewriter,
                                                           mlir::Location loc);

mlir::FailureOr<mlir::Value> createRowStripAssemblyBlueprint(mlir::Value storage,
                                                               mlir::RankedTensorType logicalType,
                                                               mlir::PatternRewriter& rewriter,
                                                               mlir::Location loc);

mlir::FailureOr<mlir::Value> applyRowStripRelu(mlir::Value storage,
                                               mlir::RankedTensorType logicalType,
                                               mlir::PatternRewriter& rewriter,
                                               mlir::Location loc);

mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(mlir::Value storage,
                                                  mlir::RankedTensorType logicalType,
                                                  mlir::Value bias,
                                                  mlir::PatternRewriter& rewriter,
                                                  mlir::Location loc);

} // namespace onnx_mlir