Bose

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp

@@ -30,6 +30,7 @@
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ShapeTilingUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -308,7 +309,7 @@ struct PendingProjectedHostOutputFragment {
   Value originalOutput;
   ClassId sourceClass = 0;
   ProducerKey producerKey;
-  Value publicationValue;
+  unsigned publicationResultIndex = 0;
   int64_t sourceFragmentOrdinal = 0;
   int64_t sourceElementOffset = 0;
   SmallVector<int64_t, 4> offsets;
@@ -1220,36 +1221,13 @@ BlockArgument appendInput(MaterializerState& state, MaterializedClass& materiali
   return std::get<1>(*arg);
 }
 
-void refreshPendingProjectedHostOutputPublicationValues(MaterializerState& state,
-                                                        Operation* oldOwner,
-                                                        Operation* newOwner) {
-  if (!oldOwner || oldOwner == newOwner)
-    return;
-
-  for (PendingProjectedHostOutputFragment& fragment : state.pendingProjectedHostOutputFragments) {
-    auto publicationResult = dyn_cast_or_null<OpResult>(fragment.publicationValue);
-    if (!publicationResult || publicationResult.getOwner() != oldOwner)
-      publicationResult = OpResult();
-    else
-      fragment.publicationValue = newOwner->getResult(publicationResult.getResultNumber());
-
-    if (auto originalResult = dyn_cast_or_null<OpResult>(fragment.originalOutput); originalResult
-        && originalResult.getOwner() == oldOwner) {
-      fragment.originalOutput = newOwner->getResult(originalResult.getResultNumber());
-    }
-
-    if (fragment.producerKey.instance.op == oldOwner)
-      fragment.producerKey.instance.op = newOwner;
-  }
-}
-
-FailureOr<Value> appendScalarPublicationResult(MaterializerState& state,
-                                               MaterializedClass& materializedClass,
-                                               Value payload,
-                                               Location loc) {
+FailureOr<unsigned> appendScalarPublicationResult(MaterializerState& state,
+                                                  MaterializedClass& materializedClass,
+                                                  Value payload,
+                                                  Location loc) {
   auto existing = materializedClass.publicationOutputToResultIndex.find(payload);
   if (existing != materializedClass.publicationOutputToResultIndex.end())
-    return materializedClass.op->getResult(existing->second);
+    return existing->second;
 
   auto compute = dyn_cast<SpatScheduledCompute>(materializedClass.op);
   if (!compute)
@@ -1264,27 +1242,25 @@ FailureOr<Value> appendScalarPublicationResult(MaterializerState& state,
   if (failed(inserted))
     return materializedClass.op->emitError("failed to append scalar publication result");
 
-  Operation* oldOp = materializedClass.op;
   auto [result, newCompute] = *inserted;
   materializedClass.op = newCompute.getOperation();
   materializedClass.body = &newCompute.getBody().front();
-  refreshPendingProjectedHostOutputPublicationValues(state, oldOp, materializedClass.op);
   materializedClass.publicationOutputToResultIndex[payload] = result.getResultNumber();
 
   auto yieldOp = dyn_cast<SpatYieldOp>(materializedClass.body->getTerminator());
   if (!yieldOp)
     return materializedClass.op->emitError("expected spat.yield terminator while appending scalar publication result");
   state.rewriter.modifyOpInPlace(yieldOp, [&] { yieldOp->insertOperands(yieldOp.getNumOperands(), payload); });
-  return result;
+  return result.getResultNumber();
 }
 
-FailureOr<Value> appendBatchPublicationResult(MaterializerState& state,
-                                              MaterializedClass& materializedClass,
-                                              Value payload,
-                                              Location loc) {
+FailureOr<unsigned> appendBatchPublicationResult(MaterializerState& state,
+                                                 MaterializedClass& materializedClass,
+                                                 Value payload,
+                                                 Location loc) {
   auto existing = materializedClass.publicationOutputToResultIndex.find(payload);
   if (existing != materializedClass.publicationOutputToResultIndex.end())
-    return materializedClass.op->getResult(existing->second);
+    return existing->second;
 
   auto batch = dyn_cast<SpatScheduledComputeBatch>(materializedClass.op);
   if (!batch)
@@ -1305,11 +1281,9 @@ FailureOr<Value> appendBatchPublicationResult(MaterializerState& state,
   if (failed(inserted))
     return materializedClass.op->emitError("failed to append batch publication result");
 
-  Operation* oldOp = materializedClass.op;
   auto [result, outputArg, newBatch] = *inserted;
   materializedClass.op = newBatch.getOperation();
   materializedClass.body = &newBatch.getBody().front();
-  refreshPendingProjectedHostOutputPublicationValues(state, oldOp, materializedClass.op);
   materializedClass.publicationOutputToResultIndex[payload] = result.getResultNumber();
 
   auto inParallelOp = dyn_cast<SpatInParallelOp>(materializedClass.body->getTerminator());
@@ -1330,7 +1304,7 @@ FailureOr<Value> appendBatchPublicationResult(MaterializerState& state,
   Value firstOffset =
     scaleIndexByDim0Size(state, materializedClass.op, *laneArg, payloadType.getDimSize(0), loc);
   createDim0ParallelInsertSlice(state, loc, payload, outputArg, firstOffset);
-  return result;
+  return result.getResultNumber();
 }
 
 // -----------------------------------------------------------------------------
@@ -1563,7 +1537,7 @@ void attachMaterializerValueOriginNote(InFlightDiagnostic& diagnostic, Value val
 void attachMaterializedClassBodySummary(InFlightDiagnostic& diagnostic, const MaterializedClass& targetClass) {
   Block& body = *targetClass.body;
   diagnostic.attachNote(targetClass.op->getLoc())
-    << "RAPTOR_MATERIALIZER_DEBUG target class " << targetClass.id << " op '" << targetClass.op->getName()
+    << "target class " << targetClass.id << " op '" << targetClass.op->getName()
     << "' body has " << body.getNumArguments() << " block arguments and "
     << std::distance(body.begin(), body.end()) << " top-level operations";
 }
@@ -1687,7 +1661,7 @@ FailureOr<Value> rematerializeIndexValueInClass(MaterializerState& state,
 
   if (auto blockArg = dyn_cast<BlockArgument>(value)) {
     InFlightDiagnostic diagnostic = targetClass.op->emitError(
-      "RAPTOR_MATERIALIZER_DEBUG cannot rematerialize external block argument in materialized class body");
+      "cannot rematerialize external block argument in materialized class body");
     diagnostic << " currentArg#" << blockArg.getArgNumber() << " currentType=" << blockArg.getType()
                << " targetClass=" << targetClass.id << " targetOp='" << targetClass.op->getName() << "'";
     if (Operation* owner = blockArg.getOwner()->getParentOp()) {
@@ -1709,16 +1683,16 @@ FailureOr<Value> rematerializeIndexValueInClass(MaterializerState& state,
     if (mapperHadOriginalValue && mappedOriginalValue != value)
       attachMaterializerValueOriginNote(diagnostic, mappedOriginalValue, "mapper value");
     if (Operation* owner = blockArg.getOwner()->getParentOp()) {
-      attachMaterializerOperationPrintNote(diagnostic, owner, "RAPTOR_MATERIALIZER_DEBUG external block argument owner op");
-      attachMaterializerParentChainNote(diagnostic, owner, "RAPTOR_MATERIALIZER_DEBUG external block argument owner parent chain");
+      attachMaterializerOperationPrintNote(diagnostic, owner, "external block argument owner op");
+      attachMaterializerParentChainNote(diagnostic, owner, "external block argument owner parent chain");
     }
-    attachMaterializerOperationPrintNote(diagnostic, targetClass.op, "RAPTOR_MATERIALIZER_DEBUG target materialized op");
+    attachMaterializerOperationPrintNote(diagnostic, targetClass.op, "target materialized op");
     attachMaterializedClassBodySummary(diagnostic, targetClass);
     return failure();
   }
 
   InFlightDiagnostic diagnostic =
-    targetClass.op->emitError("RAPTOR_MATERIALIZER_DEBUG cannot rematerialize external index value in materialized class body");
+    targetClass.op->emitError("cannot rematerialize external index value in materialized class body");
   diagnostic << " type=" << value.getType() << " targetClass=" << targetClass.id << " targetOp='"
              << targetClass.op->getName() << "'";
   attachMaterializerValueOriginNote(diagnostic, originalValue, "original value");
@@ -1793,8 +1767,12 @@ FailureOr<Value> rematerializeTensorValueInClass(MaterializerState& state,
       strides.push_back(*localized);
     }
 
-    return tensor::ExtractSliceOp::create(state.rewriter, anchor->getLoc(), *localizedSource, offsets, sizes, strides)
-      .getResult();
+    auto resultType = dyn_cast<RankedTensorType>(extractSlice.getResult().getType());
+    if (!resultType)
+      return anchor->emitError("expected ranked tensor extract_slice while rematerializing tensor capture");
+
+    return extractStaticSliceOrIdentity(
+      state.rewriter, anchor->getLoc(), *localizedSource, resultType, offsets, sizes, strides);
   }
 
   if (auto collapseShape = value.getDefiningOp<tensor::CollapseShapeOp>()) {
@@ -2108,8 +2086,10 @@ Value scaleIndexByDim0Size(MaterializerState& state, Operation* anchor, Value in
   if (dim0Size == 1)
     return index;
 
-  Value dim0SizeValue = getOrCreateIndexConstant(state.constantFolder, anchor, dim0Size);
-  return arith::MulIOp::create(state.rewriter, loc, index, dim0SizeValue).getResult();
+  MLIRContext* context = state.func.getContext();
+  AffineExpr d0 = getAffineDimExpr(0, context);
+  AffineMap map = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, d0 * dim0Size);
+  return createOrFoldAffineApply(state.rewriter, loc, map, ValueRange {index}, anchor);
 }
 
 FailureOr<Value> scaleIndexByDim0SizeInClass(MaterializerState& state,
@@ -2123,8 +2103,7 @@ FailureOr<Value> scaleIndexByDim0SizeInClass(MaterializerState& state,
   if (dim0Size == 1)
     return *localizedIndex;
 
-  Value dim0SizeValue = getOrCreateIndexConstant(state.constantFolder, targetClass.op, dim0Size);
-  return arith::MulIOp::create(state.rewriter, loc, *localizedIndex, dim0SizeValue).getResult();
+  return scaleIndexByDim0Size(state, targetClass.op, *localizedIndex, dim0Size, loc);
 }
 
 bool sameProducerResult(ProducerKey lhs, ProducerKey rhs) {
@@ -3677,10 +3656,13 @@ FailureOr<Value> buildProjectedPackedPayload(MaterializerState& state,
     ValueRange {init},
     [&](OpBuilder&, Location, Value fragmentIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
       Value acc = iterArgs.front();
-      Value payloadFragmentCount =
-        getOrCreateIndexConstant(state.constantFolder, targetClass.op, descriptor.layout.payloadFragmentCount);
-      Value flatBase = arith::MulIOp::create(state.rewriter, loc, *localizedMessageIndex, payloadFragmentCount).getResult();
-      Value flatIndex = arith::AddIOp::create(state.rewriter, loc, flatBase, fragmentIndex).getResult();
+      MLIRContext* context = state.func.getContext();
+      AffineExpr d0 = getAffineDimExpr(0, context);
+      AffineExpr d1 = getAffineDimExpr(1, context);
+      AffineMap flatIndexMap =
+        AffineMap::get(/*dimCount=*/2, /*symbolCount=*/0, d0 * descriptor.layout.payloadFragmentCount + d1);
+      Value flatIndex = createOrFoldAffineApply(
+        state.rewriter, loc, flatIndexMap, ValueRange {*localizedMessageIndex, fragmentIndex}, targetClass.op);
 
       FailureOr<SmallVector<OpFoldResult, 4>> fragmentOffsets =
         buildProjectedFragmentOffsetsInClass(state, targetClass, descriptor, flatIndex, loc);
@@ -5618,8 +5600,8 @@ FailureOr<bool> recordProjectedScalarHostFragmentsFromPackedRun(MaterializerStat
     return failure();
   }
 
-  FailureOr<Value> publicationResult = appendScalarPublicationResult(state, sourceClass, packed, loc);
-  if (failed(publicationResult))
+  FailureOr<unsigned> publicationResultIndex = appendScalarPublicationResult(state, sourceClass, packed, loc);
+  if (failed(publicationResultIndex))
     return failure();
 
   int64_t fragmentElementCount = fragmentType.getNumElements();
@@ -5657,7 +5639,7 @@ FailureOr<bool> recordProjectedScalarHostFragmentsFromPackedRun(MaterializerStat
       originalOutput,
       sourceClass.id,
       ProducerKey {peer, resultIndex},
-      *publicationResult,
+      *publicationResultIndex,
       static_cast<int64_t>(runIndex),
       static_cast<int64_t>(runIndex) * fragmentElementCount,
       SmallVector<int64_t, 4>(*offsets),
@@ -5711,8 +5693,8 @@ FailureOr<bool> recordProjectedScalarHostFragmentsFromPackedValue(MaterializerSt
   if (fragmentType == originalOutput.getType())
     return false;
 
-  FailureOr<Value> publicationResult = appendBatchPublicationResult(state, sourceClass, packed, loc);
-  if (failed(publicationResult))
+  FailureOr<unsigned> publicationResultIndex = appendBatchPublicationResult(state, sourceClass, packed, loc);
+  if (failed(publicationResultIndex))
     return failure();
 
   if (packedType != fragmentType) {
@@ -5764,7 +5746,7 @@ FailureOr<bool> recordProjectedScalarHostFragmentsFromPackedValue(MaterializerSt
       originalOutput,
       sourceClass.id,
       key,
-      *publicationResult,
+      *publicationResultIndex,
       static_cast<int64_t>(fragmentIndex),
       static_cast<int64_t>(*publishedLaneIndex) * payloadElementCount + localFragmentOffsetWithinPublishedPayload,
       SmallVector<int64_t, 4>(*offsets),
@@ -5787,18 +5769,26 @@ LogicalResult finalizeProjectedHostOutputFragments(MaterializerState& state) {
 
   SmallVector<Value, 8> outputs;
   outputs.reserve(byOutput.size());
-  for (const auto& entry : byOutput)
-    outputs.push_back(entry.first);
-  llvm::sort(outputs, [](Value lhs, Value rhs) {
-    return reinterpret_cast<uintptr_t>(lhs.getAsOpaquePointer())
-         < reinterpret_cast<uintptr_t>(rhs.getAsOpaquePointer());
-  });
 
   auto returnOp = dyn_cast<func::ReturnOp>(state.func.getBody().front().getTerminator());
   if (!returnOp)
     return state.func.emitError("expected func.return terminator while finalizing projected host output fragments");
 
+  DenseSet<Value> seenOutputs;
+  for (Value returned : returnOp.getOperands()) {
+    if (!byOutput.contains(returned) || !seenOutputs.insert(returned).second)
+      continue;
+    outputs.push_back(returned);
+  }
+  if (outputs.size() != byOutput.size())
+    return state.func.emitError("projected host output fragments must be keyed by returned logical host outputs");
+
   for (Value originalOutput : outputs) {
+    if (isa_and_present<SpatScheduledCompute, SpatScheduledComputeBatch>(originalOutput.getDefiningOp())) {
+      return state.func.emitError(
+        "projected host output assembly must be keyed by the original logical host output, not by a materialized scheduled result");
+    }
+
     auto resultType = dyn_cast<RankedTensorType>(originalOutput.getType());
     if (!resultType || !resultType.hasStaticShape())
       return state.func.emitError("projected host output must have static ranked tensor type");
@@ -5806,13 +5796,12 @@ LogicalResult finalizeProjectedHostOutputFragments(MaterializerState& state) {
     SmallVector<PendingProjectedHostOutputFragment*, 16>& fragments = byOutput[originalOutput];
     llvm::sort(fragments, [](const PendingProjectedHostOutputFragment* lhs,
                              const PendingProjectedHostOutputFragment* rhs) {
-      if (lhs->publicationValue != rhs->publicationValue)
-        return reinterpret_cast<uintptr_t>(lhs->publicationValue.getAsOpaquePointer())
-             < reinterpret_cast<uintptr_t>(rhs->publicationValue.getAsOpaquePointer());
-      if (lhs->sourceFragmentOrdinal != rhs->sourceFragmentOrdinal)
-        return lhs->sourceFragmentOrdinal < rhs->sourceFragmentOrdinal;
       if (lhs->sourceClass != rhs->sourceClass)
         return lhs->sourceClass < rhs->sourceClass;
+      if (lhs->publicationResultIndex != rhs->publicationResultIndex)
+        return lhs->publicationResultIndex < rhs->publicationResultIndex;
+      if (lhs->sourceFragmentOrdinal != rhs->sourceFragmentOrdinal)
+        return lhs->sourceFragmentOrdinal < rhs->sourceFragmentOrdinal;
       return std::lexicographical_compare(lhs->offsets.begin(),
                                           lhs->offsets.end(),
                                           rhs->offsets.begin(),
@@ -5821,7 +5810,7 @@ LogicalResult finalizeProjectedHostOutputFragments(MaterializerState& state) {
 
     state.rewriter.setInsertionPoint(returnOp);
     Location loc = fragments.front()->loc;
-    SmallVector<Value, 16> reconciliatorOperands;
+    SmallVector<Value, 16> blueprintOperands;
     SmallVector<int64_t, 16> fragmentOperandIndices;
     SmallVector<int64_t, 16> fragmentSourceOffsets;
     SmallVector<int64_t, 64> flatOffsets;
@@ -5830,12 +5819,23 @@ LogicalResult finalizeProjectedHostOutputFragments(MaterializerState& state) {
     DenseMap<Value, int64_t> operandIndicesByValue;
 
     for (PendingProjectedHostOutputFragment* fragmentRecord : fragments) {
-      Value operand = fragmentRecord->publicationValue;
+      if (fragmentRecord->sourceClass >= state.classes.size())
+        return state.func.emitError("projected host output fragment references an invalid source class");
+
+      MaterializedClass& sourceClass = state.classes[fragmentRecord->sourceClass];
+      if (fragmentRecord->publicationResultIndex >= sourceClass.op->getNumResults()) {
+        return sourceClass.op->emitError("projected host output fragment references an invalid publication result")
+               << " sourceClass=" << sourceClass.id
+               << " resultIndex=" << fragmentRecord->publicationResultIndex
+               << " resultCount=" << sourceClass.op->getNumResults();
+      }
+
+      Value operand = sourceClass.op->getResult(fragmentRecord->publicationResultIndex);
 
       auto [operandIt, inserted] =
-        operandIndicesByValue.try_emplace(operand, static_cast<int64_t>(reconciliatorOperands.size()));
+        operandIndicesByValue.try_emplace(operand, static_cast<int64_t>(blueprintOperands.size()));
       if (inserted)
-        reconciliatorOperands.push_back(operand);
+        blueprintOperands.push_back(operand);
       fragmentOperandIndices.push_back(operandIt->second);
       fragmentSourceOffsets.push_back(fragmentRecord->sourceElementOffset);
       llvm::append_range(flatOffsets, fragmentRecord->offsets);
@@ -5847,12 +5847,12 @@ LogicalResult finalizeProjectedHostOutputFragments(MaterializerState& state) {
         return state.func.emitError("projected host output assembly requires static ranked tensor operands");
     }
 
-    if (reconciliatorOperands.empty())
+    if (blueprintOperands.empty())
       return state.func.emitError("missing projected host output fragments");
 
-    Value input = reconciliatorOperands.front();
-    ValueRange extraFragments = ValueRange(reconciliatorOperands).drop_front();
-    auto reconciliator = spatial::SpatReconciliatorOp::create(
+    Value input = blueprintOperands.front();
+    ValueRange extraFragments = ValueRange(blueprintOperands).drop_front();
+    auto blueprint = spatial::SpatBlueprintOp::create(
       state.rewriter,
       loc,
       resultType,
@@ -5870,7 +5870,7 @@ LogicalResult finalizeProjectedHostOutputFragments(MaterializerState& state) {
       state.rewriter.getStringAttr("disjoint"),
       state.rewriter.getStringAttr("complete"));
 
-    state.hostReplacements[originalOutput] = reconciliator.getOutput();
+    state.hostReplacements[originalOutput] = blueprint.getOutput();
   }
 
   return success();
@@ -6284,6 +6284,32 @@ LogicalResult cloneComputeTemplateBody(MaterializerState& state,
         mapper.map(operand, *localized);
     }
 
+    if (auto extract = dyn_cast<tensor::ExtractSliceOp>(&op)) {
+      auto remapFoldResult = [&](OpFoldResult value) -> OpFoldResult {
+        if (auto mappedValue = dyn_cast_if_present<Value>(value))
+          return mapper.lookupOrDefault(mappedValue);
+        return value;
+      };
+
+      SmallVector<OpFoldResult, 4> offsets;
+      SmallVector<OpFoldResult, 4> sizes;
+      SmallVector<OpFoldResult, 4> strides;
+      offsets.reserve(extract.getMixedOffsets().size());
+      sizes.reserve(extract.getMixedSizes().size());
+      strides.reserve(extract.getMixedStrides().size());
+
+      llvm::append_range(offsets, llvm::map_range(extract.getMixedOffsets(), remapFoldResult));
+      llvm::append_range(sizes, llvm::map_range(extract.getMixedSizes(), remapFoldResult));
+      llvm::append_range(strides, llvm::map_range(extract.getMixedStrides(), remapFoldResult));
+
+      auto resultType = cast<RankedTensorType>(extract.getType());
+      Value localizedSource = mapper.lookupOrDefault(extract.getSource());
+      Value localizedExtract = extractStaticSliceOrIdentity(
+        state.rewriter, extract.getLoc(), localizedSource, resultType, offsets, sizes, strides);
+      mapper.map(extract.getResult(), localizedExtract);
+      continue;
+    }
+
     Operation* cloned = state.rewriter.clone(op, mapper);
     if (failed(mapClonedRegionBlockArguments(op, *cloned, mapper)))
       return failure();
@@ -6350,18 +6376,20 @@ FailureOr<Value> materializeProjectedExtractReplacement(MaterializerState& state
       if (failed(localizedIv))
         return failure();
       Value iv = *localizedIv;
-      Value lowerBound =
-        getOrCreateIndexConstant(state.constantFolder, targetClass.op, replacement.layout.loopLowerBounds[index]);
-      Value step = getOrCreateIndexConstant(state.constantFolder, targetClass.op, replacement.layout.loopSteps[index]);
-      Value tripCount =
-        getOrCreateIndexConstant(state.constantFolder, targetClass.op, replacement.layout.loopTripCounts[index]);
+      MLIRContext* context = state.func.getContext();
+      AffineExpr d0 = getAffineDimExpr(0, context);
+      AffineMap normalizedMap =
+        AffineMap::get(/*dimCount=*/1,
+                       /*symbolCount=*/0,
+                       (d0 - replacement.layout.loopLowerBounds[index]).floorDiv(replacement.layout.loopSteps[index]));
+      Value normalized =
+        createOrFoldAffineApply(state.rewriter, extract.getLoc(), normalizedMap, ValueRange {iv}, targetClass.op);
 
-      Value normalized = arith::SubIOp::create(state.rewriter, extract.getLoc(), iv, lowerBound).getResult();
-      if (replacement.layout.loopSteps[index] != 1)
-        normalized = arith::DivUIOp::create(state.rewriter, extract.getLoc(), normalized, step).getResult();
-      linearizedIndex = arith::MulIOp::create(state.rewriter, extract.getLoc(), linearizedIndex, tripCount).getResult();
-      linearizedIndex =
-        arith::AddIOp::create(state.rewriter, extract.getLoc(), linearizedIndex, normalized).getResult();
+      AffineExpr d1 = getAffineDimExpr(1, context);
+      AffineMap linearizedMap = AffineMap::get(
+        /*dimCount=*/2, /*symbolCount=*/0, d0 * replacement.layout.loopTripCounts[index] + d1);
+      linearizedIndex = createOrFoldAffineApply(
+        state.rewriter, extract.getLoc(), linearizedMap, ValueRange {linearizedIndex, normalized}, targetClass.op);
     }
     return linearizedIndex;
   };
@@ -6386,12 +6414,16 @@ FailureOr<Value> materializeProjectedExtractReplacement(MaterializerState& state
     if (failed(localProjectionSlotIndex))
       return failure();
 
-    Value fragmentsPerLogicalSlot =
-      getOrCreateIndexConstant(state.constantFolder, targetClass.op, replacement.layout.fragmentsPerLogicalSlot);
-    Value base =
-      arith::MulIOp::create(state.rewriter, extract.getLoc(), *localProjectionSlotIndex, fragmentsPerLogicalSlot)
-        .getResult();
-    return arith::AddIOp::create(state.rewriter, extract.getLoc(), base, intraSlotFragmentIndex).getResult();
+    MLIRContext* context = state.func.getContext();
+    AffineExpr d0 = getAffineDimExpr(0, context);
+    AffineExpr d1 = getAffineDimExpr(1, context);
+    AffineMap packedIndexMap = AffineMap::get(
+      /*dimCount=*/2, /*symbolCount=*/0, d0 * replacement.layout.fragmentsPerLogicalSlot + d1);
+    return createOrFoldAffineApply(state.rewriter,
+                                   extract.getLoc(),
+                                   packedIndexMap,
+                                   ValueRange {*localProjectionSlotIndex, intraSlotFragmentIndex},
+                                   targetClass.op);
   };
 
   FailureOr<Value> packedFragmentIndex = computeProjectedPayloadFragmentIndex();
@@ -6445,18 +6477,18 @@ LogicalResult localizeCapturesInOperationTree(MaterializerState& state,
         localizeMaterializedClassOperand(state, targetClass, current, nestedOp, tensorContext, genericContext, mapper);
       if (failed(localized)) {
         InFlightDiagnostic diagnostic = targetClass.op->emitError(
-          "RAPTOR_MATERIALIZER_DEBUG failed to localize cloned scheduled-body operand");
+          "failed to localize cloned scheduled-body operand");
         diagnostic << " targetClass=" << targetClass.id << " nestedOp='" << nestedOp->getName()
                    << "' operand#" << operand.getOperandNumber() << " operandType=" << current.getType()
                    << " offendingIR=\"" << truncateMaterializerDebugString(stringifyOperationForMaterializerDebug(nestedOp))
                    << "\" offendingOperands=\"" << formatMaterializerOperandListInline(nestedOp, targetClass)
                    << "\" parentChain=\"" << formatMaterializerParentChainInline(nestedOp) << "\"";
         diagnostic.attachNote(nestedOp->getLoc()) << "offending nested operation";
-        attachMaterializerOperationPrintNote(diagnostic, nestedOp, "RAPTOR_MATERIALIZER_DEBUG offending nested operation IR");
-        attachMaterializerOperandListNote(diagnostic, nestedOp, targetClass, "RAPTOR_MATERIALIZER_DEBUG offending nested operation operands");
-        attachMaterializerParentChainNote(diagnostic, nestedOp, "RAPTOR_MATERIALIZER_DEBUG offending nested operation parent chain");
+        attachMaterializerOperationPrintNote(diagnostic, nestedOp, "offending nested operation IR");
+        attachMaterializerOperandListNote(diagnostic, nestedOp, targetClass, "offending nested operation operands");
+        attachMaterializerParentChainNote(diagnostic, nestedOp, "offending nested operation parent chain");
         attachMaterializerValueOriginNote(diagnostic, current, "offending operand");
-        attachMaterializerOperationPrintNote(diagnostic, targetClass.op, "RAPTOR_MATERIALIZER_DEBUG target materialized op");
+        attachMaterializerOperationPrintNote(diagnostic, targetClass.op, "target materialized op");
         attachMaterializedClassBodySummary(diagnostic, targetClass);
         return WalkResult::interrupt();
       }
@@ -6505,7 +6537,7 @@ LogicalResult localizeAllScheduledBodyCaptures(MaterializerState& state, Materia
         "final scheduled body capture localization found an unsupported external non-tensor operand");
       if (failed(localized)) {
         InFlightDiagnostic diagnostic = targetClass.op->emitError(
-          "RAPTOR_MATERIALIZER_DEBUG failed to localize final scheduled-body operand");
+          "failed to localize final scheduled-body operand");
         diagnostic << " targetClass=" << targetClass.id << " nestedOp='" << nestedOp->getName()
                    << "' operand#" << operand.getOperandNumber() << " operandType=" << current.getType()
                    << " offendingIR=\"" << truncateMaterializerDebugString(stringifyOperationForMaterializerDebug(nestedOp))

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp.rej

@@ -1,128 +0,0 @@
---- src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp	2026-06-24 18:51:29.043731129 +0000
-+++ src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp	2026-06-24 18:51:29.026726895 +0000
-@@ -4112,104 +4112,8 @@
-                                                                 Value originalOutput,
-                                                                 Location loc);
- 
--FailureOr<SmallVector<OpFoldResult, 4>> rematerializeProjectionIndexListForBatchHostOutput(
--  MaterializerState& state,
--  MaterializedClass& sourceClass,
--  ArrayRef<OpFoldResult> values,
--  IRMapping& mapper,
--  Location loc) {
--  SmallVector<OpFoldResult, 4> localized;
--  localized.reserve(values.size());
--  for (OpFoldResult value : values) {
--    FailureOr<OpFoldResult> remapped =
--      rematerializeIndexOpFoldResultInClass(state, sourceClass, value, loc, &mapper);
--    if (failed(remapped))
--      return failure();
--    localized.push_back(*remapped);
--  }
--  return localized;
--}
--
--LogicalResult createProjectionAwareBatchHostInsert(MaterializerState& state,
--                                                  MaterializedClass& sourceClass,
--                                                  Value originalOutput,
--                                                  Value payload,
--                                                  Value destination,
--                                                  ArrayRef<ProducerKey> keys,
--                                                  Location loc) {
--  auto originalResult = dyn_cast<OpResult>(originalOutput);
--  if (!originalResult)
--    return failure();
--
--  auto sourceBatch = dyn_cast_or_null<SpatComputeBatch>(originalResult.getOwner());
--  if (!sourceBatch || sourceBatch.getNumResults() == 0)
--    return failure();
--
--  FailureOr<tensor::ParallelInsertSliceOp> projection =
--    getBatchResultProjectionInsert(sourceBatch, originalResult.getResultNumber());
--  if (failed(projection))
--    return failure();
--
--  auto sourceLaneArg = sourceBatch.getLaneArgument();
--  if (!sourceLaneArg)
--    return failure();
--
--  auto materializedBatch = dyn_cast<SpatScheduledComputeBatch>(sourceClass.op);
--  if (!materializedBatch)
--    return failure();
--
--  auto materializedLaneArg = materializedBatch.getLaneArgument();
--  if (!materializedLaneArg)
--    return failure();
--
--  if (keys.size() != sourceClass.cpus.size())
--    return failure();
--
--  SmallVector<int64_t, 8> logicalLanes;
--  logicalLanes.reserve(keys.size());
--  for (ProducerKey key : keys) {
--    if (key.instance.op != sourceBatch.getOperation() || key.resultIndex != originalResult.getResultNumber())
--      return failure();
--    logicalLanes.push_back(key.instance.laneStart);
--  }
--
--  IRMapping mapper;
--  Value logicalLane = createIndexedIndexValue(state,
--                                             sourceClass.op,
--                                             ArrayRef<int64_t>(logicalLanes),
--                                             *materializedLaneArg,
--                                             loc,
--                                             static_cast<int64_t>(sourceClass.cpus.size()),
--                                             /*allowExhaustiveTiledSearch=*/false);
--  mapper.map(*sourceLaneArg, logicalLane);
--
--  FailureOr<SmallVector<OpFoldResult, 4>> offsets =
--    rematerializeProjectionIndexListForBatchHostOutput(
--      state, sourceClass, projection->getMixedOffsets(), mapper, loc);
--  if (failed(offsets))
--    return failure();
--  FailureOr<SmallVector<OpFoldResult, 4>> sizes =
--    rematerializeProjectionIndexListForBatchHostOutput(
--      state, sourceClass, projection->getMixedSizes(), mapper, loc);
--  if (failed(sizes))
--    return failure();
--  FailureOr<SmallVector<OpFoldResult, 4>> strides =
--    rematerializeProjectionIndexListForBatchHostOutput(
--      state, sourceClass, projection->getMixedStrides(), mapper, loc);
--  if (failed(strides))
--    return failure();
--
--  tensor::ParallelInsertSliceOp::create(
--    state.rewriter, loc, payload, destination, *offsets, *sizes, *strides);
--  return success();
--}
--
- LogicalResult
--setHostOutputValue(MaterializerState& state,
--                   MaterializedClass& sourceClass,
--                   Value originalOutput,
--                   Value payload,
--                   ArrayRef<ProducerKey> keys = {}) {
-+setHostOutputValue(MaterializerState& state, MaterializedClass& sourceClass, Value originalOutput, Value payload) {
-   auto resultIt = sourceClass.hostOutputToResultIndex.find(originalOutput);
-   if (resultIt == sourceClass.hostOutputToResultIndex.end())
-     return sourceClass.op->emitError("missing host result slot for materialized output")
-@@ -4253,10 +4157,6 @@
-     return batch.emitOpError("expected compute_batch output block argument while materializing batch output");
- 
-   state.rewriter.setInsertionPointToStart(&inParallelOp.getRegion().front());
--  if (succeeded(createProjectionAwareBatchHostInsert(
--        state, sourceClass, originalOutput, payload, *outputArg, keys, payload.getLoc())))
--    return success();
--
-   createDim0ParallelInsertSlice(state, payload.getLoc(), payload, *outputArg, *laneArg);
-   return success();
- }
-@@ -4276,7 +4176,7 @@
- 
-   MaterializedClass& ownerClass = state.classes[ownerIt->second];
-   if (sourceClass.id == ownerClass.id)
--    return setHostOutputValue(state, ownerClass, originalOutput, payload, keys);
-+    return setHostOutputValue(state, ownerClass, originalOutput, payload);
- 
-   // Keep the old deadlock-free communication discipline: only scalar-to-scalar
-   // host-owner forwarding is introduced here. Batch host publication remains on

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp

@@ -354,13 +354,13 @@ public:
   void runOnOperation() override {
     func::FuncOp func = getOperation();
     if (failed(verifyLogicalSpatialGraphInvariants(func))) {
-      func.emitOpError("RAPTOR_PHASE_CHECK logical Spatial graph verification failed at the start of MergeComputeNodes");
+      func.emitOpError("logical Spatial graph verification failed at the start of MergeComputeNodes");
       signalPassFailure();
       return;
     }
     mergeTriviallyConnectedComputes(func);
     if (failed(verifyLogicalSpatialGraphInvariants(func))) {
-      func.emitOpError("RAPTOR_PHASE_CHECK logical Spatial graph verification failed after trivial merge simplification");
+      func.emitOpError("logical Spatial graph verification failed after trivial merge simplification");
       signalPassFailure();
       return;
     }
@@ -378,7 +378,7 @@ public:
       return;
     }
     if (failed(verifyScheduledSpatialInvariants(func))) {
-      func.emitOpError("RAPTOR_PHASE_CHECK scheduled Spatial verification failed after merge materialization");
+      func.emitOpError("scheduled Spatial verification failed after merge materialization");
       signalPassFailure();
       return;
     }