huge refactor for high RewritePatterns usage and less ad-hoc cpp code

remove Spatial many ops in favor of tensor ops like in pim

src/PIM/Compiler/PimWeightEmitter.cpp

@@ -0,0 +1,221 @@
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <cassert>
+
+#include "Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
+#include "src/Accelerators/PIM/Compiler/PimBatchEmission.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Compiler/PimWeightEmitter.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+
+using namespace llvm;
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+struct DenseWeightView {
+  DenseElementsAttr denseAttr;
+  SmallVector<int64_t> shape;
+  SmallVector<int64_t> strides;
+  int64_t offset = 0;
+};
+
+SmallVector<int64_t> computeRowMajorStridesForShape(ArrayRef<int64_t> shape) {
+  SmallVector<int64_t> strides(shape.size(), 1);
+  for (int64_t index = static_cast<int64_t>(shape.size()) - 2; index >= 0; --index)
+    strides[index] = strides[index + 1] * shape[index + 1];
+  return strides;
+}
+
+bool allStaticSubviewParts(memref::SubViewOp subview) {
+  return llvm::all_of(subview.getStaticOffsets(), [](int64_t value) { return !ShapedType::isDynamic(value); })
+      && llvm::all_of(subview.getStaticSizes(), [](int64_t value) { return !ShapedType::isDynamic(value); })
+      && llvm::all_of(subview.getStaticStrides(), [](int64_t value) { return !ShapedType::isDynamic(value); });
+}
+
+FailureOr<DenseWeightView> resolveDenseWeightView(ModuleOp moduleOp, mlir::Value weight) {
+  SmallVector<memref::SubViewOp> subviews;
+  mlir::Value current = weight;
+  memref::GetGlobalOp getGlobalOp;
+
+  while (true) {
+    Operation* defOp = current.getDefiningOp();
+    if (!defOp)
+      return failure();
+    if ((getGlobalOp = dyn_cast<memref::GetGlobalOp>(defOp)))
+      break;
+    if (auto subview = dyn_cast<memref::SubViewOp>(defOp)) {
+      if (!allStaticSubviewParts(subview))
+        return failure();
+      subviews.push_back(subview);
+      current = subview.getSource();
+      continue;
+    }
+    if (auto cast = dyn_cast<memref::CastOp>(defOp)) {
+      current = cast.getSource();
+      continue;
+    }
+    return failure();
+  }
+
+  auto globalOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
+  if (!globalOp || !globalOp.getInitialValue())
+    return failure();
+
+  auto denseAttr = dyn_cast<DenseElementsAttr>(*globalOp.getInitialValue());
+  if (!denseAttr)
+    return failure();
+
+  DenseWeightView view;
+  view.denseAttr = denseAttr;
+  view.shape.assign(denseAttr.getType().getShape().begin(), denseAttr.getType().getShape().end());
+  view.strides = computeRowMajorStridesForShape(view.shape);
+
+  for (memref::SubViewOp subview : llvm::reverse(subviews)) {
+    SmallVector<int64_t> nextStrides;
+    nextStrides.reserve(subview.getStaticStrides().size());
+    for (auto [offset, stride, sourceStride] :
+         llvm::zip_equal(subview.getStaticOffsets(), subview.getStaticStrides(), view.strides)) {
+      view.offset += offset * sourceStride;
+      nextStrides.push_back(stride * sourceStride);
+    }
+    view.shape.assign(subview.getStaticSizes().begin(), subview.getStaticSizes().end());
+    view.strides = std::move(nextStrides);
+  }
+
+  return view;
+}
+
+SmallVector<unsigned, 8> getUsedWeightIndices(Block& block) {
+  SmallVector<unsigned, 8> indices;
+  auto addIndex = [&](unsigned weightIndex) {
+    if (!llvm::is_contained(indices, weightIndex))
+      indices.push_back(weightIndex);
+  };
+
+  block.walk([&](pim::PimMVMOp mvmOp) { addIndex(mvmOp.getWeightIndex()); });
+  block.walk([&](pim::PimVMMOp vmmOp) { addIndex(vmmOp.getWeightIndex()); });
+  llvm::sort(indices);
+  return indices;
+}
+
+SmallVector<unsigned, 8> getUsedWeightIndices(pim::PimCoreOp coreOp) {
+  return getUsedWeightIndices(coreOp.getBody().front());
+}
+
+SmallVector<Operation*> collectTopLevelCoreLikeOps(func::FuncOp funcOp) {
+  SmallVector<Operation*> coreLikeOps;
+  for (Operation& op : funcOp.getBody().front())
+    if (dyn_cast<pim::PimCoreOp>(&op) || dyn_cast<pim::PimCoreBatchOp>(&op))
+      coreLikeOps.push_back(&op);
+  return coreLikeOps;
+}
+
+} // namespace
+
+llvm::DenseMap<size_t, llvm::DenseMap<mlir::Value, std::string>>
+createAndPopulateWeightFolder(func::FuncOp funcOp, StringRef outputDirPath) {
+  ModuleOp moduleOp = funcOp->getParentOfType<ModuleOp>();
+  auto coreWeightsDirPath = outputDirPath + "/weights";
+  auto error = sys::fs::create_directory(coreWeightsDirPath);
+  assert(!error && "Error creating weights directory");
+  size_t indexFileName = 0;
+
+  int64_t xbarSize = crossbarSize.getValue();
+  llvm::DenseMap<size_t, llvm::DenseMap<mlir::Value, std::string>> mapCoreWeightToFileName;
+  llvm::DenseMap<memref::GlobalOp, std::string> mapGlobalOpToFileName;
+
+  SmallVector<Operation*> coreLikeOps = collectTopLevelCoreLikeOps(funcOp);
+
+  for (Operation* op : coreLikeOps) {
+    auto processCore = [&](pim::PimCoreOp coreOp) {
+      size_t coreId = static_cast<size_t>(coreOp.getCoreId());
+      for (unsigned index : getUsedWeightIndices(coreOp)) {
+        if (index >= coreOp.getWeights().size()) {
+          coreOp.emitWarning("Weight index " + std::to_string(index) + " is out of range");
+          assert(index < coreOp.getWeights().size() && "Weight index is out of range");
+        }
+        mlir::Value weight = coreOp.getWeights()[index];
+
+        auto weightView = resolveDenseWeightView(moduleOp, weight);
+        if (failed(weightView)) {
+          coreOp.emitWarning("Weight is not from a memref.get_global at index " + std::to_string(index));
+          assert(succeeded(weightView) && "Weight is not from a dense memref.global view");
+        }
+
+        if (mapCoreWeightToFileName[coreId].contains(weight))
+          continue;
+
+        auto getGlobalOp = weight.getDefiningOp<memref::GetGlobalOp>();
+        auto globalOp = getGlobalOp ? lookupGlobalForGetGlobal(moduleOp, getGlobalOp) : memref::GlobalOp {};
+        if (globalOp && mapGlobalOpToFileName.contains(globalOp)) {
+          auto& fileName = mapGlobalOpToFileName[globalOp];
+          mapCoreWeightToFileName[coreId].insert({weight, fileName});
+          continue;
+        }
+
+        DenseElementsAttr denseAttr = weightView->denseAttr;
+        ArrayRef<int64_t> shape = weightView->shape;
+        assert(isMatrixShape(shape) && "Weight matrix must be 2-dimensional");
+        int64_t numRows = shape[0];
+        int64_t numCols = shape[1];
+        assert(numRows <= xbarSize && numCols <= xbarSize && "Weight dimensions must not exceed crossbar size");
+
+        size_t elementByteWidth = denseAttr.getElementType().getIntOrFloatBitWidth() / 8;
+
+        std::string newFileName = "crossbar_" + std::to_string(indexFileName++) + ".bin";
+        auto weightFilePath = (coreWeightsDirPath + "/" + newFileName).str();
+        std::error_code errorCode;
+        raw_fd_ostream weightFileStream(weightFilePath, errorCode, sys::fs::OF_None);
+        if (errorCode) {
+          errs() << "Error while opening weight file `" << weightFilePath << "`: " << errorCode.message() << '\n';
+          assert(errorCode);
+        }
+
+        uint64_t zero = 0;
+        for (int64_t row = 0; row < xbarSize; row++) {
+          for (int64_t col = 0; col < xbarSize; col++) {
+            if (row < numRows && col < numCols) {
+              int64_t elementIndex = weightView->offset + row * weightView->strides[0] + col * weightView->strides[1];
+              APInt bits = denseAttr.getValues<APFloat>()[elementIndex].bitcastToAPInt();
+              uint64_t word = bits.getZExtValue();
+              weightFileStream.write(reinterpret_cast<const char*>(&word), elementByteWidth);
+            }
+            else {
+              weightFileStream.write(reinterpret_cast<const char*>(&zero), elementByteWidth);
+            }
+          }
+        }
+
+        weightFileStream.close();
+        if (globalOp)
+          mapGlobalOpToFileName.insert({globalOp, newFileName});
+        mapCoreWeightToFileName[coreId].insert({weight, newFileName});
+      }
+      return success();
+    };
+
+    if (auto coreOp = dyn_cast<pim::PimCoreOp>(op)) {
+      (void) processCore(coreOp);
+      continue;
+    }
+
+    auto coreBatchOp = cast<pim::PimCoreBatchOp>(op);
+    for (unsigned lane = 0; lane < static_cast<unsigned>(coreBatchOp.getLaneCount()); ++lane)
+      if (failed(withScalarCoreFromBatchLane(coreBatchOp, lane, processCore)))
+        return mapCoreWeightToFileName;
+  }
+  return mapCoreWeightToFileName;
+}
+
+} // namespace onnx_mlir