big refactor

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

@@ -149,11 +149,10 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe
     auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(use.getOwner());
     if (!blueprint || blueprint->getParentOp() != blueprint->getParentOfType<func::FuncOp>())
       return failure();
-    std::optional<StringRef> mode = blueprint.getMode();
     std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
     std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
     std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
-    if (!mode || *mode != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr || !sourceSlotsAttr)
+    if (!spatial::isFragmentAssembly(blueprint.getMode()) || !operandIndicesAttr || !sourceOffsetsAttr || !sourceSlotsAttr)
       return failure();
     if (!blueprint.getOutput().hasOneUse() || !isa<func::ReturnOp>(*blueprint.getOutput().getUsers().begin()))
       return failure();
@@ -418,8 +417,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
     rewriter.setInsertionPointToEnd(newBlock);
 
     if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-      std::optional<StringRef> modeAttr = blueprint.getMode();
-      if (modeAttr && *modeAttr == "fragment_assembly") {
+      if (spatial::isFragmentAssembly(blueprint.getMode())) {
         for (Operation* user : blueprint.getOutput().getUsers()) {
           if (!isa<tensor::ParallelInsertSliceOp>(user))
             return blueprint.emitOpError(
@@ -483,8 +481,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
         auto hostTargetType = cast<ShapedType>(hostTarget.getType());
         if (auto blueprint =
               insertSlice.getSource().getDefiningOp<spatial::SpatBlueprintOp>()) {
-          std::optional<StringRef> modeAttr = blueprint.getMode();
-          if (modeAttr && *modeAttr == "fragment_assembly") {
+          if (spatial::isFragmentAssembly(blueprint.getMode())) {
             FailureOr<SmallVector<FragmentAssemblyCopy, 8>> fragmentAssemblyCopies =
               collectFragmentAssemblyCopiesFromBlueprint(blueprint, mapper, /*lane=*/0, /*hostTargetIndex=*/0);
             if (failed(fragmentAssemblyCopies))

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -129,6 +129,32 @@ LogicalResult validateFragmentAssemblyMetadata(spatial::SpatBlueprintOp blueprin
   return success();
 }
 
+FailureOr<mlir::Value> reshapeContiguousRowMajorFragments(RewriterBase& rewriter,
+                                                          Location loc,
+                                                          mlir::Value source,
+                                                          RankedTensorType resultType) {
+  auto sourceType = dyn_cast<RankedTensorType>(source.getType());
+  if (!sourceType || !sourceType.hasStaticShape() || !resultType.hasStaticShape() || resultType.getRank() < 2
+      || sourceType.getRank() != resultType.getRank() + 1 || sourceType.getElementType() != resultType.getElementType()
+      || sourceType.getNumElements() != resultType.getNumElements()
+      || sourceType.getDimSize(0) != getStaticShapeElementCount(resultType.getShape().drop_back())
+      || sourceType.getDimSize(sourceType.getRank() - 1) != resultType.getDimSize(resultType.getRank() - 1)
+      || llvm::any_of(sourceType.getShape().slice(1, sourceType.getRank() - 2), [](int64_t dim) { return dim != 1; }))
+    return failure();
+
+  SmallVector<ReassociationIndices> collapse {{}, {sourceType.getRank() - 1}};
+  for (int64_t dim = 0; dim < sourceType.getRank() - 1; ++dim)
+    collapse.front().push_back(dim);
+  auto flatType = RankedTensorType::get(
+    {sourceType.getDimSize(0), sourceType.getDimSize(sourceType.getRank() - 1)}, resultType.getElementType());
+  mlir::Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, source, collapse);
+
+  SmallVector<ReassociationIndices> expand {{}, {resultType.getRank() - 1}};
+  for (int64_t dim = 0; dim < resultType.getRank() - 1; ++dim)
+    expand.front().push_back(dim);
+  return tensor::ExpandShapeOp::create(rewriter, loc, resultType, flat, expand).getResult();
+}
+
 static SmallVector<int64_t, 4> expandFlatElementIndex(int64_t flatIndex, ArrayRef<int64_t> shape) {
   SmallVector<int64_t, 4> indices(shape.size(), 0);
   for (int64_t dim = static_cast<int64_t>(shape.size()) - 1; dim >= 0; --dim) {

src/PIM/Conversion/SpatialToPim/Common.hpp

@@ -51,6 +51,11 @@ mlir::LogicalResult validateFragmentAssemblyMetadata(onnx_mlir::spatial::SpatBlu
                                                      llvm::ArrayRef<int64_t> flatSizes,
                                                      llvm::ArrayRef<int64_t> flatStrides);
 
+mlir::FailureOr<mlir::Value> reshapeContiguousRowMajorFragments(mlir::RewriterBase& rewriter,
+                                                                mlir::Location loc,
+                                                                mlir::Value source,
+                                                                mlir::RankedTensorType resultType);
+
 mlir::FailureOr<mlir::SmallVector<int64_t, 4>>
 getStaticSliceOffsetsForElementOffset(mlir::Operation* anchor,
                                       mlir::ShapedType sourceType,

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -42,12 +42,11 @@ static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
   if (!resultType || !resultType.hasStaticShape())
     return blueprint.emitOpError("fragment assembly lowering requires a static ranked tensor result");
 
-  std::optional<StringRef> modeAttr = blueprint.getMode();
   std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
   std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
   std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
   std::optional<ArrayRef<int64_t>> fragmentStridesAttr = blueprint.getFragmentStrides();
-  if (!modeAttr || *modeAttr != "fragment_assembly" || !operandIndicesAttr || !sourceSlotsAttr
+  if (!spatial::isFragmentAssembly(blueprint.getMode()) || !operandIndicesAttr || !sourceSlotsAttr
       || !sourceOffsetsAttr || !fragmentStridesAttr)
     return blueprint.emitOpError("fragment assembly lowering requires explicit fragment metadata");
 
@@ -71,6 +70,16 @@ static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
                                               flatStrides)))
     return failure();
 
+  if (blueprint.getIndexMap() == spatial::kContiguousRowMajorFragments) {
+    if (!spatial::isCanonicalContiguousRowMajorFragmentAssembly(blueprint))
+      return blueprint.emitOpError("contiguous row-major fragment physical source order or storage is not canonical"), failure();
+    Value source = mapping.lookupOrDefault(blueprint.getInput());
+    auto reshaped = reshapeContiguousRowMajorFragments(
+      rewriter, blueprint.getLoc(), source, cast<RankedTensorType>(resultType));
+    if (failed(reshaped))
+      return blueprint.emitOpError("contiguous row-major fragment storage does not match its logical result"), failure();
+    return *reshaped;
+  }
   SmallVector<int64_t> hostStrides = computeRowMajorStrides(resultType.getShape());
   SmallVector<FragmentAssemblyCopy, 8> copies;
   for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {
@@ -193,8 +202,7 @@ static bool isHostMaterializableHelperOp(Operation* op) {
   if (isa<arith::ConstantOp>(op) || op->hasTrait<OpTrait::ConstantLike>())
     return true;
   if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-    std::optional<StringRef> mode = blueprint.getMode();
-    return mode && *mode == "fragment_assembly";
+    return spatial::isFragmentAssembly(blueprint.getMode());
   }
   return isShapingOnlyOp(op) || isPureIndexComputationOp(op);
 }
@@ -281,8 +289,7 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatSchedule
   }
   for (Operation& op : block.without_terminator()) {
     if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-      std::optional<StringRef> modeAttr = blueprint.getMode();
-      if (modeAttr && *modeAttr == "fragment_assembly") {
+      if (spatial::isFragmentAssembly(blueprint.getMode())) {
         auto lowered = lowerFragmentAssemblyBlueprint(rewriter, blueprint, mapping);
         if (failed(lowered))
           return false;

src/PIM/Conversion/SpatialToPim/Patterns.cpp

@@ -22,8 +22,7 @@ struct LowerFragmentAssemblyBlueprintPattern
   LogicalResult matchAndRewrite(spatial::SpatBlueprintOp op,
                                 OpAdaptor adaptor,
                                 ConversionPatternRewriter& rewriter) const override {
-    std::optional<StringRef> modeAttr = op.getMode();
-    if (!modeAttr || *modeAttr != "fragment_assembly")
+    if (!spatial::isFragmentAssembly(op.getMode()))
       return failure();
 
     auto resultType = dyn_cast<ShapedType>(op.getOutput().getType());
@@ -49,6 +48,16 @@ struct LowerFragmentAssemblyBlueprintPattern
           op, rank, fragmentOperands.size(), operandIndices, sourceOffsets, flatOffsets, flatSizes, flatStrides)))
       return failure();
 
+    if (op.getIndexMap() == spatial::kContiguousRowMajorFragments) {
+      if (!spatial::isCanonicalContiguousRowMajorFragmentAssembly(op))
+        return op.emitOpError("contiguous row-major fragment physical source order or storage is not canonical");
+      auto reshaped = reshapeContiguousRowMajorFragments(
+        rewriter, op.getLoc(), adaptor.getInput(), cast<RankedTensorType>(resultType));
+      if (failed(reshaped))
+        return op.emitOpError("contiguous row-major fragment storage does not match its logical result");
+      rewriter.replaceOp(op, *reshaped);
+      return success();
+    }
     Value currentOutput =
       tensor::EmptyOp::create(rewriter, op.getLoc(), resultType.getShape(), resultType.getElementType()).getResult();
     for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp

@@ -158,8 +158,7 @@ analyzeTopLevelFragmentAssemblyUses(Value value) {
     auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(use.getOwner());
     if (!blueprint || blueprint->getParentOp() != blueprint->getParentOfType<func::FuncOp>())
       return failure();
-    std::optional<StringRef> mode = blueprint.getMode();
-    if (!mode || *mode != "fragment_assembly")
+    if (!spatial::isFragmentAssembly(blueprint.getMode()))
       return failure();
     if (!blueprint.getOutput().hasOneUse() || !isa<func::ReturnOp>(*blueprint.getOutput().getUsers().begin()))
       return failure();
@@ -819,8 +818,7 @@ void raptor::SpatialToPimPass::replaceReturnWithOutputBuffers(func::ReturnOp ret
     }
 
     if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-      std::optional<StringRef> mode = blueprint.getMode();
-      if (mode && *mode == "fragment_assembly") {
+      if (spatial::isFragmentAssembly(blueprint.getMode())) {
         markOpToRemove(blueprint.getOperation());
         for (Value operand : blueprint->getOperands())
           markOwnedReturnChain(operand.getDefiningOp(), markOwnedReturnChain);

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

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.hpp

@@ -18,6 +18,7 @@
 #include "Conversion/SpatialToPim/Common.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
 
 namespace onnx_mlir {
 namespace raptor {
@@ -28,7 +29,10 @@ struct SpatialToPimPass : mlir::PassWrapper<SpatialToPimPass, mlir::OperationPas
   llvm::StringRef getDescription() const override { return "Lower Spatial ops to PIM-ready format"; }
 
   SpatialToPimPass() = default;
-  SpatialToPimPass(const SpatialToPimPass& pass) {}
+  explicit SpatialToPimPass(const spatial::SpatialTargetResources& target)
+      : targetResources(target), hasTarget(true) {}
+  SpatialToPimPass(const SpatialToPimPass& pass)
+      : targetResources(pass.targetResources), hasTarget(pass.hasTarget) {}
 
   void runOnOperation() final;
 
@@ -37,6 +41,8 @@ private:
 
   llvm::SmallVector<OutputTensorFactory> outputTensors;
   llvm::SmallVector<mlir::Operation*> operationsToRemove;
+  spatial::SpatialTargetResources targetResources;
+  bool hasTarget = false;
 
   mlir::LogicalResult allocateAndInitializeCoreLocalVariables(mlir::func::FuncOp funcOp, mlir::IRRewriter& rewriter);
   mlir::LogicalResult