nnicolosi/Raptor
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fixes
Validate Operations / validate-operations (push) Has been cancelled
Details

Browse Source 
slightly faster codegen
This commit is contained in:
NiccoloN
2026-07-20 16:01:57 +02:00
parent ab54243fda
commit dbb66be93e
27 changed files with 1358 additions and 1241 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/PIM/Compiler/PimCodeGen.hpp
+28 -25
View File
@@ -24,6 +24,10 @@
namespace onnx_mlir {
struct CompiledCoreProgram;
struct CompiledTransposePlan;
class PimInstructionWriter;
struct MemEntry {
size_t address;
size_t size;
@@ -35,9 +39,7 @@ struct PhysicalSlotInfo {
size_t size = 0;
};
struct MemoryPlanArtifacts {
std::string textReport;
};
using MemoryPlanArtifacts = std::string;
struct MemoryValueKey {
mlir::Value value;
@@ -98,6 +100,7 @@ class PimMemory {
size_t nextPhysicalSlotId = 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);
PhysicalSlotInfo allocatePhysicalSlot(size_t slotSize, const MemoryValueKey& key);
@@ -110,7 +113,6 @@ public:
void allocateCore(mlir::Operation* op, std::optional<unsigned> lane = std::nullopt);
MemoryReportRow getReportRow() const;
const MemoryPlanArtifacts& getLivenessArtifacts() const { return livenessArtifacts; }
void remove(mlir::Value val);
size_t getFirstAvailableAddress() const { return firstAvailableAddress; }
MemEntry getMemEntry(const MemoryValueKey& key) const;
@@ -153,12 +155,11 @@ public:
uint64_t totalAllocaBytes);
void setTotalWeightBytes(uint64_t bytes) { totalWeightBytes = bytes; }
void flushReport();
void clean(mlir::Operation* op);
};
struct CoreEmissionJob {
mlir::Operation* coreLikeOp = nullptr;
size_t originalCoreId = 0;
const CompiledCoreProgram* program = nullptr;
size_t emittedCoreId = 0;
llvm::SmallVector<unsigned, 4> lanes;
std::optional<uint64_t> batchReportId;
@@ -166,11 +167,10 @@ struct CoreEmissionJob {
class PimCodeGen {
PimAcceleratorMemory& memory;
llvm::raw_fd_ostream& coreBinaryStream;
PimInstructionWriter& instructionWriter;
llvm::raw_fd_ostream* coreJsonStream;
const llvm::DenseMap<size_t, size_t>& emittedCoreIds;
std::optional<unsigned> batchLane;
mutable uint32_t emittedInstructionCount = 0;
mutable std::array<std::optional<int32_t>, 256> scalarRegisterValues = {};
size_t addressOf(mlir::Value value, const StaticValueKnowledge& knowledge) const {
@@ -181,30 +181,44 @@ class PimCodeGen {
void emitInstruction(const pim_binary::InstructionRecord& instruction) const;
void updateScalarRegisterCache(const pim_binary::InstructionRecord& instruction) 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(mlir::StringRef opName,
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(mlir::StringRef opName, size_t bufferAddr, size_t coreId, size_t 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,
size_t byteSize,
const StaticValueKnowledge& knowledge) const;
void emitUnaryVectorOp(pim_binary::Opcode opcode,
mlir::Value output,
mlir::Value input,
size_t byteSize,
const StaticValueKnowledge& knowledge,
int32_t r2OrImm = 0,
int32_t generic1 = 0) const;
PimCodeGen(PimAcceleratorMemory& memory,
llvm::raw_fd_ostream& coreBinary,
PimInstructionWriter& instructionWriter,
llvm::raw_fd_ostream* coreJson,
const llvm::DenseMap<size_t, size_t>& emittedCoreIds)
: memory(memory), coreBinaryStream(coreBinary), coreJsonStream(coreJson), emittedCoreIds(emittedCoreIds) {}
: memory(memory), instructionWriter(instructionWriter), coreJsonStream(coreJson), emittedCoreIds(emittedCoreIds) {}
uint32_t getEmittedInstructionCount() const { return emittedInstructionCount; }
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);
@@ -221,18 +235,7 @@ public:
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 codeGetGlobalOp(mlir::memref::GetGlobalOp getGlobalOp, const StaticValueKnowledge& knowledge) const;
void codeGenTransposeOp(pim::PimTransposeOp transposeOp, const StaticValueKnowledge& knowledge) const;
void codeGenTransposeOp(const CompiledTransposePlan& plan, const StaticValueKnowledge& knowledge) const;
};
OnnxMlirCompilerErrorCodes compileToPimCode(mlir::ModuleOp& moduleOpRef, std::string& outputDirName);