fix pool lowering

better reports (dcp merge and memory)

src/PIM/Compiler/PimCodeGen.cpp

@@ -25,6 +25,7 @@
 #include <utility>
 
 #include "Common/PimCommon.hpp"
+#include "Common/IR/CompactAsmUtils.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Compiler/PimArtifactWriter.hpp"
 #include "src/Accelerators/PIM/Compiler/PimBatchEmission.hpp"
@@ -36,6 +37,7 @@
 using namespace llvm;
 using namespace mlir;
 using namespace onnx_mlir;
+using namespace onnx_mlir::compact_asm;
 
 static size_t getValueSizeInBytes(mlir::Value value) {
   auto type = cast<ShapedType>(value.getType());
@@ -125,26 +127,29 @@ std::string formatMemory(uint64_t bytes) {
   return rss.str();
 }
 
-void PimMemory::report(llvm::raw_ostream& file) {
-  uint64_t numAlloca = 0;
-  uint64_t sizeAlloca = 0;
-  uint64_t numGlobal = 0;
-  uint64_t sizeGlobal = 0;
+static void printMemoryReportRow(raw_ostream& os, const MemoryReportRow& row) {
+  os << "\tNumber of allocas: " << row.numAlloca << "\n";
+  os << "\tAllocated memory: " << formatMemory(row.sizeAlloca) << "\n";
+  os << "\tNumber of globals: " << row.numGlobal << "\n";
+  os << "\tGlobal memory: " << formatMemory(row.sizeGlobal) << "\n";
+}
+
+MemoryReportRow PimMemory::getReportRow() const {
+  MemoryReportRow row;
   for (auto& [val, memEntry] : globalMemEntriesMap) {
     if (auto op = val.getDefiningOp()) {
-      if (auto allocaOp = dyn_cast<memref::AllocOp>(op)) {
-        numAlloca++;
-        sizeAlloca += memEntry.size;
+      if (isa<memref::AllocOp>(op)) {
+        row.numAlloca++;
+        row.sizeAlloca += memEntry.size;
       }
 
-      if (auto allocaOp = dyn_cast<memref::GetGlobalOp>(op)) {
-        numGlobal++;
-        sizeGlobal += memEntry.size;
+      if (isa<memref::GetGlobalOp>(op)) {
+        row.numGlobal++;
+        row.sizeGlobal += memEntry.size;
       }
     }
   }
-
-  file << numAlloca << " " << formatMemory(sizeAlloca) << " " << numGlobal << " " << formatMemory(sizeGlobal) << "\n";
+  return row;
 }
 
 void PimMemory::remove(mlir::Value val) {
@@ -193,17 +198,64 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value value, const StaticValu
 }
 
 void PimAcceleratorMemory::reportHost() {
-  llvm::raw_os_ostream os(fileReport);
-  os << "Host Memory\t";
-  hostMem.report(os);
-  os.flush();
+  hostReportRow = hostMem.getReportRow();
 }
 
 void PimAcceleratorMemory::reportCore(size_t coreId) {
+  coreReportRows.push_back({coreId, deviceMem.at(coreId).getReportRow()});
+}
+
+void PimAcceleratorMemory::flushReport() {
+  if (!fileReport.is_open())
+    return;
+
   llvm::raw_os_ostream os(fileReport);
-  os << "Core " << coreId << " Memory\t";
-  deviceMem.at(coreId).report(os);
+  if (hostReportRow.has_value()) {
+    os << "Host:\n";
+    printMemoryReportRow(os, *hostReportRow);
+  }
+
+  if (!coreReportRows.empty()) {
+    if (hostReportRow.has_value())
+      os << "\n";
+
+    llvm::stable_sort(coreReportRows, [](const auto& lhs, const auto& rhs) {
+      const MemoryReportRow& lhsRow = lhs.second;
+      const MemoryReportRow& rhsRow = rhs.second;
+      if (lhsRow.sizeAlloca != rhsRow.sizeAlloca)
+        return lhsRow.sizeAlloca > rhsRow.sizeAlloca;
+      if (lhsRow.numAlloca != rhsRow.numAlloca)
+        return lhsRow.numAlloca > rhsRow.numAlloca;
+      if (lhsRow.sizeGlobal != rhsRow.sizeGlobal)
+        return lhsRow.sizeGlobal > rhsRow.sizeGlobal;
+      if (lhsRow.numGlobal != rhsRow.numGlobal)
+        return lhsRow.numGlobal > rhsRow.numGlobal;
+      return lhs.first < rhs.first;
+    });
+
+    for (size_t index = 0; index < coreReportRows.size();) {
+      size_t runEnd = index + 1;
+      while (runEnd < coreReportRows.size() && coreReportRows[runEnd].second == coreReportRows[index].second)
+        ++runEnd;
+
+      llvm::SmallVector<size_t, 8> coreIds;
+      coreIds.reserve(runEnd - index);
+      for (size_t coreIndex = index; coreIndex < runEnd; ++coreIndex)
+        coreIds.push_back(coreReportRows[coreIndex].first);
+
+      os << "Core ";
+      printCompressedIntegerEntries(os, ArrayRef<size_t>(coreIds));
+      os << ":\n";
+      printMemoryReportRow(os, coreReportRows[index].second);
+      if (runEnd < coreReportRows.size())
+        os << "\n";
+
+      index = runEnd;
+    }
+  }
+
   os.flush();
+  fileReport.close();
 }
 
 void PimAcceleratorMemory::clean(mlir::Operation* op) {
@@ -867,5 +919,6 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimJson(ModuleOp& moduleOp, std::
     }
   }
 
+  memory.flushReport();
   return writeConfigJson(funcOp, memory, maxCoreId, std::move(xbarsPerArrayGroup), outputDirPath);
 }

src/PIM/Compiler/PimCodeGen.hpp

@@ -8,6 +8,7 @@
 #include "llvm/Support/raw_os_ostream.h"
 
 #include <fstream>
+#include <optional>
 
 #include "onnx-mlir/Compiler/OMCompilerTypes.h"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
@@ -20,6 +21,18 @@ struct MemEntry {
   size_t size;
 };
 
+struct MemoryReportRow {
+  uint64_t numAlloca = 0;
+  uint64_t sizeAlloca = 0;
+  uint64_t numGlobal = 0;
+  uint64_t sizeGlobal = 0;
+
+  bool operator==(const MemoryReportRow& other) const {
+    return numAlloca == other.numAlloca && sizeAlloca == other.sizeAlloca && numGlobal == other.numGlobal
+           && sizeGlobal == other.sizeGlobal;
+  }
+};
+
 class PimMemory {
   llvm::SmallVector<std::pair<MemEntry, mlir::Value>, 32> memEntries;
   llvm::SmallDenseMap<mlir::Value, MemEntry, 32>& globalMemEntriesMap;
@@ -37,7 +50,7 @@ public:
 
   void allocateHost(mlir::ModuleOp moduleOp, mlir::func::FuncOp funcOp);
   void allocateCore(mlir::Operation* op);
-  void report(llvm::raw_ostream& os);
+  MemoryReportRow getReportRow() const;
   void remove(mlir::Value val);
 
   size_t getFirstAvailableAddress() const { return firstAvailableAddress; }
@@ -52,6 +65,8 @@ public:
 private:
   llvm::SmallDenseMap<size_t, PimMemory> deviceMem;
   std::fstream fileReport;
+  std::optional<MemoryReportRow> hostReportRow;
+  llvm::SmallVector<std::pair<size_t, MemoryReportRow>, 32> coreReportRows;
 
 public:
   PimAcceleratorMemory()
@@ -72,6 +87,7 @@ public:
   size_t getValueAddress(mlir::Value value, const StaticValueKnowledge& knowledge = {}) const;
   void reportHost();
   void reportCore(size_t coreId);
+  void flushReport();
   void clean(mlir::Operation* op);
 };
 

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -1,16 +1,3 @@
-/*
- * SPDX-License-Identifier: Apache-2.0
- */
-
-//===------------------------- PimCompilerOptions.cpp --------------------===//
-//
-// Copyright 2022 The IBM Research Authors.
-//
-// =============================================================================
-//
-// Compiler Options for PIM
-//
-//===----------------------------------------------------------------------===//
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 
 #define DEBUG_TYPE "PimCompilerOptions"