104 lines
3.2 KiB
C++
104 lines
3.2 KiB
C++
#pragma once
|
|
|
|
#include "llvm-project/clang/include/clang/Basic/LLVM.h"
|
|
#include "llvm/Support/JSON.h"
|
|
|
|
#include "Common/ValueMap.hpp"
|
|
#include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/PIM/PimOps.hpp"
|
|
|
|
namespace onnx_mlir {
|
|
|
|
using namespace llvm;
|
|
using namespace mlir;
|
|
using Value = mlir::Value;
|
|
using Type = mlir::Type;
|
|
using FunctionType = mlir::FunctionType;
|
|
|
|
struct MemEntry {
|
|
size_t address;
|
|
size_t size;
|
|
};
|
|
|
|
class PimMemory {
|
|
SmallVector<std::pair<MemEntry, Value>, 32> memEntries;
|
|
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(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:
|
|
SmallDenseMap<Value, MemEntry, 32> memEntriesMap;
|
|
PimMemory hostMem;
|
|
|
|
private:
|
|
SmallDenseMap<size_t, PimMemory> deviceMem;
|
|
|
|
public:
|
|
PimAcceleratorMemory()
|
|
: hostMem(memEntriesMap) {}
|
|
|
|
PimMemory getOrCreateDeviceMem(size_t id);
|
|
|
|
size_t getValueAddress(Value value) const;
|
|
};
|
|
|
|
class PimCodeGen {
|
|
PimAcceleratorMemory& memory;
|
|
raw_fd_ostream& coreFileStream;
|
|
|
|
static json::Object createEmptyOffset();
|
|
void emitInstruction(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(StringRef opName, size_t rdAddr, size_t rdOffset, size_t rs1Addr, size_t rs1Offset, size_t size) const;
|
|
void emitCommunicationOp(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, raw_fd_ostream& coreJson)
|
|
: memory(memory), coreFileStream(coreJson) {}
|
|
|
|
void codeGenLoadOp(pim::PimMemCopyHostToDevOp loadOp) const;
|
|
void codeGenStoreOp(pim::PimMemCopyDevToHostOp storeOp) const;
|
|
void codeGenLmvOp(pim::PimMemCopyOp lmvOp) const;
|
|
|
|
void codeGenReceiveOp(pim::PimReceiveOp receiveOp) const;
|
|
void codeGenSendOp(pim::PimSendOp sendOp) const;
|
|
|
|
template <typename MVMTy>
|
|
void codeGenMVMLikeOp(size_t mvmId, MVMTy mvmLikeOp, bool transposeMatrix);
|
|
|
|
void codeGenVAddOp(pim::PimVAddOp vaddOp) const;
|
|
void codeGenVMaxOp(pim::PimVMaxOp vmaxOp) const;
|
|
void codeGenVReluOp(pim::PimVReluOp vreluOp) const;
|
|
void codeGenApplyFiltersOp(pim::PimApplyFiltersOp applyFiltersOp) const;
|
|
};
|
|
|
|
OnnxMlirCompilerErrorCodes compileToPimJson(ModuleOp& moduleOpRef, std::string& outputDirName);
|
|
|
|
} // namespace onnx_mlir
|