add PIM accelerator

src/PIM/Compiler/PimCodeGen.hpp

@@ -0,0 +1,97 @@
+#pragma once
+
+#include "llvm/Support/JSON.h"
+
+#include "Common/ValueMap.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/PIM/PimOps.hpp"
+#include "src/Accelerators/PIM/Pass/PimPasses.hpp"
+#include "src/Compiler/CompilerPasses.hpp"
+
+namespace onnx_mlir {
+
+struct MemEntry {
+  size_t address;
+  size_t size;
+};
+
+class PimMemory {
+  SmallVector<std::pair<MemEntry, Value>, 32> memEntries;
+  llvm::SmallDenseMap<Value, MemEntry, 32>& globalMemEntriesMap;
+
+  size_t maxSize = 0; // 0 for unbounded memory
+  size_t startAddress = 0;
+  size_t minAlignment = 4;
+  size_t firstAvailableAddress = 0;
+
+  MemEntry* gatherMemEntry(Value value);
+  void allocateMemoryForValue(Value value, MemEntry& memEntry);
+
+public:
+  PimMemory(llvm::SmallDenseMap<Value, MemEntry, 32>& globalMemEntriesMap)
+  : globalMemEntriesMap(globalMemEntriesMap) {}
+
+  void allocateHost(ModuleOp moduleOp, func::FuncOp funcOp);
+  void allocateCore(Operation* op);
+
+  size_t getFirstAvailableAddress() const { return firstAvailableAddress; }
+  MemEntry getMemEntry(Value value) const ;
+};
+
+class PimAcceleratorMemory {
+public:
+  llvm::SmallDenseMap<Value, MemEntry, 32> memEntriesMap;
+  PimMemory hostMem;
+
+private:
+  llvm::SmallDenseMap<size_t, PimMemory> deviceMem;
+
+public:
+  PimAcceleratorMemory()
+  : hostMem(memEntriesMap) {}
+
+  PimMemory getOrCreateDeviceMem(size_t id);
+
+  size_t getValueAddress(Value value) const;
+};
+
+class PimCodeGen {
+  PimAcceleratorMemory& memory;
+  llvm::raw_fd_ostream& coreFileStream;
+
+public:
+  PimCodeGen(PimAcceleratorMemory& memory, llvm::raw_fd_ostream& coreJson)
+  : memory(memory), coreFileStream(coreJson) {}
+
+  llvm::json::Object createSetImmediate(size_t targetRegister, size_t immediate);
+  llvm::json::Object createEmptyOffset();
+
+  void genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate);
+
+  void createRd(size_t rdAddress, size_t rdOffset);
+  void createRdRs1(size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset);
+  void createRdRs1Rs2(
+    size_t rdAddress, size_t rdOffset, size_t rs1Address, size_t rs1Offset, size_t rs2Address, size_t rs2Offset);
+
+  void codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp);
+
+  void codeGenStoreOp(pim::PimMemCopyDevToHostOp storeOp);
+
+  template <typename MVMTy>
+  void codeGenMVMLikeOp(size_t mvmId, MVMTy mvmLikeOp, bool transposeMatrix);
+
+  void codeGenReceiveOp(pim::PimReceiveOp receiveOp);
+
+  void codeGenSendOp(pim::PimSendOp sendOp);
+
+  void codeGenVAddOp(pim::PimVAddOp vaddOp);
+
+  void codeGenVMaxOp(pim::PimVMaxOp vmaxOp);
+
+  void codeGenVReluOp(pim::PimVReluOp vreluOp);
+
+  void codeGenApplyFiltersOp(pim::PimApplyFiltersOp applyFiltersOp);
+};
+
+} // namespace onnx_mlir