Merge with fast resnet

2026-07-20 18:01:58 +02:00
parent 5c94e00f43 ab54243fda
commit 961fd613bd
85 changed files with 4475 additions and 4455 deletions
@@ -13,6 +13,7 @@
 #include <utility>

 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"

@@ -374,9 +375,6 @@ extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
  auto physicalType = mlir::dyn_cast<mlir::RankedTensorType>(physicalBatch.getType());
  if (!physicalType || physicalType.getRank() != fragmentType.getRank() + 1)
    return mlir::failure();
-  mlir::SmallVector<int64_t> selectedShape {1};
-  llvm::append_range(selectedShape, fragmentType.getShape());
-  auto selectedType = mlir::RankedTensorType::get(selectedShape, fragmentType.getElementType(), fragmentType.getEncoding());
  mlir::SmallVector<mlir::OpFoldResult> offsets {slot};
  mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
  mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
@@ -385,11 +383,8 @@ extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
    sizes.push_back(rewriter.getIndexAttr(dim));
    strides.push_back(rewriter.getIndexAttr(1));
  }
-  mlir::Value selected = mlir::tensor::ExtractSliceOp::create(rewriter, loc, selectedType, physicalBatch, offsets, sizes, strides);
-  mlir::SmallVector<mlir::ReassociationIndices> reassociation {{0, 1}};
-  for (int64_t dim = 2; dim <= fragmentType.getRank(); ++dim)
-    reassociation.push_back({dim});
-  return mlir::tensor::CollapseShapeOp::create(rewriter, loc, fragmentType, selected, reassociation).getResult();
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, physicalBatch, fragmentType, {offsets, sizes, strides});
 }

 template <typename BodyFn>
@@ -3,6 +3,7 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"

 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"

 using namespace mlir;
@@ -38,6 +39,30 @@ Value createPaddedInputCompute(Value input,
  if (inputType == paddedInputType)
    return input;

+  auto producer = inputType.getRank() == 2 && paddedInputType.getRank() == 2
+    ? input.getDefiningOp<spatial::SpatGraphComputeBatch>()
+    : spatial::SpatGraphComputeBatch();
+  auto inputFragmentType = producer
+    ? spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount())
+    : FailureOr<RankedTensorType>(failure());
+  auto paddedFragmentType = producer
+    ? spatial::getGraphBatchFragmentType(paddedInputType, producer.getLaneCount())
+    : FailureOr<RankedTensorType>(failure());
+  if (producer && succeeded(inputFragmentType) && succeeded(paddedFragmentType)) {
+    auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {paddedInputType}, producer.getLaneCount(), {}, input,
+      [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
+        auto fragment = extractGraphBatchPhysicalFragment(
+          rewriter, loc, args.inputs.front(), args.lane, *inputFragmentType);
+        if (failed(fragment))
+          return failure();
+        Value padded = createZeroPaddedTensor(*fragment, *paddedFragmentType, rewriter, loc);
+        publishGraphBatchPhysicalFragment(rewriter, loc, padded, args.outputs.front(), args.lane);
+        return success();
+      });
+    if (succeeded(batch))
+      return batch->getResult(0);
+  }
+
  auto computeOp = createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, input, [&](Value computeInput) {
    Value paddedInput = createZeroPaddedTensor(computeInput, paddedInputType, rewriter, loc);
    spatial::SpatYieldOp::create(rewriter, loc, paddedInput);
@@ -79,6 +79,57 @@ materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location
  return createRowStripAssemblyBlueprint(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
 }

+static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType,
+                                               PatternRewriter& rewriter, Location loc) {
+  auto producer = input.getDefiningOp<spatial::SpatGraphComputeBatch>();
+  auto inputType = dyn_cast<RankedTensorType>(input.getType());
+  auto biasType = dyn_cast<RankedTensorType>(bias.getType());
+  if (!producer || !inputType || !biasType || !inputType.hasStaticShape() || !biasType.hasStaticShape()
+      || !resultType.hasStaticShape() || inputType.getDimSize(0) != producer.getLaneCount()
+      || biasType.getDimSize(0) != producer.getLaneCount() || resultType.getDimSize(0) != producer.getLaneCount())
+    return failure();
+  auto inputFragmentType = spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount());
+  auto outputFragmentType = spatial::getGraphBatchFragmentType(resultType, producer.getLaneCount());
+  if (failed(inputFragmentType) || failed(outputFragmentType) || inputFragmentType->getRank() != biasType.getRank()
+      || inputFragmentType->getDimSize(0) != 1 || inputFragmentType->getShape().drop_front() != biasType.getShape().drop_front()
+      || inputFragmentType->getRank() != outputFragmentType->getRank() + 1)
+    return failure();
+  for (auto [inputDim, outputDim] : llvm::zip(inputFragmentType->getShape().drop_front(), outputFragmentType->getShape()))
+    if (outputDim > inputDim)
+      return failure();
+
+  auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {resultType}, producer.getLaneCount(), {}, ValueRange {input, bias},
+    [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
+      FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, *inputFragmentType);
+      if (failed(fragment))
+        return failure();
+      MixedSliceGeometry biasSlice;
+      for (int64_t dim : inputFragmentType->getShape()) {
+        biasSlice.offsets.push_back(biasSlice.offsets.empty() ? OpFoldResult(args.lane) : rewriter.getIndexAttr(0));
+        biasSlice.sizes.push_back(rewriter.getIndexAttr(dim));
+        biasSlice.strides.push_back(rewriter.getIndexAttr(1));
+      }
+      Value biasFragment = extractMixedSliceOrIdentity(rewriter, loc, args.inputs[1], *inputFragmentType, biasSlice);
+      if (!biasFragment)
+        return failure();
+      Value added = spatial::SpatVAddOp::create(rewriter, loc, *inputFragmentType, *fragment, biasFragment);
+      MixedSliceGeometry outputSlice;
+      outputSlice.offsets.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(0));
+      outputSlice.sizes.push_back(rewriter.getIndexAttr(1));
+      outputSlice.strides.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(1));
+      for (int64_t dim : outputFragmentType->getShape())
+        outputSlice.sizes.push_back(rewriter.getIndexAttr(dim));
+      Value output = extractMixedSliceOrIdentity(rewriter, loc, added, *outputFragmentType, outputSlice);
+      if (!output)
+        return failure();
+      publishGraphBatchPhysicalFragment(rewriter, loc, output, args.outputs.front(), args.lane);
+      return success();
+    });
+  if (failed(batch))
+    return failure();
+  return batch->getResult(0);
+}
+
 struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass)

@@ -267,6 +318,13 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
          signalPassFailure();
          return;
        }
+        if (planOp.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
+          FailureOr<Value> lowered = lowerDenseBatchBiasAdd(planOp.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
+          if (succeeded(lowered)) {
+            rewriter.replaceOp(planOp, *lowered);
+            continue;
+          }
+        }
        auto computeOp = createSpatCompute<2>(rewriter,
                                              planOp.getLoc(),
                                              planOp.getOutput().getType(),
@@ -20,6 +20,7 @@

 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
@@ -1355,20 +1356,11 @@ static Value createWeightTile(Value packedWeights,
  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
                                   rewriter.getIndexAttr(tiling.tileInputRows),
                                   rewriter.getIndexAttr(tiling.tileOutputChannels)};
-  auto sliceType =
-    RankedTensorType::get({1, tiling.tileInputRows, tiling.tileOutputChannels}, packedWeightType.getElementType());
-  Value slice = tensor::ExtractSliceOp::create(
-    rewriter, loc, sliceType, packedWeights, offsets, sizes, getUnitStrides(rewriter, 3));
  auto collapsedType =
    RankedTensorType::get({tiling.tileInputRows, tiling.tileOutputChannels}, packedWeightType.getElementType());
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         collapsedType,
-                                         slice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, packedWeights, collapsedType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }

 static Value createBiasTile(
@@ -1676,8 +1668,6 @@ struct ConvGemmPlan {
  int64_t effectiveMaxParallelPixels;
  int64_t packedNumRows;

-  RankedTensorType im2colType;
-  RankedTensorType im2colRowType;
  RankedTensorType gemmInputRowsType;
  RankedTensorType wFlatType;
  RankedTensorType wTransType;
@@ -1718,52 +1708,6 @@ static PreparedConvInput prepareInputForIm2Col(const ConvLoweringState& state,
  return {paddedInputOp.getResult(0), paddedType};
 }

-static Value createPaddedRows(Value rows,
-                              RankedTensorType rowsType,
-                              int64_t paddedRows,
-                              PatternRewriter& rewriter,
-                              Location loc) {
-  if (rowsType.getDimSize(0) == paddedRows)
-    return rows;
-
-  auto paddedType =
-    RankedTensorType::get({paddedRows, rowsType.getDimSize(1)}, rowsType.getElementType(), rowsType.getEncoding());
-  return createZeroPaddedTensor(
-    rows, paddedType, {0, 0}, {paddedRows - rowsType.getDimSize(0), 0}, rewriter, loc);
-}
-
-static Value packRowsForParallelGemm(
-  Value rows, RankedTensorType rowsType, int64_t packFactor, PatternRewriter& rewriter, Location loc) {
-  if (packFactor == 1)
-    return rows;
-
-  const int64_t paddedNumRows = ceilIntegerDivide(rowsType.getDimSize(0), packFactor) * packFactor;
-  const int64_t packedNumRows = paddedNumRows / packFactor;
-  const int64_t rowWidth = rowsType.getDimSize(1);
-  auto groupedType =
-    RankedTensorType::get({packedNumRows, packFactor, rowWidth}, rowsType.getElementType(), rowsType.getEncoding());
-  auto packedType =
-    RankedTensorType::get({packedNumRows, packFactor * rowWidth}, rowsType.getElementType(), rowsType.getEncoding());
-
-  Value padded = createPaddedRows(rows, rowsType, paddedNumRows, rewriter, loc);
-  Value grouped = tensor::ExpandShapeOp::create(rewriter,
-                                                loc,
-                                                groupedType,
-                                                padded,
-                                                SmallVector<ReassociationIndices> {
-                                                  {0, 1},
-                                                  {2}
-  });
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         packedType,
-                                         grouped,
-                                         SmallVector<ReassociationIndices> {
-                                           {0},
-                                           {1, 2}
-  });
-}
-
 static Value unpackRowsFromParallelGemm(Value packedRows,
                                        RankedTensorType packedRowsType,
                                        int64_t unpackedRows,
@@ -2197,8 +2141,6 @@ buildConvGemmPlan(const ConvLoweringState& state,

  auto elemType = state.xType.getElementType();
  auto outElemType = state.outType.getElementType();
-  plan.im2colType = RankedTensorType::get({plan.chunkNumPatches, plan.patchSize}, elemType);
-  plan.im2colRowType = RankedTensorType::get({1, plan.patchSize}, elemType);
  plan.gemmInputRowsType =
    RankedTensorType::get({plan.packedNumRows, plan.effectiveMaxParallelPixels * plan.patchSize}, elemType);
  plan.wFlatType = RankedTensorType::get({state.numChannelsOut, plan.patchSize}, state.wType.getElementType());
@@ -2216,44 +2158,99 @@ static Value createIm2colRows(const ConvLoweringState& state,
                              const ConvGemmPlan& plan,
                              PatternRewriter& rewriter,
                              Location loc) {
-  constexpr size_t numInputs = 1;
-  auto im2colComputeOp =
-    createSpatCompute<numInputs>(rewriter, loc, TypeRange {plan.gemmInputRowsType}, {}, preparedInput.value, [&](Value xArg) {
-      auto elemType = preparedInput.type.getElementType();
-      // Keep the standard im2col view of convolution, flipped so filters sit in
-      // B / crossbar columns:
-      //   A (im2col):  [numPatches, patchSize]  -- one row per output spatial position
-      //   B (weights): [patchSize, cOut]
-      //   Gemm output: [numPatches, cOut]
-      Value im2colInit = tensor::EmptyOp::create(rewriter, loc, plan.im2colType.getShape(), elemType);
+  if (plan.gemmInputRowsType.getDimSize(1) > crossbarSize.getValue()) {
+    assert(plan.effectiveMaxParallelPixels == 1 && "multi-crossbar im2col rows cannot pack pixels");
+    auto compute = createSpatCompute<1>(
+      rewriter, loc, TypeRange {plan.gemmInputRowsType}, {}, preparedInput.value, [&](Value input) {
+        auto elemType = preparedInput.type.getElementType();
+        Value empty = tensor::EmptyOp::create(rewriter, loc, plan.gemmInputRowsType.getShape(), elemType);
+        Operation *anchor = rewriter.getInsertionBlock()->getParentOp();
+        Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0);
+        Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1);
+        Value upper = getOrCreateIndexConstant(rewriter, anchor, plan.chunkNumPatches);
+        auto patchType = RankedTensorType::get(
+          {1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
+        auto rowType = RankedTensorType::get({plan.patchSize}, elemType);
+        auto loop = buildNormalizedScfFor(
+          rewriter, loc, c0, upper, c1, ValueRange {empty},
+          [&](OpBuilder &, Location nestedLoc, Value patchIndex, ValueRange iterArgs,
+              SmallVectorImpl<Value> &yielded) {
+            Value batchIndex = affineAddFloorDivConst(
+              rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchor);
+            Value batchPatchIndex = affineAddModConst(
+              rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchor);
+            Value outHeight = affineFloorDivConst(
+              rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchor);
+            Value outWidth = affineModConst(
+              rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchor);
+            Value patch = createConvInputPatch(
+              input, patchType, batchIndex, c0,
+              affineMulConst(rewriter, nestedLoc, outHeight, state.strideHeight, anchor),
+              affineMulConst(rewriter, nestedLoc, outWidth, state.strideWidth, anchor),
+              state.dilationHeight, state.dilationWidth, rewriter, nestedLoc);
+            Value row = tensor::CollapseShapeOp::create(
+              rewriter, nestedLoc, rowType, patch,
+              SmallVector<ReassociationIndices> {{0, 1, 2, 3}});
+            Value next = tensor::InsertSliceOp::create(
+              rewriter, nestedLoc, row, iterArgs.front(),
+              SmallVector<OpFoldResult> {patchIndex, rewriter.getIndexAttr(0)},
+              SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(plan.patchSize)},
+              getUnitStrides(rewriter, 2));
+            yielded.push_back(next);
+            return success();
+          });
+        if (failed(loop))
+          return failure();
+        spatial::SpatYieldOp::create(rewriter, loc, loop->results.front());
+        return success();
+      });
+    assert(succeeded(compute) && "Conv im2col compute construction must succeed");
+    return compute->getResult(0);
+  }
+
+  auto elemType = preparedInput.type.getElementType();
+  auto packedRowType = RankedTensorType::get(
+    {plan.effectiveMaxParallelPixels * plan.patchSize}, elemType, plan.gemmInputRowsType.getEncoding());
+  auto zeroAttr = DenseElementsAttr::get(packedRowType, rewriter.getZeroAttr(elemType));
+  Value zeroRow = getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, packedRowType);
+  auto im2colComputeOp = createSpatComputeBatch(
+    rewriter,
+    loc,
+    TypeRange {plan.gemmInputRowsType},
+    plan.packedNumRows,
+    {},
+    ValueRange {preparedInput.value, zeroRow},
+    [&](detail::SpatComputeBatchBodyArgs args) {
      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
      Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
      Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
+      Value cPack = getOrCreateIndexConstant(rewriter, anchorOp, plan.effectiveMaxParallelPixels);
      Value cNumPatches = getOrCreateIndexConstant(rewriter, anchorOp, plan.chunkNumPatches);
+      Value laneStart = affineMulConst(rewriter, loc, args.lane, plan.effectiveMaxParallelPixels, anchorOp);
+      Value remaining = arith::SubIOp::create(rewriter, loc, cNumPatches, laneStart);
+      Value isPartial = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::ult, remaining, cPack);
+      Value lanePatches = arith::SelectOp::create(rewriter, loc, isPartial, remaining, cPack);
+      auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
+      auto patchRowType = RankedTensorType::get({plan.patchSize}, elemType);

-      auto im2colLoop = buildNormalizedScfFor(
+      auto rowLoop = buildNormalizedScfFor(
        rewriter,
        loc,
        c0,
-        cNumPatches,
+        lanePatches,
        c1,
-        ValueRange {im2colInit},
-        [&](OpBuilder&, Location nestedLoc, Value patchIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
-          Value im2colAcc = iterArgs.front();
+        ValueRange {args.inputs[1]},
+        [&](OpBuilder&, Location nestedLoc, Value copyIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+          Value patchIndex = arith::AddIOp::create(rewriter, nestedLoc, laneStart, copyIndex);
          Value batchIndex =
            affineAddFloorDivConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp);
          Value batchPatchIndex =
            affineAddModConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp);
          Value outHeightIndex = affineFloorDivConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp);
          Value outWidthIndex = affineModConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp);
-          Value inputHeightOffset =
-            affineMulConst(rewriter, nestedLoc, outHeightIndex, state.strideHeight, anchorOp);
-          Value inputWidthOffset =
-            affineMulConst(rewriter, nestedLoc, outWidthIndex, state.strideWidth, anchorOp);
-
-          auto patchType =
-            RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
-          Value patch = createConvInputPatch(xArg,
+          Value inputHeightOffset = affineMulConst(rewriter, nestedLoc, outHeightIndex, state.strideHeight, anchorOp);
+          Value inputWidthOffset = affineMulConst(rewriter, nestedLoc, outWidthIndex, state.strideWidth, anchorOp);
+          Value patch = createConvInputPatch(args.inputs.front(),
                                             patchType,
                                             batchIndex,
                                             c0,
@@ -2263,34 +2260,27 @@ static Value createIm2colRows(const ConvLoweringState& state,
                                             state.dilationWidth,
                                             rewriter,
                                             nestedLoc);
-          Value row = tensor::CollapseShapeOp::create(rewriter,
-                                                      nestedLoc,
-                                                      plan.im2colRowType,
-                                                      patch,
-                                                      SmallVector<ReassociationIndices> {
-                                                        {0},
-                                                        {1, 2, 3}
+          Value patchRow = tensor::CollapseShapeOp::create(rewriter,
+                                                           nestedLoc,
+                                                           patchRowType,
+                                                           patch,
+                                                           SmallVector<ReassociationIndices> {
+                                                             {0, 1, 2, 3}
          });
-
-          SmallVector<OpFoldResult> rowOffsets {patchIndex, rewriter.getIndexAttr(0)};
-          SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(plan.patchSize)};
-          Value next = tensor::InsertSliceOp::create(
-            rewriter, nestedLoc, row, im2colAcc, rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
+          Value rowOffset = affineMulConst(rewriter, nestedLoc, copyIndex, plan.patchSize, anchorOp);
+          Value next = tensor::InsertSliceOp::create(rewriter,
+                                                     nestedLoc,
+                                                     patchRow,
+                                                     iterArgs.front(),
+                                                     SmallVector<OpFoldResult> {rowOffset},
+                                                     SmallVector<OpFoldResult> {rewriter.getIndexAttr(plan.patchSize)},
+                                                     getUnitStrides(rewriter, 1));
          yielded.push_back(next);
          return success();
        });
-      if (failed(im2colLoop))
+      if (failed(rowLoop))
        return failure();
-
-      Value gemmInputRows = im2colLoop->results.front();
-      // Pack N old im2col rows into one longer row so one GEMM can cover N
-      // pixels in parallel. The corresponding packed weight matrix contains N
-      // block-diagonal copies of W^T, and the packed output must be unpacked
-      // back to one row per spatial patch.
-      if (plan.effectiveMaxParallelPixels != 1)
-        gemmInputRows = packRowsForParallelGemm(gemmInputRows, plan.im2colType, plan.effectiveMaxParallelPixels, rewriter, loc);
-
-      spatial::SpatYieldOp::create(rewriter, loc, gemmInputRows);
+      publishGraphBatchPhysicalFragment(rewriter, loc, rowLoop->results.front(), args.outputs.front(), args.lane);
      return success();
    });

@@ -17,6 +17,7 @@
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
@@ -501,9 +502,9 @@ static Value extractReductionPiece(Value partialPiecesArg,
  SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
  SmallVector<OpFoldResult> pieceOffsets {
    createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-  auto selectedType = RankedTensorType::get({numOutRows, 1, static_cast<int64_t>(crossbarSize.getValue())}, pieceType.getElementType());
-  Value selected = tensor::ExtractSliceOp::create(rewriter, loc, selectedType, partialPiecesArg, pieceOffsets, pieceSizes, unitStrides);
-  return tensor::CollapseShapeOp::create(rewriter, loc, pieceType, selected, SmallVector<ReassociationIndices> {{0, 1}, {2}});
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, partialPiecesArg, pieceType,
+    {pieceOffsets, pieceSizes, unitStrides});
 }

 static Value reducePartialPiecesForHSlice(Value partialPiecesArg,
@@ -9,6 +9,7 @@

 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
@@ -299,22 +300,14 @@ static Value extractBatchedATile(Value a,
                                 RankedTensorType aTileType,
                                 PatternRewriter& rewriter,
                                 Location loc) {
-  auto aSliceType = RankedTensorType::get({1, 1, aTileType.getDimSize(1)}, aTileType.getElementType());
  Value sourceBatchIndex =
    mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
  SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), row, kOffset};
  SmallVector<OpFoldResult> sizes {
    rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(aTileType.getDimSize(1))};
-  auto slice =
-    tensor::ExtractSliceOp::create(rewriter, loc, aSliceType, a, offsets, sizes, getUnitStrides(rewriter, 3));
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         aTileType,
-                                         slice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, a, aTileType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }

 static Value extractBatchedBTile(Value b,
@@ -326,24 +319,15 @@ static Value extractBatchedBTile(Value b,
                                 RankedTensorType bTileType,
                                 PatternRewriter& rewriter,
                                 Location loc) {
-  auto bSliceType =
-    RankedTensorType::get({1, bTileType.getDimSize(0), bTileType.getDimSize(1)}, bTileType.getElementType());
  Value sourceBatchIndex =
    mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
  SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), kOffset, hOffset};
  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
                                   rewriter.getIndexAttr(bTileType.getDimSize(0)),
                                   rewriter.getIndexAttr(bTileType.getDimSize(1))};
-  auto slice =
-    tensor::ExtractSliceOp::create(rewriter, loc, bSliceType, b, offsets, sizes, getUnitStrides(rewriter, 3));
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         bTileType,
-                                         slice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, b, bTileType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }

 static Value getBatchLaneIndex(
@@ -448,22 +432,14 @@ static Value extractDynamicBatchedRowVector(Value matrix,
                                            RankedTensorType vectorType,
                                            PatternRewriter& rewriter,
                                            Location loc) {
-  auto rowSliceType = RankedTensorType::get({1, 1, vectorType.getDimSize(1)}, vectorType.getElementType());
  Value sourceBatchIndex =
    mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
  SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), row, rewriter.getIndexAttr(0)};
  SmallVector<OpFoldResult> sizes {
    rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1))};
-  auto rowSlice =
-    tensor::ExtractSliceOp::create(rewriter, loc, rowSliceType, matrix, offsets, sizes, getUnitStrides(rewriter, 3));
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         vectorType,
-                                         rowSlice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, matrix, vectorType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }

 static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
@@ -519,7 +495,6 @@ static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
  const int64_t numOutRows = outType.getDimSize(1);
  const int64_t numOutCols = outType.getDimSize(2);
  auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
-  auto outputScalarType = RankedTensorType::get({1, 1, 1}, outType.getElementType());

  auto computeOp = createSpatCompute<1>(
    rewriter, loc, TypeRange {outType}, {}, ValueRange {scalarPieces}, [&](Value pieces) -> LogicalResult {
@@ -545,20 +520,12 @@ static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
          FailureOr<Value> scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType);
          if (failed(scalar))
            return failure();
-          Value expanded = tensor::ExpandShapeOp::create(rewriter,
-                                                         nestedLoc,
-                                                         outputScalarType,
-                                                         *scalar,
-                                                         SmallVector<ReassociationIndices> {
-                                                           {0},
-                                                           {1, 2}
-          });
          SmallVector<OpFoldResult> outputOffsets {batch, row, column};
          SmallVector<OpFoldResult> outputSizes = {
            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
          Value next =
            tensor::InsertSliceOp::create(
-              rewriter, nestedLoc, expanded, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
+              rewriter, nestedLoc, *scalar, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
              .getResult();
          yielded.push_back(next);
          return success();
@@ -591,9 +558,9 @@ static Value extractBatchedReductionPiece(Value partialPiecesArg,
  Value pieceOffset = arith::AddIOp::create(rewriter, loc, batchAndHSlice, kOffset);
  SmallVector<OpFoldResult> offsets {pieceOffset, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
-  auto selectedType = RankedTensorType::get({numOutRows, 1, static_cast<int64_t>(crossbarSize.getValue())}, pieceType.getElementType());
-  Value selected = tensor::ExtractSliceOp::create(rewriter, loc, selectedType, partialPiecesArg, offsets, sizes, getUnitStrides(rewriter, 3));
-  return tensor::CollapseShapeOp::create(rewriter, loc, pieceType, selected, SmallVector<ReassociationIndices> {{0, 1}, {2}});
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, partialPiecesArg, pieceType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }

 static Value reduceBatchedPartialPiecesForHSlice(Value partialPiecesArg,
@@ -640,8 +607,6 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
      const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(2), crossbarSize.getValue());
      auto pieceType = RankedTensorType::get({numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
                                             partialPiecesType.getElementType());
-      auto outputSliceType = RankedTensorType::get({1, numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
-                                                   partialPiecesType.getElementType());

      Value outputInit =
        tensor::EmptyOp::create(rewriter, loc, paddedOutType.getShape(), paddedOutType.getElementType()).getResult();
@@ -671,14 +636,6 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
              Value outputAcc = hIterArgs.front();
              Value reduced = reduceBatchedPartialPiecesForHSlice(
                partialPiecesArg, batch, hSlice, pieceType, numKSlices, numOutHSlices, numOutRows, rewriter, hLoc);
-              Value expandedReduced = tensor::ExpandShapeOp::create(rewriter,
-                                                                    hLoc,
-                                                                    outputSliceType,
-                                                                    reduced,
-                                                                    SmallVector<ReassociationIndices> {
-                                                                      {0, 1},
-                                                                      {2}
-              });
              Value hOffset = affineMulConst(
                rewriter, hLoc, hSlice, crossbarSize.getValue(), rewriter.getInsertionBlock()->getParentOp());
              SmallVector<OpFoldResult> outputOffsets {batch, rewriter.getIndexAttr(0), hOffset};
@@ -687,7 +644,7 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
                                                     rewriter.getIndexAttr(crossbarSize.getValue())};
              Value next =
                tensor::InsertSliceOp::create(
-                  rewriter, hLoc, expandedReduced, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
+                  rewriter, hLoc, reduced, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
                  .getResult();
              hYielded.push_back(next);
              return success();
@@ -56,12 +56,14 @@ collectFragmentAssemblyCopiesFromBlueprint(spatial::SpatBlueprintOp blueprint,
    return blueprint.emitOpError("fragment assembly lowering requires static ranked tensor results");

  std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
+  std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
  std::optional<ArrayRef<int64_t>> fragmentStridesAttr = blueprint.getFragmentStrides();
-  if (!operandIndicesAttr || !fragmentStridesAttr)
+  if (!operandIndicesAttr || !sourceSlotsAttr || !fragmentStridesAttr)
    return blueprint.emitOpError(
-      "fragment assembly lowering requires explicit operand indices and unit strides");
+      "fragment assembly lowering requires explicit operand indices, source slots, and unit strides");

  ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+  ArrayRef<int64_t> sourceSlots = *sourceSlotsAttr;
  std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
  if (!sourceOffsetsAttr)
    return blueprint.emitOpError("fragment assembly lowering requires explicit source offsets");
@@ -115,7 +117,11 @@ collectFragmentAssemblyCopiesFromBlueprint(spatial::SpatBlueprintOp blueprint,
            copy.sourceType = sourceType;
            copy.hostTargetIndex = hostTargetIndex;
            copy.lane = lane;
-            copy.sourceByteOffset = (sourceOffsets[fragmentIndex] + relativeSourceOffset) * static_cast<int64_t>(elementSize);
+            copy.sourceByteOffset =
+              (getFragmentAssemblySourceElementOffset(
+                 sourceType, sourceSlots[fragmentIndex], sourceOffsets[fragmentIndex])
+               + relativeSourceOffset)
+              * static_cast<int64_t>(elementSize);
            copy.hostByteOffset = hostElementOffset * static_cast<int64_t>(elementSize);
            copy.byteSize = chunkElements * static_cast<int64_t>(elementSize);
            copies.push_back(copy);
@@ -183,8 +189,8 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe
      if (operandIndices[fragmentIndex] != static_cast<int64_t>(use.getOperandNumber()))
        continue;

-      int64_t sourceElementOffset =
-        sourceSlots[fragmentIndex] * payloadElementCount + sourceOffsets[fragmentIndex];
+      int64_t sourceElementOffset = getFragmentAssemblySourceElementOffset(
+        packedResultType, sourceSlots[fragmentIndex], sourceOffsets[fragmentIndex]);
      int64_t lane = sourceElementOffset / payloadElementCount;
      if (lane < 0 || lane >= static_cast<int64_t>(laneCount))
        return failure();
@@ -193,6 +193,14 @@ forEachContiguousDestinationChunk(ArrayRef<int64_t> destShape,
  return visit(visit, 0);
 }

+int64_t getFragmentAssemblySourceElementOffset(RankedTensorType sourceType,
+                                               int64_t sourceSlot,
+                                               int64_t sourceOffset) {
+  assert(sourceType.getRank() > 0 && sourceType.hasStaticShape()
+         && "fragment assembly source must have a static leading slot dimension");
+  return sourceSlot * (sourceType.getNumElements() / sourceType.getDimSize(0)) + sourceOffset;
+}
+
 static mlir::Value
 createSteppedOffset(OpBuilder& builder, Location loc, mlir::Value start, mlir::Value index,
                    int64_t stepBytes, Operation *constantAnchor) {
@@ -65,6 +65,10 @@ forEachContiguousDestinationChunk(llvm::ArrayRef<int64_t> destShape,
                                  llvm::function_ref<mlir::LogicalResult(llvm::ArrayRef<int64_t>, int64_t, int64_t)>
                                    callback);

+int64_t getFragmentAssemblySourceElementOffset(mlir::RankedTensorType sourceType,
+                                               int64_t sourceSlot,
+                                               int64_t sourceOffset);
+
 struct FragmentAssemblyCopy {
  mlir::Value source;
  mlir::RankedTensorType sourceType;
@@ -44,13 +44,15 @@ static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,

  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 || !sourceOffsetsAttr
-      || !fragmentStridesAttr)
+  if (!modeAttr || *modeAttr != "fragment_assembly" || !operandIndicesAttr || !sourceSlotsAttr
+      || !sourceOffsetsAttr || !fragmentStridesAttr)
    return blueprint.emitOpError("fragment assembly lowering requires explicit fragment metadata");

  ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+  ArrayRef<int64_t> sourceSlots = *sourceSlotsAttr;
  ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
  ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
  ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
@@ -102,7 +104,10 @@ static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
            copy.source = source;
            copy.sourceType = sourceType;
            copy.sourceByteOffset =
-              (sourceOffsets[fragmentIndex] + relativeSourceOffset) * static_cast<int64_t>(elementSize);
+              (getFragmentAssemblySourceElementOffset(
+                 sourceType, sourceSlots[fragmentIndex], sourceOffsets[fragmentIndex])
+               + relativeSourceOffset)
+              * static_cast<int64_t>(elementSize);
            copy.hostByteOffset = hostElementOffset * static_cast<int64_t>(elementSize);
            copy.byteSize = chunkElements * static_cast<int64_t>(elementSize);
            copies.push_back(copy);
@@ -1,7 +1,8 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
-
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"

@@ -18,6 +19,30 @@ static void copyRaptorDebugAttrs(Operation* source, Operation* target) {
  }
 }

+static Value createDestinationByteOffset(PatternRewriter& rewriter,
+                                         tensor::InsertSliceOp insert) {
+  auto destinationType = cast<RankedTensorType>(insert.getDestType());
+  SmallVector<int64_t> strides = computeRowMajorStrides(destinationType.getShape());
+  int64_t elementBytes = getElementTypeSizeInBytes(destinationType.getElementType());
+  Value total = arith::ConstantIndexOp::create(rewriter, insert.getLoc(), 0);
+  for (auto [dimension, offset] : llvm::enumerate(insert.getMixedOffsets())) {
+    int64_t scale = strides[dimension] * elementBytes;
+    Value component;
+    if (auto attribute = dyn_cast<Attribute>(offset)) {
+      component = arith::ConstantIndexOp::create(
+        rewriter, insert.getLoc(), cast<IntegerAttr>(attribute).getInt() * scale);
+    } else {
+      component = cast<Value>(offset);
+      if (scale != 1)
+        component = arith::MulIOp::create(
+          rewriter, insert.getLoc(), component,
+          arith::ConstantIndexOp::create(rewriter, insert.getLoc(), scale));
+    }
+    total = arith::AddIOp::create(rewriter, insert.getLoc(), total, component);
+  }
+  return total;
+}
+
 struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
  using OpRewritePattern::OpRewritePattern;

@@ -40,7 +65,21 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
      rewriter.eraseOp(op);
      return success();
    }
-    auto outputType = cast<ShapedType>(op.getResult().getType());
+    auto outputType = cast<RankedTensorType>(op.getResult().getType());
+    tensor::InsertSliceOp destinationInsert;
+    if (op->hasOneUse()) {
+      auto insert = dyn_cast<tensor::InsertSliceOp>(*op->getUsers().begin());
+      auto destinationType = insert
+        ? dyn_cast<RankedTensorType>(insert.getDestType()) : RankedTensorType();
+      if (insert && insert.getSource() == op.getOutput()
+          && insert.getSourceType() == outputType
+          && insert->getBlock() == op->getBlock() && destinationType
+          && destinationType.hasStaticShape()
+          && isContiguousSubviewWithDynamicOffsets(
+            destinationType.getShape(), insert.getMixedOffsets(),
+            insert.getStaticSizes(), insert.getStaticStrides()))
+        destinationInsert = insert;
+    }
    Value outputBuffer =
      tensor::EmptyOp::create(rewriter, op.getLoc(), outputType.getShape(), outputType.getElementType()).getResult();
    auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getResult());
@@ -50,7 +89,19 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
      rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, *sizeAttr, op.getSourceCoreId());
    copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
    Value received = receive.getOutput();
-    rewriter.replaceOp(op, received);
+    if (!destinationInsert) {
+      rewriter.replaceOp(op, received);
+      return success();
+    }
+
+    rewriter.setInsertionPoint(destinationInsert);
+    Value targetOffset = createDestinationByteOffset(rewriter, destinationInsert);
+    Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
+    auto copy = pim::PimMemCopyOp::create(
+      rewriter, op.getLoc(), destinationInsert.getDestType(), targetOffset, zero,
+      destinationInsert.getDest(), received, *sizeAttr);
+    rewriter.replaceOp(destinationInsert, copy.getOutput());
+    rewriter.eraseOp(op);
    return success();
  }
 };
@@ -608,11 +608,12 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
    for (auto [blueprint, operandNumber] : *fragmentAssemblyUses) {
      rewriter.setInsertionPointAfterValue(storedValue);
      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>> stridesAttr = blueprint.getFragmentStrides();
-      if (!operandIndicesAttr || !sourceOffsetsAttr || !stridesAttr) {
+      if (!operandIndicesAttr || !sourceSlotsAttr || !sourceOffsetsAttr || !stridesAttr) {
        blueprint.emitOpError(
-          "fragment assembly lowering requires explicit operand, source-offset, and stride metadata");
+          "fragment assembly lowering requires explicit operand, source-slot, source-offset, and stride metadata");
        return ReturnPathLoweringResult::Failure;
      }

@@ -626,6 +627,7 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
      }

      ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+      ArrayRef<int64_t> sourceSlots = *sourceSlotsAttr;
      ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
      ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
      ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
@@ -668,7 +670,9 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
                                   elementSize,
                                   producerOp,
                                   "fragment assembly host offset");
-            auto sourceOffset = getCheckedByteOffset(sourceOffsets[fragmentIndex] + relativeSourceOffset,
+            int64_t sourceElementOffset = getFragmentAssemblySourceElementOffset(
+              sourceType, sourceSlots[fragmentIndex], sourceOffsets[fragmentIndex]);
+            auto sourceOffset = getCheckedByteOffset(sourceElementOffset + relativeSourceOffset,
                                                     elementSize,
                                                     producerOp,
                                                     "fragment assembly source offset");
@@ -12,7 +12,7 @@ include "src/Accelerators/PIM/Dialect/Pim/Pim.td"
 def spatToPimVMM : Pat<
  (SpatVMMOp:$srcOpRes $weight, $vector),
  (PimVMMOp $weight, $vector,
-    (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
+    (NativeCodeCall<"tensor::EmptyOp::create($_builder, $_loc, cast<ShapedType>($0.getType()).getShape(), cast<ShapedType>($0.getType()).getElementType())"> $srcOpRes))
 >;

 def spatToPimVVDMul : Pat<
@@ -1,4 +1,3 @@
-#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
@@ -175,11 +174,11 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {

  RewritePatternSet coreBodyPatterns(ctx);
  populateCoreBodyPatterns(coreBodyPatterns);
-  populateAffineToStdConversionPatterns(coreBodyPatterns);
  FrozenRewritePatternSet frozenCoreBodyPatterns(std::move(coreBodyPatterns));

  ConversionTarget coreBodyTarget(*ctx);
-  coreBodyTarget.addLegalDialect<PimDialect,
+  coreBodyTarget.addLegalDialect<affine::AffineDialect,
+                                 PimDialect,
                                 tensor::TensorDialect,
                                 arith::ArithDialect,
                                 bufferization::BufferizationDialect,
@@ -226,7 +225,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
  eraseUnusedTensorPackingOps(funcOp, rewriter);

  ConversionTarget communicationTarget(*ctx);
-  communicationTarget.addLegalDialect<PimDialect,
+  communicationTarget.addLegalDialect<affine::AffineDialect,
+                                      PimDialect,
                                      tensor::TensorDialect,
                                      arith::ArithDialect,
                                      bufferization::BufferizationDialect,