big refactor

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -29,13 +29,13 @@
 #include "Common/IR/ConstantUtils.hpp"
 #include "Common/PimCommon.hpp"
 #include "Common/Support/CheckedArithmetic.hpp"
-#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/Common.hpp"
 #include "Conversion/SpatialToPim/Patterns.hpp"
 #include "Dialect/Pim/PimOps.hpp"
 #include "Dialect/Spatial/SpatialOps.hpp"
-#include "Pass/PIMPasses.h"
+#include "Passes/PIMPasses.h"
 #include "SpatialToPimPass.hpp"
 
 using namespace mlir;
@@ -66,17 +66,20 @@ createZeroPaddedTensor(IRRewriter& rewriter, Location loc, Value value, RankedTe
   return padOp.getResult();
 }
 
-static FailureOr<Value> padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector) {
+static FailureOr<Value> padHVectorInputToCrossbarSize(IRRewriter& rewriter,
+                                                      Location loc,
+                                                      Value vector,
+                                                      int64_t crossbarSize) {
   auto vectorType = cast<RankedTensorType>(vector.getType());
   ArrayRef<int64_t> shape = vectorType.getShape();
   assert(isHVectorShape(shape) && "expected a horizontal vector");
-  assert(shape[1] <= static_cast<int64_t>(crossbarSize) && "vector width must fit in one crossbar");
+  assert(shape[1] <= crossbarSize && "vector width must fit in one crossbar");
 
-  if (shape[1] == static_cast<int64_t>(crossbarSize))
+  if (shape[1] == crossbarSize)
     return vector;
 
   auto paddedType = RankedTensorType::get(
-    {shape[0], static_cast<int64_t>(crossbarSize)}, vectorType.getElementType(), vectorType.getEncoding());
+    {shape[0], crossbarSize}, vectorType.getElementType(), vectorType.getEncoding());
   return createZeroPaddedTensor(rewriter, loc, vector, paddedType);
 }
 
@@ -84,6 +87,11 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   outputTensors.clear();
   operationsToRemove.clear();
   ModuleOp moduleOp = getOperation();
+  if (!hasTarget || failed(targetResources.verify())) {
+    moduleOp.emitError("Spatial-to-PIM lowering requires valid injected target resources");
+    signalPassFailure();
+    return;
+  }
   MLIRContext* ctx = moduleOp.getContext();
 
   auto entryFunc = getPimEntryFunc(moduleOp);
@@ -265,15 +273,16 @@ LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func:
     ArrayRef<int64_t> outputShape = outputType.getShape();
     assert(isHVectorShape(outputShape) && "expected a horizontal vector output");
     auto weightType = cast<RankedTensorType>(vmmOp.getWeight().getType());
-    const int64_t xbarDim = static_cast<int64_t>(crossbarSize);
+    const int64_t xbarDim = static_cast<int64_t>(targetResources.matrixShape.columns);
     const int64_t paddedOutputWidth = ceilIntegerDivide(outputShape[1], xbarDim) * xbarDim;
     assert(weightType.getRank() == 2 && weightType.getDimSize(1) == paddedOutputWidth
            && "expected VMM weight width to match the padded output width");
-    assert(paddedOutputWidth / xbarDim <= static_cast<int64_t>(crossbarCountInCore)
+    assert(paddedOutputWidth / xbarDim <= static_cast<int64_t>(targetResources.matrixUnitsPerProcessor)
            && "output width must fit in one core");
 
     rewriter.setInsertionPoint(vmmOp);
-    auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput());
+    auto paddedInput = padHVectorInputToCrossbarSize(
+      rewriter, vmmOp.getLoc(), vmmOp.getInput(), xbarDim);
     if (failed(paddedInput)) {
       hasFailure = true;
       return WalkResult::interrupt();
@@ -375,4 +384,9 @@ void raptor::SpatialToPimPass::eraseOpsToRemove() {
 
 std::unique_ptr<Pass> createSpatialToPimPass() { return std::make_unique<raptor::SpatialToPimPass>(); }
 
+std::unique_ptr<Pass> createSpatialToPimPass(
+    const spatial::SpatialTargetResources& target) {
+  return std::make_unique<raptor::SpatialToPimPass>(target);
+}
+
 } // namespace onnx_mlir