From a490e41916df870cffe6a4acbaa7439d372a68a2 Mon Sep 17 00:00:00 2001 From: NiccoloN Date: Wed, 25 Feb 2026 17:55:44 +0100 Subject: [PATCH] fix dense attributes with single element writing to memory --- src/PIM/Compiler/PimCodeGen.cpp | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/src/PIM/Compiler/PimCodeGen.cpp b/src/PIM/Compiler/PimCodeGen.cpp index cecccd9..25d0a6f 100644 --- a/src/PIM/Compiler/PimCodeGen.cpp +++ b/src/PIM/Compiler/PimCodeGen.cpp @@ -475,9 +475,21 @@ OnnxMlirCompilerErrorCodes compileModuleToPIMJSON(ModuleOp& moduleOp, std::strin auto denseAttr = dyn_cast(*initialValue); if (!denseAttr) return; - auto memEntry = memory.hostMem.getMemEntry(getGlobalOp.getResult()); - auto rawData = denseAttr.getRawData(); - std::memcpy(memoryBuffer.data() + memEntry.address, rawData.data(), std::min(rawData.size(), memEntry.size)); + + MemEntry memEntry = memory.hostMem.getMemEntry(getGlobalOp.getResult()); + ArrayRef rawData = denseAttr.getRawData(); + char* dst = memoryBuffer.data() + memEntry.address; + + if (denseAttr.isSplat()) { + size_t elementSize = rawData.size(); + assert(elementSize * getGlobalOp.getType().getNumElements() == memEntry.size && "Data size mismatch"); + for (size_t offset = 0; offset < memEntry.size; offset += elementSize) + std::memcpy(dst + offset, rawData.data(), std::min(elementSize, memEntry.size - offset)); + } + else { + assert(rawData.size() == memEntry.size && "Data size mismatch"); + std::memcpy(dst, rawData.data(), rawData.size()); + } }); memoryFileStream.write(memoryBuffer.data(), memoryBuffer.size()); memoryFileStream.close();