big refactor

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) {