Raptor/src/PIM/Compiler/PimCodeGen.hpp

#pragma once

#include "mlir/IR/Operation.h"

#include "llvm-project/clang/include/clang/Basic/LLVM.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/Hashing.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/JSON.h"
#include "llvm/Support/raw_os_ostream.h"

#include <array>
#include <fstream>
#include <limits>
#include <optional>
#include <string>

#include "onnx-mlir/Compiler/OMCompilerTypes.h"
#include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
#include "src/Accelerators/PIM/Compiler/PimBinaryFormat.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

namespace onnx_mlir {

struct CompiledCoreProgram;
struct CompiledTransposePlan;
class PimInstructionWriter;

struct MemEntry {
  size_t address;
  size_t size;
};

struct MemoryValueKey {
  mlir::Value value;
  std::optional<unsigned> lane;

  bool operator==(const MemoryValueKey& other) const { return value == other.value && lane == other.lane; }
};

struct CompiledLocalMemoryEntry {
  mlir::Value value;
  MemEntry memory;
};

struct CompiledCoreMemoryPlan {
  llvm::SmallVector<CompiledLocalMemoryEntry, 32> entries;
  uint64_t logicalAllocationCount = 0;
  uint64_t logicalBytes = 0;
  size_t arenaSize = 0;
};

struct MemoryReportRow {
  uint64_t hostObjectCount = 0;
  uint64_t hostBytes = 0;
  uint64_t logicalLocalAllocationCount = 0;
  uint64_t logicalLocalBytes = 0;
  uint64_t physicalLocalBytes = 0;

  bool operator==(const MemoryReportRow& other) const {
    return hostObjectCount == other.hostObjectCount && hostBytes == other.hostBytes
        && logicalLocalAllocationCount == other.logicalLocalAllocationCount
        && logicalLocalBytes == other.logicalLocalBytes && physicalLocalBytes == other.physicalLocalBytes;
  }
};

enum class MemoryReportKind {
  None,
  Alloca,
  Global,
  Input
};

struct PendingMemEntry {
  MemEntry memEntry;
  MemoryValueKey key;
  MemoryReportKind reportKind = MemoryReportKind::None;
};

struct MemoryReportEntry {
  enum class Kind {
    Core,
    Batch
  };

  Kind kind = Kind::Core;
  uint64_t id = 0;
  llvm::SmallVector<int32_t, 8> coreIds;
  MemoryReportRow row;
};

class PimMemory {
  llvm::SmallVector<PendingMemEntry, 32> memEntries;
  llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& globalMemEntriesMap;
  llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32> ownedMemEntriesMap;
  MemoryReportRow reportRow;
  std::optional<size_t> localArenaSize;

  size_t minAlignment = 4;
  size_t firstAvailableAddress = 0;

  MemEntry* gatherMemEntry(mlir::Value value, std::optional<unsigned> lane = std::nullopt);
  size_t allocateAddress(size_t size, const MemoryValueKey& key);
  void allocateGatheredMemory();
  void allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memEntry, MemoryReportKind reportKind);

public:
  PimMemory(llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& globalMemEntriesMap)
  : globalMemEntriesMap(globalMemEntriesMap) {}

  void allocateHost(mlir::ModuleOp moduleOp, mlir::func::FuncOp funcOp);
  void allocateCore(const CompiledCoreMemoryPlan& plan, std::optional<unsigned> lane = std::nullopt);
  MemoryReportRow getReportRow() const;

  size_t getFirstAvailableAddress() const { return firstAvailableAddress; }
  MemEntry getMemEntry(const MemoryValueKey& key) const;
};

class PimAcceleratorMemory {
public:
  llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32> memEntriesMap;
  PimMemory hostMem;

private:
  llvm::SmallDenseMap<size_t, PimMemory> deviceMem;
  std::fstream fileReport;
  std::optional<MemoryReportRow> hostReportRow;
  llvm::SmallVector<MemoryReportEntry, 32> reportEntries;
  uint64_t totalWeightBytes = 0;
  mutable llvm::DenseMap<mlir::Value, CompiledIndexExpr> compiledIndexExprs;
  mutable llvm::DenseMap<mlir::Value, CompiledAddressExpr> compiledAddressExprs;

public:
  PimAcceleratorMemory();
  PimAcceleratorMemory(const llvm::SmallDenseMap<MemoryValueKey, MemEntry, 32>& initialMemEntries, bool enableReport);

  PimMemory& getOrCreateDeviceMem(size_t id);

  size_t getValueAddress(mlir::Value value,
                         const StaticValueKnowledge& knowledge = {},
                         std::optional<unsigned> lane = std::nullopt) const;
  llvm::FailureOr<int64_t> getIndexValue(mlir::Value value, const StaticValueKnowledge& knowledge = {}) const;
  void reportHost();
  void recordCoreReport(size_t coreId, const MemoryReportRow& row);
  void recordBatchReport(uint64_t batchId, llvm::ArrayRef<int32_t> coreIds, const MemoryReportRow& perCoreRow);
  void setTotalWeightBytes(uint64_t bytes) { totalWeightBytes = bytes; }
  void flushReport();
};

struct CoreEmissionJob {
  mlir::Operation* coreLikeOp = nullptr;
  const CompiledCoreProgram* program = nullptr;
  const CompiledCoreMemoryPlan* memoryPlan = nullptr;
  size_t physicalCoreId = 0;
  llvm::SmallVector<unsigned, 4> lanes;
  std::optional<uint64_t> batchReportId;
};

class PimCodeGen {
  PimAcceleratorMemory& memory;
  PimInstructionWriter& instructionWriter;
  llvm::raw_fd_ostream* coreJsonStream;
  std::optional<unsigned> batchLane;
  mutable std::array<std::optional<int32_t>, 256> scalarRegisterValues = {};
  mutable std::optional<std::array<int32_t, 2>> vectorBitwidths;

  size_t addressOf(mlir::Value value, const StaticValueKnowledge& knowledge) const {
    return memory.getValueAddress(value, knowledge, batchLane);
  }
  void emitInstruction(const pim_binary::InstructionRecord& instruction) const;
  void updateScalarRegisterCache(const pim_binary::InstructionRecord& instruction) const;
  void ensureVectorBitwidth(int32_t inputBitwidth, int32_t outputBitwidth) const;

  void genSetRegisterImmediate(uint8_t registerNumber, int32_t immediate) const;
  void genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const;
  void setupRd(size_t rdAddress, size_t rdOffset) const;
  void setupRdRs1(size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset) const;
  void setupRdRs1Rs2(
    size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset, size_t rs2Address, size_t rs2Offset) const;

  void emitMemCopyOp(pim_binary::Opcode opcode,
                     size_t rdAddr,
                     size_t rdOffset,
                     size_t rs1Addr,
                     size_t rs1Offset,
                     size_t size,
                     mlir::StringRef sizeFieldName = "size") const;
  void emitCommunicationOp(pim_binary::Opcode opcode, size_t bufferAddr, size_t coreId, size_t size) const;
  void emitMvmOp(size_t groupId, size_t rdAddr, size_t rdOffset, size_t rs1Addr, size_t rs1Offset) const;

public:
  void emitBinaryVectorOp(pim_binary::Opcode opcode,
                          mlir::Value output,
                          mlir::Value lhs,
                          mlir::Value rhs,
                          const StaticValueKnowledge& knowledge) const;
  void emitUnaryVectorOp(pim_binary::Opcode opcode,
                         mlir::Value output,
                         mlir::Value input,
                         const StaticValueKnowledge& knowledge,
                         int32_t r2OrImm = 0,
                         int32_t generic1 = 0) const;

  PimCodeGen(PimAcceleratorMemory& memory, PimInstructionWriter& instructionWriter, llvm::raw_fd_ostream* coreJson)
  : memory(memory), instructionWriter(instructionWriter), coreJsonStream(coreJson) {}

  void setBatchLane(std::optional<unsigned> lane) { batchLane = lane; }
  llvm::FailureOr<int64_t> indexOf(mlir::Value value, const StaticValueKnowledge& knowledge) const {
    return memory.getIndexValue(value, knowledge);
  }

  void codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp, const StaticValueKnowledge& knowledge) const;
  void codeGenStoreOp(pim::PimMemCopyDevToHostOp storeOp, const StaticValueKnowledge& knowledge) const;
  void codeGenLmvOp(pim::PimMemCopyOp lmvOp, const StaticValueKnowledge& knowledge) const;

  void codeGenReceiveOp(pim::PimReceiveOp receiveOp, const StaticValueKnowledge& knowledge) const;
  void codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge& knowledge) const;
  void codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKnowledge& knowledge) const;

  template <typename MVMTy>
  void codeGenMVMLikeOp(size_t mvmId, MVMTy mvmLikeOp, bool transposeMatrix, const StaticValueKnowledge& knowledge);

  void codeGenTransposeOp(const CompiledTransposePlan& plan, const StaticValueKnowledge& knowledge) const;
};

OnnxMlirCompilerErrorCodes compileToPimCode(mlir::ModuleOp& moduleOpRef, std::string& outputDirName);

} // namespace onnx_mlir

namespace llvm {

template <>
struct DenseMapInfo<onnx_mlir::MemoryValueKey> {
  static onnx_mlir::MemoryValueKey getEmptyKey() { return {DenseMapInfo<mlir::Value>::getEmptyKey(), 0}; }

  static onnx_mlir::MemoryValueKey getTombstoneKey() { return {DenseMapInfo<mlir::Value>::getTombstoneKey(), 0}; }

  static unsigned getHashValue(const onnx_mlir::MemoryValueKey& key) {
    return hash_combine(key.value, key.lane.value_or(std::numeric_limits<unsigned>::max()));
  }

  static bool isEqual(const onnx_mlir::MemoryValueKey& lhs, const onnx_mlir::MemoryValueKey& rhs) { return lhs == rhs; }
};

} // namespace llvm