#pragma once

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

#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"

#include <string>

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

namespace onnx_mlir {

struct WeightFileRequest {
  size_t coreId = 0;
  llvm::SmallVector<ResolvedWeightView, 8> weights;
};

llvm::DenseMap<size_t, llvm::SmallVector<std::string, 8>>
createAndPopulateWeightFolder(llvm::ArrayRef<WeightFileRequest> requests, llvm::StringRef outputDirPath);

} // namespace onnx_mlir