114 lines
4.6 KiB
C++
114 lines
4.6 KiB
C++
#pragma once
|
|
|
|
#include "llvm-project/clang/include/clang/Basic/LLVM.h"
|
|
#include "llvm/Support/JSON.h"
|
|
|
|
#include "onnx-mlir/Compiler/OMCompilerTypes.h"
|
|
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
|
|
|
|
namespace onnx_mlir {
|
|
|
|
struct MemEntry {
|
|
size_t address;
|
|
size_t size;
|
|
};
|
|
|
|
class PimMemory {
|
|
llvm::SmallVector<std::pair<MemEntry, mlir::Value>, 32> memEntries;
|
|
llvm::SmallDenseMap<mlir::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(mlir::Value value);
|
|
void allocateMemoryForValue(mlir::Value value, MemEntry& memEntry);
|
|
|
|
public:
|
|
PimMemory(llvm::SmallDenseMap<mlir::Value, MemEntry, 32>& globalMemEntriesMap)
|
|
: globalMemEntriesMap(globalMemEntriesMap) {}
|
|
|
|
void allocateHost(mlir::ModuleOp moduleOp, mlir::func::FuncOp funcOp);
|
|
void allocateCore(mlir::Operation* op);
|
|
|
|
size_t getFirstAvailableAddress() const { return firstAvailableAddress; }
|
|
MemEntry getMemEntry(mlir::Value value) const;
|
|
};
|
|
|
|
class PimAcceleratorMemory {
|
|
public:
|
|
llvm::SmallDenseMap<mlir::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(mlir::Value value, const StaticValueKnowledge& knowledge = {}) const;
|
|
};
|
|
|
|
class PimCodeGen {
|
|
PimAcceleratorMemory& memory;
|
|
llvm::raw_fd_ostream& coreFileStream;
|
|
|
|
size_t addressOf(mlir::Value value, const StaticValueKnowledge& knowledge) const {
|
|
return memory.getValueAddress(value, knowledge);
|
|
}
|
|
|
|
static llvm::json::Object createEmptyOffset();
|
|
void emitInstruction(llvm::json::Object instruction) 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(mlir::StringRef opName,
|
|
size_t rdAddr,
|
|
size_t rdOffset,
|
|
size_t rs1Addr,
|
|
size_t rs1Offset,
|
|
size_t size,
|
|
mlir::StringRef sizeFieldName = "size") const;
|
|
void emitCommunicationOp(mlir::StringRef opName, 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:
|
|
PimCodeGen(PimAcceleratorMemory& memory, llvm::raw_fd_ostream& coreJson)
|
|
: memory(memory), coreFileStream(coreJson) {}
|
|
|
|
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;
|
|
|
|
template <typename MVMTy>
|
|
void codeGenMVMLikeOp(size_t mvmId, MVMTy mvmLikeOp, bool transposeMatrix, const StaticValueKnowledge& knowledge);
|
|
|
|
void codeGenVVAddOp(pim::PimVVAddOp vvaddOp, const StaticValueKnowledge& knowledge) const;
|
|
void codeGenVVSubOp(pim::PimVVSubOp vvsubOp, const StaticValueKnowledge& knowledge) const;
|
|
void codeGenVVMulOp(pim::PimVVMulOp vvmulOp, const StaticValueKnowledge& knowledge) const;
|
|
void codeGenVVMaxOp(pim::PimVVMaxOp vvmaxOp, const StaticValueKnowledge& knowledge) const;
|
|
void codeGenVVDMulOp(pim::PimVVDMulOp vvdmulOp, const StaticValueKnowledge& knowledge) const;
|
|
void codeGenVAvgOp(pim::PimVAvgOp vavgOp, const StaticValueKnowledge& knowledge) const;
|
|
void codeGenVReluOp(pim::PimVReluOp vreluOp, const StaticValueKnowledge& knowledge) const;
|
|
void codeGenVTanhOp(pim::PimVTanhOp vtanhOp, const StaticValueKnowledge& knowledge) const;
|
|
void codeGenVSigmOp(pim::PimVSigmOp vsigmOp, const StaticValueKnowledge& knowledge) const;
|
|
void codeGenVSoftmaxOp(pim::PimVSoftmaxOp vsoftmaxOp, const StaticValueKnowledge& knowledge) const;
|
|
void codeGenTransposeOp(pim::PimTransposeOp transposeOp, const StaticValueKnowledge& knowledge) const;
|
|
};
|
|
|
|
OnnxMlirCompilerErrorCodes compileToPimJson(mlir::ModuleOp& moduleOpRef, std::string& outputDirName);
|
|
|
|
} // namespace onnx_mlir
|