#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

#include <string>

using namespace mlir;

namespace onnx_mlir {
namespace pim {

BlockArgument PimCoreOp::getWeightArgument(unsigned idx) { return getBody().front().getArgument(idx); }

void PimCoreOp::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
  if (region.empty())
    return;

  for (unsigned index = 0; index < getWeights().size(); ++index)
    setNameFn(getWeightArgument(index), ("w" + std::to_string(index)).c_str());
}

BlockArgument PimCoreBatchOp::getLaneArgument() { return getBody().front().getArgument(0); }

BlockArgument PimCoreBatchOp::getWeightArgument(unsigned idx) { return getBody().front().getArgument(1 + idx); }

BlockArgument PimCoreBatchOp::getInputArgument(unsigned idx) {
  return getBody().front().getArgument(1 + getWeights().size() + idx);
}

void PimCoreBatchOp::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
  if (region.empty())
    return;

  setNameFn(getLaneArgument(), "lane");
  for (unsigned index = 0; index < getWeights().size(); ++index)
    setNameFn(getWeightArgument(index), ("w" + std::to_string(index)).c_str());
  for (unsigned index = 0; index < getInputs().size(); ++index)
    setNameFn(getInputArgument(index), ("in" + std::to_string(index)).c_str());
}

void PimDialect::initialize() {
  addOperations<
#define GET_OP_LIST
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.cpp.inc"

    >();
}

} // namespace pim
} // namespace onnx_mlir

//===----------------------------------------------------------------------===//
// TableGen'd op method definitions
//===----------------------------------------------------------------------===//

#define GET_OP_CLASSES
#include "src/Accelerators/PIM/Dialect/Pim/PimDialect.cpp.inc"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.cpp.inc"