meno diamantini

src/PIM/CMakeLists.txt

@@ -117,7 +117,6 @@ add_pim_library(OMPIMAccel
   SpatialOps
   PimOps
   OMONNXToSpatial
-  OMSpatialToGraphviz
   OMSpatialToPim
   OMPimCommon
   OMPimBufferization

src/PIM/Common/IR/CoreBlockUtils.cpp

@@ -74,6 +74,21 @@ walkPimCoreBlock(mlir::Block& block,
       continue;
     }
 
+    if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(op)) {
+      auto condition = resolveIndexValue(ifOp.getCondition(), knowledge);
+      if (failed(condition)) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
+        hasFailure = true;
+        continue;
+      }
+
+      mlir::Region& selectedRegion = *condition != 0 ? ifOp.getThenRegion() : ifOp.getElseRegion();
+      if (!selectedRegion.empty())
+        if (failed(walkPimCoreBlock(selectedRegion.front(), knowledge, callback)))
+          hasFailure = true;
+      continue;
+    }
+
     if (failed(callback(op, knowledge)))
       hasFailure = true;
   }
@@ -128,6 +143,22 @@ mlir::LogicalResult walkPimCoreBlockStructurally(
       continue;
     }
 
+    if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(op)) {
+      if (failed(resolveIndexValue(ifOp.getCondition(), knowledge))) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM verification");
+        hasFailure = true;
+        continue;
+      }
+
+      if (!ifOp.getThenRegion().empty())
+        if (failed(walkPimCoreBlockStructurally(ifOp.getThenRegion().front(), knowledge, callback)))
+          hasFailure = true;
+      if (!ifOp.getElseRegion().empty())
+        if (failed(walkPimCoreBlockStructurally(ifOp.getElseRegion().front(), knowledge, callback)))
+          hasFailure = true;
+      continue;
+    }
+
     if (failed(callback(op, knowledge)))
       hasFailure = true;
   }

src/PIM/Compiler/PimCodeGen.cpp

@@ -414,31 +414,35 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value value,
                                              const StaticValueKnowledge& knowledge,
                                              std::optional<unsigned> lane) const {
   value = resolveCachedAlias(value, knowledge);
-  auto compiledIt = compiledAddressExprs.find(value);
+
-  if (compiledIt == compiledAddressExprs.end()) {
+  FailureOr<ResolvedContiguousAddress> resolvedAddress = resolveContiguousAddress(value, knowledge);
-    auto compiledExpr = compileContiguousAddressExpr(value);
+  if (failed(resolvedAddress)) {
-    if (failed(compiledExpr)) {
+    auto compiledIt = compiledAddressExprs.find(value);
-      errs() << "Failed to compile contiguous address for value: ";
+    if (compiledIt == compiledAddressExprs.end()) {
+      auto compiledExpr = compileContiguousAddressExpr(value);
+      if (failed(compiledExpr)) {
+        errs() << "Failed to compile contiguous address for value: ";
+        value.print(errs());
+        errs() << " : " << value.getType();
+        errs() << "\n";
+        llvm_unreachable("Failed to compile contiguous address");
+      }
+      compiledIt = compiledAddressExprs.try_emplace(value, *compiledExpr).first;
+    }
+
+    resolvedAddress = compiledIt->second.evaluate(knowledge, lane);
+    if (failed(resolvedAddress)) {
+      errs() << "Failed to evaluate contiguous address for value: ";
       value.print(errs());
       errs() << " : " << value.getType();
       errs() << "\n";
-      llvm_unreachable("Failed to compile contiguous address");
+      if (auto* definingOp = value.getDefiningOp()) {
+        errs() << "Defining op:\n";
+        definingOp->print(errs());
+        errs() << "\n";
+      }
+      llvm_unreachable("Failed to resolve contiguous address");
     }
-    compiledIt = compiledAddressExprs.try_emplace(value, *compiledExpr).first;
-  }
-
-  auto resolvedAddress = compiledIt->second.evaluate(knowledge, lane);
-  if (failed(resolvedAddress)) {
-    errs() << "Failed to evaluate contiguous address for value: ";
-    value.print(errs());
-    errs() << " : " << value.getType();
-    errs() << "\n";
-    if (auto* definingOp = value.getDefiningOp()) {
-      errs() << "Defining op:\n";
-      definingOp->print(errs());
-      errs() << "\n";
-    }
-    llvm_unreachable("Failed to resolve contiguous address");
   }
 
   MemoryValueKey key = getMemoryValueKey(resolvedAddress->base, lane);
@@ -1114,7 +1118,8 @@ enum class CompiledCoreOpKind : uint8_t {
 struct CompiledCoreNode {
   enum class Kind : uint8_t {
     Op,
-    Loop
+    Loop,
+    If
   };
 
   Kind kind = Kind::Op;
@@ -1123,7 +1128,10 @@ struct CompiledCoreNode {
   CompiledIndexExpr lowerBound;
   CompiledIndexExpr upperBound;
   CompiledIndexExpr step;
+  CompiledIndexExpr condition;
   std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> loopBody;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> thenBody;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> elseBody;
 };
 
 static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
@@ -1201,6 +1209,28 @@ compileCoreEmissionPlan(Block& block, Operation* weightOwner, llvm::SmallVectorI
       continue;
     }
 
+    if (auto ifOp = dyn_cast<mlir::scf::IfOp>(op)) {
+      auto condition = compileIndexExpr(ifOp.getCondition());
+      if (failed(condition)) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
+        return failure();
+      }
+
+      CompiledCoreNode ifNode;
+      ifNode.kind = CompiledCoreNode::Kind::If;
+      ifNode.op = ifOp.getOperation();
+      ifNode.condition = *condition;
+      ifNode.thenBody = std::make_unique<llvm::SmallVector<CompiledCoreNode, 8>>();
+      if (failed(compileCoreEmissionPlan(ifOp.getThenRegion().front(), weightOwner, *ifNode.thenBody)))
+        return failure();
+      ifNode.elseBody = std::make_unique<llvm::SmallVector<CompiledCoreNode, 8>>();
+      if (!ifOp.getElseRegion().empty())
+        if (failed(compileCoreEmissionPlan(ifOp.getElseRegion().front(), weightOwner, *ifNode.elseBody)))
+          return failure();
+      plan.push_back(std::move(ifNode));
+      continue;
+    }
+
     auto opKind = classifyCompiledCoreOpKind(op);
     if (failed(opKind)) {
       InFlightDiagnostic diag = op.emitError() << "unsupported codegen for op '" << op.getName().getStringRef() << "'";
@@ -1263,6 +1293,26 @@ static LogicalResult executeCompiledCorePlan(
       continue;
     }
 
+    if (node.kind == CompiledCoreNode::Kind::If) {
+      auto condition = node.condition.evaluate(knowledge);
+      auto ifOp = cast<mlir::scf::IfOp>(node.op);
+      if (failed(condition)) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
+        return failure();
+      }
+
+      const auto& selectedBody = *condition != 0 ? node.thenBody : node.elseBody;
+      if (selectedBody && failed(executeCompiledCorePlan(*selectedBody,
+                                                         coreCodeGen,
+                                                         knowledge,
+                                                         resolveWeightSlot,
+                                                         processedOperations,
+                                                         batchLane,
+                                                         batchLaneCount)))
+        return failure();
+      continue;
+    }
+
     switch (node.opKind) {
     case CompiledCoreOpKind::Load:
       coreCodeGen.codeGenLoadOp(cast<pim::PimMemCopyHostToDevOp>(node.op), knowledge);

src/PIM/Conversion/CMakeLists.txt

@@ -1,3 +1,2 @@
 add_subdirectory(ONNXToSpatial)
-add_subdirectory(SpatialToGraphviz)
+add_subdirectory(SpatialToPim)
-add_subdirectory(SpatialToPim)

src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt

@@ -31,8 +31,10 @@ add_pim_library(OMONNXToSpatial
   SpatialLayoutPlanningPass.cpp
   LowerSpatialPlansPass.cpp
   Common/AttributeUtils.cpp
+  Common/BiasAddUtils.cpp
   Common/ComputeRegionBuilder.cpp
   Common/MatrixProductLowering.cpp
+  Common/RowStripLayoutUtils.cpp
   Common/ShapeTilingUtils.cpp
   Common/WeightMaterialization.cpp
 

src/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.cpp

@@ -0,0 +1,112 @@
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+LogicalResult isSupportedBiasAddShape(RankedTensorType biasType, RankedTensorType resultType) {
+  if (!biasType || !resultType || !biasType.hasStaticShape() || !resultType.hasStaticShape())
+    return failure();
+  if (resultType.getRank() != 4)
+    return failure();
+  if (biasType.getElementType() != resultType.getElementType())
+    return failure();
+
+  const int64_t channels = resultType.getDimSize(1);
+  ArrayRef<int64_t> shape = biasType.getShape();
+  if (shape.empty())
+    return success();
+  if (shape.size() == 1)
+    return success(shape[0] == channels);
+  if (shape.size() == 2)
+    return success(shape[0] == 1 && shape[1] == channels);
+  if (shape.size() == 4)
+    return success(shape[0] == 1 && shape[1] == channels && shape[2] == 1 && shape[3] == 1);
+  return failure();
+}
+
+FailureOr<SmallVector<Attribute>> getBiasChannelValues(DenseElementsAttr denseAttr, RankedTensorType resultType) {
+  auto biasType = dyn_cast<RankedTensorType>(denseAttr.getType());
+  if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType)))
+    return failure();
+
+  const int64_t channels = resultType.getDimSize(1);
+  if (denseAttr.isSplat()) {
+    return SmallVector<Attribute>(channels, denseAttr.getSplatValue<Attribute>());
+  }
+
+  SmallVector<Attribute> flattened(denseAttr.getValues<Attribute>());
+  if (biasType.getRank() == 1)
+    return flattened;
+  if (biasType.getRank() == 2)
+    return flattened;
+
+  SmallVector<Attribute> channelValues;
+  channelValues.reserve(channels);
+  const int64_t channelStride = biasType.getDimSize(2) * biasType.getDimSize(3);
+  for (int64_t channel = 0; channel < channels; ++channel)
+    channelValues.push_back(flattened[channel * channelStride]);
+  return channelValues;
+}
+
+bool isSupportedBiasAddValue(Value bias, RankedTensorType resultType, DenseElementsAttr* denseAttr) {
+  auto attr = getHostConstDenseElementsAttr(bias);
+  if (!attr)
+    return false;
+  auto biasType = dyn_cast<RankedTensorType>(attr.getType());
+  if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType)))
+    return false;
+  if (failed(getBiasChannelValues(attr, resultType)))
+    return false;
+  if (denseAttr)
+    *denseAttr = attr;
+  return true;
+}
+
+FailureOr<BiasAddPlanCandidate> classifyBiasAddPlanCandidate(Value lhs, Value rhs, RankedTensorType resultType) {
+  auto lhsType = dyn_cast<RankedTensorType>(lhs.getType());
+  auto rhsType = dyn_cast<RankedTensorType>(rhs.getType());
+  if (!lhsType || !rhsType)
+    return failure();
+  if (lhsType == resultType && isSupportedBiasAddValue(rhs, resultType))
+    return BiasAddPlanCandidate {lhs, rhs};
+  if (rhsType == resultType && isSupportedBiasAddValue(lhs, resultType))
+    return BiasAddPlanCandidate {rhs, lhs};
+  return failure();
+}
+
+FailureOr<Value>
+materializeDenseBiasAddTensor(Value bias, RankedTensorType resultType, RewriterBase& rewriter, Location loc) {
+  DenseElementsAttr denseAttr;
+  if (!isSupportedBiasAddValue(bias, resultType, &denseAttr))
+    return failure();
+
+  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, resultType);
+  if (failed(channelValues))
+    return failure();
+
+  SmallVector<Attribute> resultValues;
+  resultValues.reserve(resultType.getNumElements());
+  const int64_t batches = resultType.getDimSize(0);
+  const int64_t channels = resultType.getDimSize(1);
+  const int64_t height = resultType.getDimSize(2);
+  const int64_t width = resultType.getDimSize(3);
+  for (int64_t n = 0; n < batches; ++n)
+    for (int64_t c = 0; c < channels; ++c)
+      for (int64_t h = 0; h < height; ++h)
+        for (int64_t w = 0; w < width; ++w)
+          resultValues.push_back((*channelValues)[c]);
+
+  auto resultAttr = DenseElementsAttr::get(resultType, resultValues);
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), resultAttr, resultType);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LogicalResult.h"
+
+namespace onnx_mlir {
+
+struct BiasAddPlanCandidate {
+  mlir::Value data;
+  mlir::Value bias;
+};
+
+mlir::LogicalResult isSupportedBiasAddShape(mlir::RankedTensorType biasType, mlir::RankedTensorType resultType);
+bool isSupportedBiasAddValue(mlir::Value bias,
+                             mlir::RankedTensorType resultType,
+                             mlir::DenseElementsAttr* denseAttr = nullptr);
+mlir::FailureOr<llvm::SmallVector<mlir::Attribute>>
+getBiasChannelValues(mlir::DenseElementsAttr denseAttr, mlir::RankedTensorType resultType);
+mlir::FailureOr<BiasAddPlanCandidate> classifyBiasAddPlanCandidate(mlir::Value lhs,
+                                                                   mlir::Value rhs,
+                                                                   mlir::RankedTensorType resultType);
+mlir::FailureOr<mlir::Value> materializeDenseBiasAddTensor(mlir::Value bias,
+                                                           mlir::RankedTensorType resultType,
+                                                           mlir::RewriterBase& rewriter,
+                                                           mlir::Location loc);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp

@@ -0,0 +1,239 @@
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
+  return RankedTensorType::get({logicalType.getDimSize(0), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
+                               logicalType.getElementType(),
+                               logicalType.getEncoding());
+}
+
+RankedTensorType getRowStripStorageType(RankedTensorType logicalType) {
+  return RankedTensorType::get({logicalType.getDimSize(2), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
+                               logicalType.getElementType(),
+                               logicalType.getEncoding());
+}
+
+std::pair<SmallVector<int64_t>, SmallVector<int64_t>> buildRowStripMetadata(RankedTensorType type) {
+  SmallVector<int64_t> offsets;
+  SmallVector<int64_t> sizes;
+  const int64_t channels = type.getDimSize(1);
+  const int64_t height = type.getDimSize(2);
+  const int64_t width = type.getDimSize(3);
+  offsets.reserve(height * 4);
+  sizes.reserve(height * 4);
+  for (int64_t row = 0; row < height; ++row) {
+    offsets.append({0, 0, row, 0});
+    sizes.append({1, channels, 1, width});
+  }
+  return {offsets, sizes};
+}
+
+SmallVector<OpFoldResult> buildRowStripFragmentOffsets(PatternRewriter& rewriter, OpFoldResult row) {
+  return {row, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+}
+
+SmallVector<OpFoldResult> buildRowStripFragmentSizes(PatternRewriter& rewriter, RankedTensorType logicalType) {
+  return {rewriter.getIndexAttr(1),
+          rewriter.getIndexAttr(logicalType.getDimSize(1)),
+          rewriter.getIndexAttr(1),
+          rewriter.getIndexAttr(logicalType.getDimSize(3))};
+}
+
+Value extractRowStripFragment(Value storage,
+                              RankedTensorType logicalType,
+                              OpFoldResult row,
+                              PatternRewriter& rewriter,
+                              Location loc) {
+  return tensor::ExtractSliceOp::create(rewriter,
+                                        loc,
+                                        getRowStripFragmentType(logicalType),
+                                        storage,
+                                        buildRowStripFragmentOffsets(rewriter, row),
+                                        buildRowStripFragmentSizes(rewriter, logicalType),
+                                        getUnitStrides(rewriter, 4));
+}
+
+void insertRowStripFragment(Value fragment,
+                            Value output,
+                            RankedTensorType logicalType,
+                            OpFoldResult row,
+                            PatternRewriter& rewriter,
+                            Location loc) {
+  createParallelInsertSliceIntoBatchOutput(rewriter,
+                                           loc,
+                                           fragment,
+                                           output,
+                                           buildRowStripFragmentOffsets(rewriter, row),
+                                           buildRowStripFragmentSizes(rewriter, logicalType),
+                                           getUnitStrides(rewriter, 4));
+}
+
+FailureOr<Value> createPerChannelConstantFragment(DenseElementsAttr denseAttr,
+                                                  RankedTensorType fragmentType,
+                                                  PatternRewriter& rewriter) {
+  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, fragmentType);
+  if (failed(channelValues))
+    return failure();
+
+  SmallVector<Attribute> values;
+  values.reserve(fragmentType.getNumElements());
+  for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n)
+    for (int64_t channel = 0; channel < fragmentType.getDimSize(1); ++channel)
+      for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h)
+        for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w)
+          values.push_back((*channelValues)[channel]);
+
+  auto attr = DenseElementsAttr::get(fragmentType, values);
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), attr, fragmentType);
+}
+
+FailureOr<Value> createRowStripStorageFromRows(Value rows,
+                                               RankedTensorType logicalType,
+                                               PatternRewriter& rewriter,
+                                               Location loc) {
+  auto rowsType = dyn_cast<RankedTensorType>(rows.getType());
+  if (!rowsType || !rowsType.hasStaticShape() || rowsType.getRank() != 2)
+    return failure();
+  if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() != 4)
+    return failure();
+  if (logicalType.getDimSize(0) != 1)
+    return failure();
+  if (rowsType.getElementType() != logicalType.getElementType())
+    return failure();
+
+  const int64_t channels = logicalType.getDimSize(1);
+  const int64_t height = logicalType.getDimSize(2);
+  const int64_t width = logicalType.getDimSize(3);
+  if (rowsType.getDimSize(0) != height * width)
+    return failure();
+  if (rowsType.getDimSize(1) != channels)
+    return failure();
+
+  auto rowSliceType = RankedTensorType::get({width, channels}, logicalType.getElementType(), rowsType.getEncoding());
+  auto channelWidthType = RankedTensorType::get({channels, width}, logicalType.getElementType(), rowsType.getEncoding());
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
+  auto batchOp = createSpatComputeBatch(
+    rewriter, loc, TypeRange {storageType}, height, {}, ValueRange {rows}, [&](detail::SpatComputeBatchBodyArgs args) {
+      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+      Value rowStart = affineMulConst(rewriter, loc, args.lane, width, anchorOp);
+      SmallVector<OpFoldResult> rowOffsets {rowStart, rewriter.getIndexAttr(0)};
+      SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(width), rewriter.getIndexAttr(channels)};
+      Value rowSlice = tensor::ExtractSliceOp::create(
+        rewriter, loc, rowSliceType, args.inputs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
+      Value channelWidth = ONNXTransposeOp::create(
+        rewriter, loc, channelWidthType, rowSlice, rewriter.getI64ArrayAttr({1, 0})).getResult();
+      Value fragment = tensor::ExpandShapeOp::create(
+        rewriter, loc, fragmentType, channelWidth, SmallVector<ReassociationIndices> {{0, 1}, {2, 3}});
+      insertRowStripFragment(fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+      return success();
+    });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+FailureOr<Value>
+materializeRowStripStorageToDense(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
+  auto storageType = dyn_cast<RankedTensorType>(storage.getType());
+  if (!storageType || storageType != getRowStripStorageType(logicalType))
+    return failure();
+
+  auto batchOp = createSpatComputeBatch(
+    rewriter, loc, TypeRange {logicalType}, logicalType.getDimSize(2), {}, ValueRange {storage},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      Value fragment = extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
+      createParallelInsertSliceIntoBatchOutput(rewriter,
+                                               loc,
+                                               fragment,
+                                               args.outputs.front(),
+                                               SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
+                                                                          rewriter.getIndexAttr(0),
+                                                                          args.lane,
+                                                                          rewriter.getIndexAttr(0)},
+                                               buildRowStripFragmentSizes(rewriter, logicalType),
+                                               getUnitStrides(rewriter, 4));
+      return success();
+    });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+FailureOr<Value>
+applyRowStripRelu(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
+  auto batchOp = createSpatComputeBatch(rewriter,
+                                        loc,
+                                        TypeRange {storageType},
+                                        logicalType.getDimSize(2),
+                                        {},
+                                        ValueRange {storage},
+                                        [&](detail::SpatComputeBatchBodyArgs args) {
+                                          Value fragment =
+                                            extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
+                                          fragment = spatial::SpatReluOp::create(rewriter, loc, fragmentType, fragment).getResult();
+                                          insertRowStripFragment(
+                                            fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+                                          return success();
+                                        });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+FailureOr<Value>
+applyRowStripBiasAdd(Value storage, RankedTensorType logicalType, Value bias, PatternRewriter& rewriter, Location loc) {
+  DenseElementsAttr denseAttr;
+  if (!isSupportedBiasAddValue(bias, logicalType, &denseAttr))
+    return failure();
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
+  auto batchOp = createSpatComputeBatch(rewriter,
+                                        loc,
+                                        TypeRange {storageType},
+                                        logicalType.getDimSize(2),
+                                        {},
+                                        ValueRange {storage},
+                                        [&](detail::SpatComputeBatchBodyArgs args) {
+                                          Value fragment =
+                                            extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
+                                          Value constant;
+                                          if (denseAttr.isSplat()) {
+                                            constant = getOrCreateConstant(
+                                              rewriter,
+                                              rewriter.getInsertionBlock()->getParentOp(),
+                                              DenseElementsAttr::get(fragmentType, denseAttr.getSplatValue<Attribute>()),
+                                              fragmentType);
+                                          }
+                                          else {
+                                            FailureOr<Value> perChannel =
+                                              createPerChannelConstantFragment(denseAttr, fragmentType, rewriter);
+                                            if (failed(perChannel))
+                                              return failure();
+                                            constant = *perChannel;
+                                          }
+                                          fragment =
+                                            spatial::SpatVAddOp::create(rewriter, loc, fragmentType, fragment, constant).getResult();
+                                          insertRowStripFragment(
+                                            fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+                                          return success();
+                                        });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp

@@ -0,0 +1,69 @@
+#pragma once
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+
+namespace onnx_mlir {
+
+inline constexpr llvm::StringLiteral kRowStripIndexMap = "nchw_row_strip_fragments";
+
+struct RowStripPhysicalValue {
+  mlir::Value storage;
+  mlir::RankedTensorType logicalType;
+  llvm::SmallVector<int64_t, 16> fragmentOffsets;
+  llvm::SmallVector<int64_t, 16> fragmentSizes;
+};
+
+std::pair<llvm::SmallVector<int64_t>, llvm::SmallVector<int64_t>>
+buildRowStripMetadata(mlir::RankedTensorType type);
+
+mlir::RankedTensorType getRowStripFragmentType(mlir::RankedTensorType logicalType);
+
+mlir::RankedTensorType getRowStripStorageType(mlir::RankedTensorType logicalType);
+
+llvm::SmallVector<mlir::OpFoldResult> buildRowStripFragmentOffsets(mlir::PatternRewriter& rewriter,
+                                                                   mlir::OpFoldResult row);
+
+llvm::SmallVector<mlir::OpFoldResult> buildRowStripFragmentSizes(mlir::PatternRewriter& rewriter,
+                                                                 mlir::RankedTensorType logicalType);
+
+mlir::Value extractRowStripFragment(mlir::Value storage,
+                                    mlir::RankedTensorType logicalType,
+                                    mlir::OpFoldResult row,
+                                    mlir::PatternRewriter& rewriter,
+                                    mlir::Location loc);
+
+void insertRowStripFragment(mlir::Value fragment,
+                            mlir::Value output,
+                            mlir::RankedTensorType logicalType,
+                            mlir::OpFoldResult row,
+                            mlir::PatternRewriter& rewriter,
+                            mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> createPerChannelConstantFragment(mlir::DenseElementsAttr denseAttr,
+                                                              mlir::RankedTensorType fragmentType,
+                                                              mlir::PatternRewriter& rewriter);
+
+mlir::FailureOr<mlir::Value> createRowStripStorageFromRows(mlir::Value rows,
+                                                           mlir::RankedTensorType logicalType,
+                                                           mlir::PatternRewriter& rewriter,
+                                                           mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> materializeRowStripStorageToDense(mlir::Value storage,
+                                                               mlir::RankedTensorType logicalType,
+                                                               mlir::PatternRewriter& rewriter,
+                                                               mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> applyRowStripRelu(mlir::Value storage,
+                                               mlir::RankedTensorType logicalType,
+                                               mlir::PatternRewriter& rewriter,
+                                               mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(mlir::Value storage,
+                                                  mlir::RankedTensorType logicalType,
+                                                  mlir::Value bias,
+                                                  mlir::PatternRewriter& rewriter,
+                                                  mlir::Location loc);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -13,7 +13,9 @@
 #include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
@@ -29,14 +31,6 @@ namespace {
 static constexpr StringLiteral kDenseLayout = "dense_nchw";
 static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";
 
-struct RowStripPhysicalValue {
-  Value physicalValue;
-  RankedTensorType logicalType;
-  SmallVector<int64_t, 16> fragmentOffsets;
-  SmallVector<int64_t, 16> fragmentSizes;
-  std::string indexMap;
-};
-
 static FailureOr<RowStripPhysicalValue> getRowStripValue(llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
                                                          Value value) {
   auto it = rowStripValues.find(value);
@@ -46,112 +40,42 @@ static FailureOr<RowStripPhysicalValue> getRowStripValue(llvm::DenseMap<Value, R
 }
 
 static FailureOr<RowStripPhysicalValue> buildRowStripValue(spatial::SpatBlueprintOp blueprint,
-                                                           Value physicalValue) {
+                                                           Value storage) {
   auto logicalType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
   if (!logicalType)
     return blueprint.emitOpError("requires ranked logical output type"), failure();
   RowStripPhysicalValue value;
-  value.physicalValue = physicalValue;
+  value.storage = storage;
   value.logicalType = logicalType;
   value.fragmentOffsets.append(blueprint.getFragmentOffsets().begin(), blueprint.getFragmentOffsets().end());
   value.fragmentSizes.append(blueprint.getFragmentSizes().begin(), blueprint.getFragmentSizes().end());
-  value.indexMap = blueprint.getIndexMap().str();
+  if (blueprint.getIndexMap() != kRowStripIndexMap)
+    return blueprint.emitOpError("requires the canonical row-strip index map"), failure();
+  auto storageType = dyn_cast<RankedTensorType>(storage.getType());
+  if (!storageType || storageType != getRowStripStorageType(logicalType))
+    return blueprint.emitOpError("requires physical row-strip fragment storage"), failure();
   return value;
 }
 
 static FailureOr<Value>
 lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) {
-  auto packedType = cast<RankedTensorType>(input.physicalValue.getType());
+  return applyRowStripRelu(input.storage, input.logicalType, rewriter, planOp.getLoc());
-  auto computeOp =
+}
-    createSpatCompute<1>(rewriter, planOp.getLoc(), TypeRange {packedType}, {}, input.physicalValue, [&](Value x) {
+
-      auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), packedType, x);
+static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
-      spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult());
+                                             spatial::SpatBiasAddPlanOp planOp,
-    });
+                                             PatternRewriter& rewriter) {
-  return computeOp.getResult(0);
+  return applyRowStripBiasAdd(input.storage, input.logicalType, planOp.getBias(), rewriter, planOp.getLoc());
 }
 
 static FailureOr<Value>
 materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
-  auto packedType = dyn_cast<RankedTensorType>(rowStripValue.physicalValue.getType());
-  if (!packedType || packedType.getRank() != 3 || !packedType.hasStaticShape())
-    return failure();
   if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
     return failure();
-  if (rowStripValue.indexMap != "packed_hwc_rows_to_nchw")
+  auto [expectedOffsets, expectedSizes] = buildRowStripMetadata(rowStripValue.logicalType);
+  if (!llvm::equal(rowStripValue.fragmentOffsets, expectedOffsets) || !llvm::equal(rowStripValue.fragmentSizes, expectedSizes))
     return failure();
-
+  return materializeRowStripStorageToDense(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
-  const int64_t rank = rowStripValue.logicalType.getRank();
-  const int64_t fragmentCount = rowStripValue.fragmentOffsets.size() / rank;
-  const int64_t packedWidth = packedType.getDimSize(1);
-  const int64_t packedChannels = packedType.getDimSize(2);
-  if (fragmentCount != packedType.getDimSize(0))
-    return failure();
-  for (int64_t fragmentIndex = 0; fragmentIndex < fragmentCount; ++fragmentIndex) {
-    if (rowStripValue.fragmentOffsets[fragmentIndex * rank + 0] != 0
-        || rowStripValue.fragmentOffsets[fragmentIndex * rank + 1] != 0
-        || rowStripValue.fragmentOffsets[fragmentIndex * rank + 2] != fragmentIndex
-        || rowStripValue.fragmentOffsets[fragmentIndex * rank + 3] != 0)
-      return failure();
-    if (rowStripValue.fragmentSizes[fragmentIndex * rank + 0] != 1
-        || rowStripValue.fragmentSizes[fragmentIndex * rank + 1] != packedChannels
-        || rowStripValue.fragmentSizes[fragmentIndex * rank + 2] != 1
-        || rowStripValue.fragmentSizes[fragmentIndex * rank + 3] != packedWidth)
-      return failure();
-  }
-
-  auto packedSliceType =
-    RankedTensorType::get({1, packedWidth, packedChannels}, packedType.getElementType(), packedType.getEncoding());
-  auto expandedType =
-    RankedTensorType::get({1, 1, packedWidth, packedChannels}, packedType.getElementType(), packedType.getEncoding());
-  auto logicalFragmentType =
-    RankedTensorType::get({1, packedChannels, 1, packedWidth}, packedType.getElementType(), packedType.getEncoding());
-  auto batchOp = createSpatComputeBatch(
-    rewriter,
-    loc,
-    TypeRange {rowStripValue.logicalType},
-    fragmentCount,
-    {},
-    ValueRange {rowStripValue.physicalValue},
-    [&](detail::SpatComputeBatchBodyArgs args) {
-      SmallVector<OpFoldResult> packedOffsets {args.lane, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> packedSizes {
-        rewriter.getIndexAttr(1), rewriter.getIndexAttr(packedWidth), rewriter.getIndexAttr(packedChannels)};
-      Value packedSlice = tensor::ExtractSliceOp::create(
-        rewriter, loc, packedSliceType, args.inputs.front(), packedOffsets, packedSizes, getUnitStrides(rewriter, 3));
-
-      Value expanded = tensor::ExpandShapeOp::create(rewriter,
-                                                     loc,
-                                                     expandedType,
-                                                     packedSlice,
-                                                     SmallVector<ReassociationIndices> {
-                                                       {0, 1},
-                                                       {2},
-                                                       {3}
-      });
-      Value transposeInit =
-        tensor::EmptyOp::create(rewriter, loc, logicalFragmentType.getShape(), logicalFragmentType.getElementType());
-      Value logicalFragment =
-        linalg::TransposeOp::create(rewriter, loc, expanded, transposeInit, SmallVector<int64_t> {0, 3, 1, 2})
-          .getResult()[0];
-
-      SmallVector<OpFoldResult> logicalOffsets {
-        rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> logicalSizes {rewriter.getIndexAttr(1),
-                                              rewriter.getIndexAttr(packedChannels),
-                                              rewriter.getIndexAttr(1),
-                                              rewriter.getIndexAttr(packedWidth)};
-      createParallelInsertSliceIntoBatchOutput(rewriter,
-                                               loc,
-                                               logicalFragment,
-                                               args.outputs.front(),
-                                               logicalOffsets,
-                                               logicalSizes,
-                                               getUnitStrides(rewriter, 4));
-      return success();
-    });
-  if (failed(batchOp))
-    return failure();
-  return batchOp->getResult(0);
 }
 
 struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
@@ -194,7 +118,7 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
           rewriter.setInsertionPoint(planOp);
           FailureOr<Value> lowered = lowerSelectedConv2DPlan(
             planOp,
-            succeeded(rowStripInput) ? std::optional<Value> {rowStripInput->physicalValue} : std::nullopt,
+            succeeded(rowStripInput) ? std::optional<Value> {rowStripInput->storage} : std::nullopt,
             /*emitRowStripLayout=*/true,
             rewriter);
           if (failed(lowered)) {
@@ -266,6 +190,64 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         rewriter.replaceOp(planOp, computeOp.getResults());
         continue;
       }
+      if (auto planOp = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
+        if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
+          auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
+            auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
+            return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
+          });
+          if (outputBlueprint == planOp.getResult().getUsers().end()) {
+            planOp.emitOpError("row-strip bias_add plan requires a row-strip blueprint result");
+            signalPassFailure();
+            return;
+          }
+
+          FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
+          rewriter.setInsertionPoint(planOp);
+          FailureOr<Value> lowered = lowerRowStripBiasAdd(*input, planOp, rewriter);
+          if (failed(lowered)) {
+            planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan");
+            signalPassFailure();
+            return;
+          }
+          auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
+          FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
+          if (failed(output)) {
+            signalPassFailure();
+            return;
+          }
+          rowStripValues[blueprint.getResult()] = *output;
+          eraseAfterLowering.insert(planOp);
+          eraseAfterLowering.insert(blueprint);
+          continue;
+        }
+
+        auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+        if (!resultType) {
+          planOp.emitOpError("requires ranked output type");
+          signalPassFailure();
+          return;
+        }
+        rewriter.setInsertionPoint(planOp);
+        FailureOr<Value> denseBias = materializeDenseBiasAddTensor(planOp.getBias(), resultType, rewriter, planOp.getLoc());
+        if (failed(denseBias)) {
+          planOp.emitOpError("failed to materialize dense Conv-style bias");
+          signalPassFailure();
+          return;
+        }
+        auto computeOp = createSpatCompute<2>(rewriter,
+                                              planOp.getLoc(),
+                                              planOp.getOutput().getType(),
+                                              {},
+                                              ValueRange {planOp.getInput(), *denseBias},
+                                              [&](Value x, Value y) {
+                                                auto added = spatial::SpatVAddOp::create(
+                                                  rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y);
+                                                spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult());
+                                              });
+        rewriter.replaceOp(planOp, computeOp.getResults());
+        continue;
+      }
       if (auto materializeOp = dyn_cast<spatial::SpatMaterializeLayoutOp>(&op)) {
         if (materializeOp.getSourcePhysicalLayout() == kDenseLayout
             && materializeOp.getTargetPhysicalLayout() == kDenseLayout) {
@@ -385,6 +367,7 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
       if (isa<ONNXEntryPointOp>(op))
         return;
       if (isa<spatial::SpatConv2DPlanOp,
+              spatial::SpatBiasAddPlanOp,
               spatial::SpatReluPlanOp,
               spatial::SpatBlueprintOp,
               spatial::SpatMaterializeLayoutOp>(op)

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -45,11 +45,12 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
   SmallVector<spatial::SpatGraphCompute> computes(funcOp.getOps<spatial::SpatGraphCompute>());
   SmallVector<spatial::SpatGraphComputeBatch> computeBatches(funcOp.getOps<spatial::SpatGraphComputeBatch>());
   SmallVector<spatial::SpatConv2DPlanOp> convPlans(funcOp.getOps<spatial::SpatConv2DPlanOp>());
+  SmallVector<spatial::SpatBiasAddPlanOp> biasAddPlans(funcOp.getOps<spatial::SpatBiasAddPlanOp>());
   SmallVector<spatial::SpatReluPlanOp> reluPlans(funcOp.getOps<spatial::SpatReluPlanOp>());
   SmallVector<spatial::SpatBlueprintOp> blueprints(funcOp.getOps<spatial::SpatBlueprintOp>());
   SmallVector<spatial::SpatMaterializeLayoutOp> materializers(funcOp.getOps<spatial::SpatMaterializeLayoutOp>());
-  if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !reluPlans.empty() || !blueprints.empty()
+  if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !biasAddPlans.empty() || !reluPlans.empty()
-      || !materializers.empty()) {
+      || !blueprints.empty() || !materializers.empty()) {
     return;
   }
 

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp

@@ -132,6 +132,7 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
             spatial::SpatGraphCompute,
             spatial::SpatGraphComputeBatch,
             spatial::SpatConv2DPlanOp,
+            spatial::SpatBiasAddPlanOp,
             spatial::SpatReluPlanOp,
             spatial::SpatBlueprintOp,
             spatial::SpatMaterializeLayoutOp>(&op)) {

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp

@@ -23,7 +23,9 @@
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.hpp"
@@ -221,10 +223,6 @@ struct DistributedConvReportTotals {
 };
 
 static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter& rewriter);
-static FailureOr<Value> createRowStripPackedRows(Value rows,
-                                                 const ConvLoweringState& state,
-                                                 PatternRewriter& rewriter,
-                                                 Location loc);
 
 static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp, Value x, Value w, Value b);
 
@@ -467,59 +465,6 @@ classifyDistributedBinaryConsumer(Operation* user,
   return std::nullopt;
 }
 
-static bool covers(RowInterval acquired, RowInterval needed) {
-  return acquired.begin <= needed.begin && acquired.end >= needed.end;
-}
-
-static bool canConsumeRowStripHwcInput(const ConvLoweringState& state, StringRef& failureReason) {
-  if (state.batchSize != 1) {
-    failureReason = "unsupported_batch";
-    return false;
-  }
-  if (state.group != 1) {
-    failureReason = "unsupported_groups";
-    return false;
-  }
-  if (isDepthwiseConv(state.group, state.numChannelsIn, state.numChannelsOut, state.numChannelsInPerGroup)) {
-    failureReason = "unsupported_depthwise";
-    return false;
-  }
-  if (state.strideHeight != 1 || state.strideWidth != 1) {
-    failureReason = "unsupported_stride";
-    return false;
-  }
-  if (state.dilationHeight != 1 || state.dilationWidth != 1) {
-    failureReason = "unsupported_dilation";
-    return false;
-  }
-  if (state.padHeightBegin != state.padHeightEnd || state.padWidthBegin != state.padWidthEnd) {
-    failureReason = "unsupported_padding";
-    return false;
-  }
-  if (state.padHeightBegin != 1 || state.padWidthBegin != 1) {
-    failureReason = "unsupported_padding";
-    return false;
-  }
-  if (state.wHeight != 3 || state.wWidth != 3) {
-    failureReason = "unsupported_kernel";
-    return false;
-  }
-  if (state.outHeight != state.xHeight || state.outWidth != state.xWidth) {
-    failureReason = "unsupported_output_shape";
-    return false;
-  }
-  if (!getHostConstDenseElementsAttr(state.w)) {
-    failureReason = "non_constant_weights";
-    return false;
-  }
-  if (state.hasBias && !getHostConstDenseElementsAttr(state.b)) {
-    failureReason = "non_constant_bias";
-    return false;
-  }
-  failureReason = "";
-  return true;
-}
-
 static std::string stringifyDistributedTensorOpKind(DistributedTensorOpKind kind) {
   switch (kind) {
   case DistributedTensorOpKind::Relu: return "Relu";
@@ -2493,21 +2438,15 @@ static Value rewriteStreamedConv(const ConvLoweringState& state,
 
 } // namespace standard
 
-static RankedTensorType getRowStripFragmentType(RankedTensorType tensorType, int64_t width) {
-  return RankedTensorType::get(
-    {tensorType.getDimSize(0), tensorType.getDimSize(1), 1, width}, tensorType.getElementType(), tensorType.getEncoding());
-}
-
 static SmallVector<DistributedFragmentInfo, 8> buildRowStripFragments(RankedTensorType tensorType) {
   SmallVector<DistributedFragmentInfo, 8> fragments;
-  const int64_t height = tensorType.getDimSize(2);
+  auto [offsets, sizes] = buildRowStripMetadata(tensorType);
-  const int64_t width = tensorType.getDimSize(3);
+  const int64_t rank = tensorType.getRank();
-  const int64_t channels = tensorType.getDimSize(1);
+  fragments.reserve(offsets.size() / rank);
-  fragments.reserve(height);
+  for (int64_t row = 0; row < static_cast<int64_t>(offsets.size() / rank); ++row) {
-  for (int64_t row = 0; row < height; ++row) {
     fragments.push_back(DistributedFragmentInfo {
-      {0, 0, row, 0},
+      {offsets.begin() + row * rank, offsets.begin() + (row + 1) * rank},
-      {1, channels, 1, width},
+      {sizes.begin() + row * rank, sizes.begin() + (row + 1) * rank},
       {1, 1, 1, 1},
       row,
     });
@@ -2527,101 +2466,266 @@ static DistributedTensorInfo makeDistributedTensorInfo(Value storage, RankedTens
   return info;
 }
 
-static Value createPerChannelConstantFragment(DenseElementsAttr denseAttr,
-                                              RankedTensorType fragmentType,
-                                              PatternRewriter& rewriter) {
-  auto denseType = cast<RankedTensorType>(denseAttr.getType());
-  SmallVector<Attribute> channelValues;
-  channelValues.reserve(fragmentType.getDimSize(1));
-  SmallVector<Attribute> flattened(denseAttr.getValues<Attribute>());
-  if (denseType.getRank() == 1) {
-    channelValues = flattened;
-  }
-  else if (denseType.getRank() == 2) {
-    channelValues = flattened;
-  }
-  else {
-    for (int64_t channel = 0; channel < denseType.getDimSize(1); ++channel)
-      channelValues.push_back(flattened[channel]);
-  }
-
-  SmallVector<Attribute> values;
-  values.reserve(fragmentType.getNumElements());
-  for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n)
-    for (int64_t channel = 0; channel < fragmentType.getDimSize(1); ++channel)
-      for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h)
-        for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w)
-          values.push_back(channelValues[channel]);
-
-  auto attr = DenseElementsAttr::get(fragmentType, values);
-  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), attr, fragmentType);
-}
-
 static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter& rewriter) {
   auto zeroAttr = DenseElementsAttr::get(gemmResultType, rewriter.getZeroAttr(gemmResultType.getElementType()));
   return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, gemmResultType);
 }
 
-static bool canDirectLowerRowStripConv(const ConvLoweringState& state, StringRef& failureReason) {
+static bool canConsumeNchwRowStripFragments(const ConvLoweringState& state, StringRef& failureReason) {
-  if (!canConsumeRowStripHwcInput(state, failureReason))
+  if (state.batchSize != 1) {
-    return false;
+    failureReason = "batch_not_one";
-
+    return false;
-  ConvGeometry geometry = buildConvGeometry(state);
+  }
-  if (state.numChannelsOut > geometry.xbarSize) {
+  if (state.group != 1) {
-    failureReason = "unsupported_output_channels";
+    failureReason = "grouped_conv";
+    return false;
+  }
+  if (!state.xType.hasStaticShape() || !state.wType.hasStaticShape() || !state.outType.hasStaticShape()) {
+    failureReason = "dynamic_shape";
+    return false;
+  }
+  if (!isa<FloatType>(state.xType.getElementType())) {
+    failureReason = "non_float_input";
+    return false;
+  }
+  if (state.strideHeight != 1 || state.strideWidth != 1) {
+    failureReason = "stride_not_one";
+    return false;
+  }
+  if (state.dilationHeight != 1 || state.dilationWidth != 1) {
+    failureReason = "dilation_not_one";
+    return false;
+  }
+  if (state.wHeight != 3 || state.wWidth != 3) {
+    failureReason = "kernel_not_3x3";
+    return false;
+  }
+  if (state.padHeightBegin != 1 || state.padHeightEnd != 1 || state.padWidthBegin != 1 || state.padWidthEnd != 1) {
+    failureReason = "padding_not_1";
+    return false;
+  }
+  if (state.outHeight != state.xHeight || state.outWidth != state.xWidth) {
+    failureReason = "not_same_spatial_shape";
+    return false;
+  }
+  if (!getHostConstDenseElementsAttr(state.w)) {
+    failureReason = "non_constant_weight";
+    return false;
+  }
+  if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType)) {
+    failureReason = "unsupported_bias";
+    return false;
+  }
+  ConvGeometry geometry = buildConvGeometry(state);
+  if (geometry.c > geometry.xbarSize) {
+    failureReason = "output_channels_exceed_crossbar";
     return false;
   }
-
-  failureReason = "";
   return true;
 }
 
-static FailureOr<Value> createRowStripPackedRows(Value rows,
+static Value createZeroTensorConstant(RankedTensorType type, PatternRewriter& rewriter) {
-                                                 const ConvLoweringState& state,
+  auto zeroAttr = DenseElementsAttr::get(type, rewriter.getZeroAttr(type.getElementType()));
-                                                 PatternRewriter& rewriter,
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, type);
-                                                 Location loc) {
-  auto rowsType = dyn_cast<RankedTensorType>(rows.getType());
-  if (!rowsType || !rowsType.hasStaticShape() || rowsType.getRank() != 2)
-    return failure();
-
-  if (state.batchSize != 1)
-    return failure();
-  if (state.outType.getRank() != 4 || !state.outType.hasStaticShape())
-    return failure();
-
-  const int64_t outHeight = state.outType.getDimSize(2);
-  const int64_t outWidth = state.outType.getDimSize(3);
-  const int64_t outChannels = state.outType.getDimSize(1);
-  if (rowsType.getDimSize(0) != outHeight * outWidth || rowsType.getDimSize(1) != outChannels)
-    return failure();
-
-  auto packedType = RankedTensorType::get({outHeight, outWidth, outChannels}, rowsType.getElementType(), rowsType.getEncoding());
-  auto packedRows =
-    createSpatCompute<1>(rewriter, loc, TypeRange {packedType}, {}, rows, [&](Value rowValues) {
-      Value packed = tensor::ExpandShapeOp::create(
-        rewriter, loc, packedType, rowValues, SmallVector<ReassociationIndices> {{0, 1}, {2}});
-      spatial::SpatYieldOp::create(rewriter, loc, packed);
-    });
-  return packedRows.getResult(0);
 }
 
-static FailureOr<Value> createConvOutputFromRowStripHwc(Value inputHwc,
+static FailureOr<Value> createPaddedBiasRowConstant(const ConvLoweringState& state,
-                                                        const ConvLoweringState& state,
+                                                    int64_t paddedChannels,
-                                                        PatternRewriter& rewriter,
+                                                    PatternRewriter& rewriter) {
-                                                        Location loc) {
+  DenseElementsAttr denseAttr;
-  auto inputType = dyn_cast<RankedTensorType>(inputHwc.getType());
+  if (!isSupportedBiasAddValue(state.b, state.outType, &denseAttr))
-  if (!inputType || !inputType.hasStaticShape() || inputType.getRank() != 3)
     return failure();
-  if (inputType.getDimSize(0) != state.xHeight || inputType.getDimSize(1) != state.xWidth
+  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, state.outType);
-      || inputType.getDimSize(2) != state.numChannelsIn)
+  if (failed(channelValues))
+    return failure();
+
+  auto biasType = RankedTensorType::get({1, paddedChannels}, state.outType.getElementType());
+  SmallVector<Attribute> values(biasType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(biasType.getElementType())));
+  for (int64_t channel = 0; channel < state.numChannelsOut; ++channel)
+    values[channel] = (*channelValues)[channel];
+  auto biasAttr = DenseElementsAttr::get(biasType, values);
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), biasAttr, biasType);
+}
+
+static Value createHorizontallyPaddedRowStripFragment(Value fragment,
+                                                      const ConvLoweringState& state,
+                                                      PatternRewriter& rewriter,
+                                                      Location loc) {
+  auto paddedType = RankedTensorType::get({1, state.numChannelsIn, 1, state.xWidth + 2},
+                                          state.xType.getElementType(),
+                                          state.xType.getEncoding());
+  return createZeroPaddedTensor(fragment, paddedType, {0, 0, 0, 1}, {0, 0, 0, 1}, rewriter, loc);
+}
+
+static Value createRowStripWindowSourceRowTable(const ConvLoweringState& state, PatternRewriter& rewriter) {
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  auto tableType = RankedTensorType::get({state.outHeight * state.wHeight}, rewriter.getIndexType());
+  SmallVector<Attribute> values;
+  values.reserve(tableType.getNumElements());
+  for (int64_t outputRow = 0; outputRow < state.outHeight; ++outputRow) {
+    for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) {
+      int64_t sourceRow = outputRow + kernelRow - state.padHeightBegin;
+      sourceRow = std::clamp(sourceRow, int64_t {0}, state.xHeight - 1);
+      values.push_back(rewriter.getIndexAttr(sourceRow));
+    }
+  }
+
+  return getOrCreateConstant(rewriter, anchorOp, DenseElementsAttr::get(tableType, values), tableType);
+}
+
+static Value createRowStripWindowTableIndex(Value outputHeight,
+                                            Value kernelRow,
+                                            const ConvLoweringState& state,
+                                            PatternRewriter& rewriter,
+                                            Location loc) {
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  MLIRContext* ctx = rewriter.getContext();
+  AffineExpr outputRowExpr = getAffineDimExpr(0, ctx);
+  AffineExpr kernelRowExpr = getAffineDimExpr(1, ctx);
+  return createOrFoldAffineApply(
+    rewriter, loc, outputRowExpr * state.wHeight + kernelRowExpr, ValueRange {outputHeight, kernelRow}, anchorOp);
+}
+
+static Value extractProjectedRowStripWindowRow(Value rowStripStorage,
+                                               Value sourceRowTable,
+                                               const ConvLoweringState& state,
+                                               Value outputHeight,
+                                               Value kernelRow,
+                                               PatternRewriter& rewriter,
+                                               Location loc) {
+  Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc);
+  Value sourceRow = tensor::ExtractOp::create(rewriter, loc, sourceRowTable, ValueRange {tableIndex}).getResult();
+  return extractRowStripFragment(rowStripStorage, state.xType, sourceRow, rewriter, loc);
+}
+
+static FailureOr<Value> createRowStripWindowMaskTable(const ConvLoweringState& state, PatternRewriter& rewriter) {
+  auto elementType = state.xType.getElementType();
+  auto floatType = dyn_cast<FloatType>(elementType);
+  if (!floatType)
+    return failure();
+
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  auto tableType = RankedTensorType::get({state.outHeight * state.wHeight, state.numChannelsIn, 1, state.xWidth},
+                                         elementType,
+                                         state.xType.getEncoding());
+  Attribute zero = rewriter.getZeroAttr(elementType);
+  Attribute one = rewriter.getFloatAttr(floatType, 1.0);
+  SmallVector<Attribute> values;
+  values.reserve(tableType.getNumElements());
+  for (int64_t outputRow = 0; outputRow < state.outHeight; ++outputRow) {
+    for (int64_t kernelRow = 0; kernelRow < state.wHeight; ++kernelRow) {
+      int64_t sourceRow = outputRow + kernelRow - state.padHeightBegin;
+      Attribute value = (sourceRow < 0 || sourceRow >= state.xHeight) ? zero : one;
+      for (int64_t channel = 0; channel < state.numChannelsIn; ++channel)
+        for (int64_t width = 0; width < state.xWidth; ++width)
+          values.push_back(value);
+    }
+  }
+
+  return getOrCreateConstant(rewriter, anchorOp, DenseElementsAttr::get(tableType, values), tableType);
+}
+
+static Value extractProjectedRowStripWindowMask(Value maskTable,
+                                                const ConvLoweringState& state,
+                                                Value outputHeight,
+                                                Value kernelRow,
+                                                PatternRewriter& rewriter,
+                                                Location loc) {
+  Value tableIndex = createRowStripWindowTableIndex(outputHeight, kernelRow, state, rewriter, loc);
+  auto fragmentType = getRowStripFragmentType(state.xType);
+  SmallVector<OpFoldResult> offsets {
+    tableIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
+                                   rewriter.getIndexAttr(state.numChannelsIn),
+                                   rewriter.getIndexAttr(1),
+                                   rewriter.getIndexAttr(state.xWidth)};
+  return tensor::ExtractSliceOp::create(rewriter,
+                                        loc,
+                                        fragmentType,
+                                        maskTable,
+                                        offsets,
+                                        sizes,
+                                        getUnitStrides(rewriter, 4));
+}
+
+static FailureOr<Value> createNchwRowStripConvWindow(Value rowStripStorage,
+                                                     const ConvLoweringState& state,
+                                                     Value outputHeight,
+                                                     PatternRewriter& rewriter,
+                                                     Location loc) {
+  auto fragmentType = getRowStripFragmentType(state.xType);
+  auto paddedWindowType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.xWidth + 2},
+                                                state.xType.getElementType(),
+                                                state.xType.getEncoding());
+  Value sourceRowTable = createRowStripWindowSourceRowTable(state, rewriter);
+  FailureOr<Value> maskTable = createRowStripWindowMaskTable(state, rewriter);
+  if (failed(maskTable))
+    return failure();
+  Value initWindow = createZeroTensorConstant(paddedWindowType, rewriter);
+
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  Value window = initWindow;
+  for (int64_t kernelRowIndex = 0; kernelRowIndex < state.wHeight; ++kernelRowIndex) {
+    Value kernelRow = getOrCreateIndexConstant(rewriter, anchorOp, kernelRowIndex);
+    Value sourceRow =
+      extractProjectedRowStripWindowRow(rowStripStorage, sourceRowTable, state, outputHeight, kernelRow, rewriter, loc);
+    Value mask = extractProjectedRowStripWindowMask(*maskTable, state, outputHeight, kernelRow, rewriter, loc);
+    Value semanticRow = spatial::SpatVMulOp::create(rewriter, loc, fragmentType, sourceRow, mask).getResult();
+    Value paddedRow = createHorizontallyPaddedRowStripFragment(semanticRow, state, rewriter, loc);
+    window = tensor::InsertSliceOp::create(rewriter,
+                                           loc,
+                                           paddedRow,
+                                           window,
+                                           SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
+                                                                      rewriter.getIndexAttr(0),
+                                                                      rewriter.getIndexAttr(kernelRowIndex),
+                                                                      rewriter.getIndexAttr(0)},
+                                           SmallVector<OpFoldResult> {rewriter.getIndexAttr(1),
+                                                                      rewriter.getIndexAttr(state.numChannelsIn),
+                                                                      rewriter.getIndexAttr(1),
+                                                                      rewriter.getIndexAttr(state.xWidth + 2)},
+                                           getUnitStrides(rewriter, 4));
+  }
+  return window;
+}
+
+static FailureOr<Value> createNchwRowStripConvPatchRow(Value paddedWindow,
+                                                       const ConvLoweringState& state,
+                                                       Value outputWidth,
+                                                       PatternRewriter& rewriter,
+                                                       Location loc) {
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth},
+                                         state.xType.getElementType(),
+                                         state.xType.getEncoding());
+  auto rowType = RankedTensorType::get({1, patchSize}, state.xType.getElementType(), state.xType.getEncoding());
+  Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
+  Value patch = createConvInputPatch(paddedWindow,
+                                     patchType,
+                                     c0,
+                                     c0,
+                                     c0,
+                                     outputWidth,
+                                     state.dilationHeight,
+                                     state.dilationWidth,
+                                     rewriter,
+                                     loc);
+  return tensor::CollapseShapeOp::create(
+           rewriter, loc, rowType, patch, SmallVector<ReassociationIndices> {{0}, {1, 2, 3}})
+    .getResult();
+}
+
+static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStripStorage,
+                                                                  const ConvLoweringState& state,
+                                                                  PatternRewriter& rewriter,
+                                                                  Location loc) {
+  auto inputType = dyn_cast<RankedTensorType>(rowStripStorage.getType());
+  if (!inputType || inputType != getRowStripStorageType(state.xType))
     return failure();
 
   StringRef failureReason;
-  if (!canDirectLowerRowStripConv(state, failureReason))
+  if (!canConsumeNchwRowStripFragments(state, failureReason))
-    return failure();
-
-  ConvRowDemand demand = buildConvRowDemand(RowInterval {0, state.outHeight}, state);
-  if (!covers(demand.acquiredInputRows, demand.neededInputRows))
     return failure();
 
   ConvGeometry geometry = buildConvGeometry(state);
@@ -2629,121 +2733,61 @@ static FailureOr<Value> createConvOutputFromRowStripHwc(Value inputHwc,
   const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
   const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
   const int64_t paddedK = numKSlices * xbarDim;
-  auto elementType = inputType.getElementType();
+  auto elementType = state.outType.getElementType();
-  auto paddedInputType = RankedTensorType::get({state.xHeight + state.padHeightBegin + state.padHeightEnd,
-                                                state.xWidth + state.padWidthBegin + state.padWidthEnd,
-                                                state.numChannelsIn},
-                                               elementType,
-                                               inputType.getEncoding());
-  auto paddedPatchType =
-    RankedTensorType::get({state.wHeight, state.wWidth, 1}, elementType, inputType.getEncoding());
-  auto flatPatchType = RankedTensorType::get({state.wHeight * state.wWidth}, elementType, inputType.getEncoding());
-  auto rowChunkType = RankedTensorType::get({1, state.wHeight * state.wWidth}, elementType, inputType.getEncoding());
   auto rowType = RankedTensorType::get({1, state.numChannelsOut}, state.outType.getElementType());
-  auto packedOutputType =
+  auto outputPixelType = RankedTensorType::get({1, state.numChannelsOut, 1, 1}, elementType);
-    RankedTensorType::get({state.outHeight, state.outWidth, state.numChannelsOut}, state.outType.getElementType());
-  auto packedOutputSliceType =
-    RankedTensorType::get({1, 1, state.numChannelsOut}, state.outType.getElementType());
   auto paddedRowType = RankedTensorType::get({1, xbarDim}, state.outType.getElementType());
   auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType, inputType.getEncoding());
   auto paddedWeightTileType = RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType());
+  auto outputStorageType = getRowStripStorageType(state.outType);
   auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
   if (!weightDenseAttr)
     return failure();
   Value paddedWeights = standard::createPaddedInputKTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
-
+  FailureOr<Value> paddedBias = failure();
-  Value paddedBias;
-  if (state.hasBias) {
-    Value biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc);
-    auto biasMatrixType = cast<RankedTensorType>(biasMatrix.getType());
-    auto paddedBiasType = RankedTensorType::get({1, xbarDim}, state.outType.getElementType());
-    if (auto biasDenseAttr = getHostConstDenseElementsAttr(state.b))
-      paddedBias = standard::createPaddedConstantMatrix(biasDenseAttr, biasMatrixType, paddedBiasType, rewriter);
-    else
-      paddedBias = materializeOrComputeUnary(
-        biasMatrix, paddedBiasType, rewriter, loc, [&](Value biasValue) {
-          return standard::createPaddedConvMatrix(biasValue, biasMatrixType, paddedBiasType, rewriter, loc);
-        });
-  }
-
-  auto paddedInputOp =
-    createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, inputHwc, [&](Value hwcInputArg) {
-      Value paddedInput = createZeroPaddedTensor(hwcInputArg,
-                                                 paddedInputType,
-                                                 {state.padHeightBegin, state.padWidthBegin, 0},
-                                                 {state.padHeightEnd, state.padWidthEnd, 0},
-                                                 rewriter,
-                                                 loc);
-      spatial::SpatYieldOp::create(rewriter, loc, paddedInput);
-    });
-
-  SmallVector<Value> batchInputs {paddedInputOp.getResult(0)};
   if (state.hasBias)
-    batchInputs.push_back(paddedBias);
+    paddedBias = createPaddedBiasRowConstant(state, xbarDim, rewriter);
+  if (state.hasBias && failed(paddedBias))
+    return failure();
+
   auto batchOp = createSpatComputeBatch(
     rewriter,
     loc,
-    TypeRange {packedOutputType},
+    TypeRange {outputStorageType},
     state.outHeight,
     ValueRange {paddedWeights},
-    batchInputs,
+    state.hasBias ? ValueRange {rowStripStorage, *paddedBias} : ValueRange {rowStripStorage},
     [&](detail::SpatComputeBatchBodyArgs args) {
       Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
       Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
       Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
       Value cNumKSlices = getOrCreateIndexConstant(rewriter, anchorOp, numKSlices);
       Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth);
-      Value cNumChannels = getOrCreateIndexConstant(rewriter, anchorOp, state.numChannelsIn);
+      Value cXbar = getOrCreateIndexConstant(rewriter, anchorOp, xbarDim);
-      Value localHeightOffset = args.lane;
+      auto fragmentType = getRowStripFragmentType(state.outType);
-      Value packedRowInit =
+      FailureOr<Value> inputWindow = createNchwRowStripConvWindow(args.inputs.front(), state, args.lane, rewriter, loc);
-        tensor::EmptyOp::create(rewriter, loc, ArrayRef<int64_t> {1, state.outWidth, state.numChannelsOut}, elementType);
+      if (failed(inputWindow))
+        return failure();
+      Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType);
       auto widthLoop = buildNormalizedScfFor(
         rewriter,
         loc,
         c0,
         cOutWidth,
         c1,
-        ValueRange {packedRowInit},
+        ValueRange {fragmentInit},
         [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl<Value>& widthYielded) {
-          Value localWidthOffset = widthIndex;
+          FailureOr<Value> patchRow =
-          Value rowInit = tensor::EmptyOp::create(rewriter, widthLoc, ArrayRef<int64_t> {1, patchSize}, elementType);
+            createNchwRowStripConvPatchRow(*inputWindow, state, widthIndex, rewriter, widthLoc);
-          auto rowLoop = buildNormalizedScfFor(
+          if (failed(patchRow))
-            rewriter,
-            widthLoc,
-            c0,
-            cNumChannels,
-            c1,
-            ValueRange {rowInit},
-            [&](OpBuilder&, Location rowLoc, Value channel, ValueRange rowIterArgs, SmallVectorImpl<Value>& rowYielded) {
-              SmallVector<OpFoldResult> patchOffsets {localHeightOffset, localWidthOffset, channel};
-              SmallVector<OpFoldResult> patchSizes {
-                rewriter.getIndexAttr(state.wHeight), rewriter.getIndexAttr(state.wWidth), rewriter.getIndexAttr(1)};
-              Value channelPatch = tensor::ExtractSliceOp::create(
-                rewriter, rowLoc, paddedPatchType, args.inputs.front(), patchOffsets, patchSizes, getUnitStrides(rewriter, 3));
-              Value flatPatch = tensor::CollapseShapeOp::create(
-                rewriter, rowLoc, flatPatchType, channelPatch, SmallVector<ReassociationIndices> {{0, 1, 2}});
-              Value rowChunk = tensor::ExpandShapeOp::create(
-                rewriter, rowLoc, rowChunkType, flatPatch, SmallVector<ReassociationIndices> {{0, 1}});
-              Value flatOffset = affineMulConst(
-                rewriter, rowLoc, channel, state.wHeight * state.wWidth, anchorOp);
-              SmallVector<OpFoldResult> rowOffsets {rewriter.getIndexAttr(0), flatOffset};
-              SmallVector<OpFoldResult> rowSizes {
-                rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.wHeight * state.wWidth)};
-              Value nextRow = tensor::InsertSliceOp::create(
-                rewriter, rowLoc, rowChunk, rowIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
-              rowYielded.push_back(nextRow);
-              return success();
-            });
-          if (failed(rowLoop))
             return failure();
 
-          Value paddedRow = rowLoop->results.front();
+          Value paddedRow = *patchRow;
           if (patchSize != paddedK)
             paddedRow = createZeroPaddedTensor(
               paddedRow, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, widthLoc);
 
-          auto zeroAttr = DenseElementsAttr::get(paddedRowType, rewriter.getZeroAttr(state.outType.getElementType()));
+          Value zeroRow = createZeroTensorConstant(paddedRowType, rewriter);
-          Value zeroRow = getOrCreateConstant(rewriter, anchorOp, zeroAttr, paddedRowType);
           auto kLoop = buildNormalizedScfFor(
             rewriter,
             widthLoc,
@@ -2752,7 +2796,7 @@ static FailureOr<Value> createConvOutputFromRowStripHwc(Value inputHwc,
             c1,
             ValueRange {zeroRow},
             [&](OpBuilder&, Location reduceLoc, Value kSlice, ValueRange reduceIterArgs, SmallVectorImpl<Value>& reduceYielded) {
-              Value kOffset = affineMulConst(rewriter, reduceLoc, kSlice, xbarDim, anchorOp);
+              Value kOffset = arith::MulIOp::create(rewriter, reduceLoc, kSlice, cXbar);
               SmallVector<OpFoldResult> aOffsets {rewriter.getIndexAttr(0), kOffset};
               SmallVector<OpFoldResult> aSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)};
               Value aTile = tensor::ExtractSliceOp::create(
@@ -2776,9 +2820,7 @@ static FailureOr<Value> createConvOutputFromRowStripHwc(Value inputHwc,
 
           Value rowResult = kLoop->results.front();
           if (state.hasBias)
-            rowResult =
+            rowResult = spatial::SpatVAddOp::create(rewriter, widthLoc, paddedRowType, rowResult, args.inputs[1]).getResult();
-              spatial::SpatVAddOp::create(rewriter, widthLoc, paddedRowType, rowResult, args.inputs[1]).getResult();
-
           Value outputRow = rowResult;
           if (state.numChannelsOut != xbarDim) {
             SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
@@ -2790,25 +2832,23 @@ static FailureOr<Value> createConvOutputFromRowStripHwc(Value inputHwc,
 
           Value outputFragment = tensor::ExpandShapeOp::create(rewriter,
                                                                widthLoc,
-                                                               packedOutputSliceType,
+                                                               outputPixelType,
                                                                outputRow,
-                                                               SmallVector<ReassociationIndices> {{0}, {1, 2}});
+                                                               SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
-          SmallVector<OpFoldResult> rowOffsets {rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
+          SmallVector<OpFoldResult> rowOffsets {
+            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex};
           SmallVector<OpFoldResult> rowSizes {
-            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
+            rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut), rewriter.getIndexAttr(1),
-          Value nextPackedRow = tensor::InsertSliceOp::create(
+            rewriter.getIndexAttr(1)};
-            rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 3));
+          Value nextFragment = tensor::InsertSliceOp::create(
-          widthYielded.push_back(nextPackedRow);
+            rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 4));
+          widthYielded.push_back(nextFragment);
           return success();
         });
       if (failed(widthLoop))
         return failure();
 
-      SmallVector<OpFoldResult> batchOffsets {args.lane, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+      insertRowStripFragment(widthLoop->results.front(), args.outputs.front(), state.outType, args.lane, rewriter, loc);
-      SmallVector<OpFoldResult> batchSizes {
-        rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.outWidth), rewriter.getIndexAttr(state.numChannelsOut)};
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, widthLoop->results.front(), args.outputs.front(), batchOffsets, batchSizes, getUnitStrides(rewriter, 3));
       return success();
     });
   if (failed(batchOp))
@@ -2816,19 +2856,22 @@ static FailureOr<Value> createConvOutputFromRowStripHwc(Value inputHwc,
   return batchOp->getResult(0);
 }
 
-static FailureOr<Value> createConvRowsFromRowStripInput(const ConvLoweringState& state,
+static FailureOr<Value> createConvOutputFromRowStripInput(const ConvLoweringState& state,
-                                                        [[maybe_unused]] const ConvLoweringDecision& decision,
+                                                          [[maybe_unused]] const ConvLoweringDecision& decision,
-                                                        Value rowStripInput,
+                                                          Value rowStripInput,
-                                                        PatternRewriter& rewriter,
+                                                          PatternRewriter& rewriter,
-                                                        Location loc) {
+                                                          Location loc) {
-  return createConvOutputFromRowStripHwc(rowStripInput, state, rewriter, loc);
+  return createConvOutputFromNchwRowStripFragments(rowStripInput, state, rewriter, loc);
 }
 
 static Value createFragmentConstant(const DistributedTensorStep& step,
                                     RankedTensorType fragmentType,
                                     PatternRewriter& rewriter) {
-  if (step.constantKind == DistributedTensorConstantKind::PerChannel)
+  if (step.constantKind == DistributedTensorConstantKind::PerChannel) {
-    return createPerChannelConstantFragment(step.constantAttr, fragmentType, rewriter);
+    FailureOr<Value> constant = createPerChannelConstantFragment(step.constantAttr, fragmentType, rewriter);
+    assert(succeeded(constant) && "distributed per-channel constants are classified before lowering");
+    return *constant;
+  }
 
   Attribute splatValue = step.constantAttr.getSplatValue<Attribute>();
   return getOrCreateConstant(rewriter,
@@ -2938,67 +2981,22 @@ static Value createFragmentReciprocalConstant(const DistributedTensorStep& step,
   }
 }
 
-[[maybe_unused]] static FailureOr<DistributedTensorInfo> createDistributedTensorFromRows(Value rows,
-                                                                                         RankedTensorType logicalType,
-                                                                                         PatternRewriter& rewriter,
-                                                                                         Location loc) {
-  const int64_t width = logicalType.getDimSize(3);
-  const int64_t height = logicalType.getDimSize(2);
-  auto rowsType = cast<RankedTensorType>(rows.getType());
-  auto rowSliceType =
-    RankedTensorType::get({width, logicalType.getDimSize(1)}, logicalType.getElementType(), rowsType.getEncoding());
-  auto channelWidthType =
-    RankedTensorType::get({logicalType.getDimSize(1), width}, logicalType.getElementType(), rowsType.getEncoding());
-  auto fragmentType = getRowStripFragmentType(logicalType, width);
-  auto batchOp = createSpatComputeBatch(
-    rewriter, loc, TypeRange {logicalType}, height, {}, ValueRange {rows}, [&](detail::SpatComputeBatchBodyArgs args) {
-      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
-      Value rowStart = affineMulConst(rewriter, loc, args.lane, width, anchorOp);
-      SmallVector<OpFoldResult> rowOffsets {rowStart, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(width), rewriter.getIndexAttr(logicalType.getDimSize(1))};
-      Value rowSlice = tensor::ExtractSliceOp::create(
-        rewriter, loc, rowSliceType, args.inputs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
-      Value channelWidth = ONNXTransposeOp::create(
-        rewriter, loc, channelWidthType, rowSlice, rewriter.getI64ArrayAttr({1, 0})).getResult();
-      Value fragment = tensor::ExpandShapeOp::create(rewriter,
-                                                     loc,
-                                                     fragmentType,
-                                                     channelWidth,
-                                                     SmallVector<ReassociationIndices> {{0, 1}, {2, 3}});
-      SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), args.lane,
-                                               rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(logicalType.getDimSize(1)),
-                                             rewriter.getIndexAttr(1), rewriter.getIndexAttr(width)};
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, fragment, args.outputs.front(), outputOffsets, outputSizes, getUnitStrides(rewriter, 4));
-      return success();
-    });
-  if (failed(batchOp))
-    return failure();
-  return makeDistributedTensorInfo(batchOp->getResult(0), logicalType);
-}
-
 [[maybe_unused]] static FailureOr<DistributedTensorInfo> applyDistributedPreservingStep(const DistributedTensorInfo& inputInfo,
                                                                                         const DistributedTensorStep& step,
                                                                                         PatternRewriter& rewriter,
                                                                                         Location loc) {
   auto logicalType = inputInfo.logicalType;
-  const int64_t width = logicalType.getDimSize(3);
+  auto fragmentType = getRowStripFragmentType(logicalType);
-  auto fragmentType = getRowStripFragmentType(logicalType, width);
+  auto storageType = getRowStripStorageType(logicalType);
   auto batchOp = createSpatComputeBatch(rewriter,
                                         loc,
-                                        TypeRange {logicalType},
+                                        TypeRange {storageType},
                                         inputInfo.laneCount,
                                         {},
                                         ValueRange {inputInfo.storage},
                                         [&](detail::SpatComputeBatchBodyArgs args) {
-                                          SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
+                                          Value fragment =
-                                                                             args.lane, rewriter.getIndexAttr(0)};
+                                            extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
-                                          SmallVector<OpFoldResult> sizes {
-                                            rewriter.getIndexAttr(1), rewriter.getIndexAttr(logicalType.getDimSize(1)),
-                                            rewriter.getIndexAttr(1), rewriter.getIndexAttr(width)};
-                                          Value fragment = tensor::ExtractSliceOp::create(
-                                            rewriter, loc, fragmentType, args.inputs.front(), offsets, sizes, getUnitStrides(rewriter, 4));
                                           switch (step.kind) {
                                           case DistributedTensorOpKind::Relu:
                                             fragment = spatial::SpatReluOp::create(rewriter, loc, fragmentType, fragment).getResult();
@@ -3034,8 +3032,8 @@ static Value createFragmentReciprocalConstant(const DistributedTensorStep& step,
                                           case DistributedTensorOpKind::Conv:
                                             return failure();
                                           }
-                                          createParallelInsertSliceIntoBatchOutput(
+                                          insertRowStripFragment(
-                                            rewriter, loc, fragment, args.outputs.front(), offsets, sizes, getUnitStrides(rewriter, 4));
+                                            fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
                                           return success();
                                         });
   if (failed(batchOp))
@@ -3743,6 +3741,8 @@ LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) {
     return failure();
   if (state->outType.getRank() != 4 || !state->outType.hasStaticShape())
     return failure();
+  if (state->hasBias && !isSupportedBiasAddValue(state->b, state->outType))
+    return failure();
 
   FailureOr<PimConvLoweringType> requestedStrategy = resolveRequestedConvLoweringStrategy(planOp.getOperation());
   if (failed(requestedStrategy))
@@ -3786,7 +3786,7 @@ LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp) {
     return failure();
 
   StringRef failureReason;
-  return canDirectLowerRowStripConv(*state, failureReason) ? success() : failure();
+  return canConsumeNchwRowStripFragments(*state, failureReason) ? success() : failure();
 }
 
 FailureOr<Value>
@@ -3822,17 +3822,33 @@ lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
     if (rowStripInput) {
       if (failed(canConsumeAndProduceRowStrip(planOp)))
         return planOp.emitOpError("selected row-strip input/output layout is not supported for this Conv plan"), failure();
-      return createConvRowsFromRowStripInput(*state, decision, *rowStripInput, rewriter, planOp.getLoc());
+      return createConvOutputFromRowStripInput(*state, decision, *rowStripInput, rewriter, planOp.getLoc());
     }
     if (failed(canLowerConvPlanToRowStrip(planOp)))
       return planOp.emitOpError("selected row-strip layout is not supported for this Conv plan"), failure();
-    FailureOr<Value> rows = createConvRowsForStrategy(*state, decision, rewriter, planOp.getLoc());
+    ConvLoweringState rowState = *state;
+    const bool applyBiasAfterStorage = rowState.hasBias;
+    Value originalBias = rowState.b;
+    if (applyBiasAfterStorage) {
+      if (!isSupportedBiasAddValue(originalBias, rowState.outType))
+        return planOp.emitOpError("selected row-strip Conv bias must be host-constant scalar/per-channel NCHW"),
+               failure();
+      rowState.b = Value();
+      rowState.hasBias = false;
+    }
+
+    FailureOr<Value> rows = createConvRowsForStrategy(rowState, decision, rewriter, planOp.getLoc());
     if (failed(rows))
       return failure();
-    FailureOr<Value> packedRows = createRowStripPackedRows(*rows, *state, rewriter, planOp.getLoc());
+    FailureOr<Value> rowStripStorage = createRowStripStorageFromRows(*rows, state->outType, rewriter, planOp.getLoc());
-    if (failed(packedRows))
+    if (failed(rowStripStorage))
-      return planOp.emitOpError("failed to pack Conv rows into the selected row-strip physical layout"), failure();
+      return planOp.emitOpError("failed to build row-strip fragment storage for the selected Conv plan"), failure();
-    return *packedRows;
+    if (applyBiasAfterStorage) {
+      rowStripStorage = applyRowStripBiasAdd(*rowStripStorage, rowState.outType, originalBias, rewriter, planOp.getLoc());
+      if (failed(rowStripStorage))
+        return planOp.emitOpError("failed to apply row-strip Conv bias per fragment"), failure();
+    }
+    return *rowStripStorage;
   }
 
   if (decision.strategy == PimConvLoweringDepthwise)

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp

@@ -5,7 +5,7 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -47,38 +47,28 @@ static FailureOr<Value> materializeBroadcastedConstantTensor(Value value,
     return failure();
 
   const int64_t rankOffset = static_cast<int64_t>(resultShape.size() - sourceShape.size());
-  for (int64_t i = 0; i < static_cast<int64_t>(resultShape.size()); ++i) {
-    const int64_t sourceIndex = i - rankOffset;
-    const int64_t sourceDim = sourceIndex < 0 ? 1 : sourceShape[sourceIndex];
-    const int64_t resultDim = resultShape[i];
-    if (sourceDim != 1 && sourceDim != resultDim)
-      return failure();
-  }
-
-  SmallVector<Attribute> sourceValues(denseAttr.getValues<Attribute>());
   SmallVector<int64_t> sourceStrides = computeRowMajorStrides(sourceShape);
   SmallVector<int64_t> resultStrides = computeRowMajorStrides(resultShape);
-
+  SmallVector<Attribute> sourceValues(denseAttr.getValues<Attribute>());
   SmallVector<Attribute> resultValues;
   resultValues.reserve(resultType.getNumElements());
 
   for (int64_t flatIndex = 0; flatIndex < resultType.getNumElements(); ++flatIndex) {
     int64_t remaining = flatIndex;
     int64_t sourceFlatIndex = 0;
-
     for (int64_t i = 0; i < static_cast<int64_t>(resultShape.size()); ++i) {
       const int64_t resultIndex = resultStrides.empty() ? 0 : remaining / resultStrides[i];
       remaining = resultStrides.empty() ? 0 : remaining % resultStrides[i];
-
       const int64_t sourceIndex = i - rankOffset;
       if (sourceIndex < 0)
         continue;
-
       const int64_t sourceDim = sourceShape[sourceIndex];
+      const int64_t resultDim = resultShape[i];
+      if (sourceDim != 1 && sourceDim != resultDim)
+        return failure();
       const int64_t mappedIndex = sourceDim == 1 ? 0 : resultIndex;
       sourceFlatIndex += mappedIndex * sourceStrides[sourceIndex];
     }
-
     resultValues.push_back(sourceValues[sourceFlatIndex]);
   }
 
@@ -106,7 +96,7 @@ static FailureOr<Value> materializeReciprocalTensor(Value value,
   if (failed(broadcastedValue))
     return failure();
 
-  auto denseAttr = dyn_cast<DenseFPElementsAttr>(getDenseConstantAttr(*broadcastedValue));
+  auto denseAttr = dyn_cast<DenseFPElementsAttr>(getHostConstDenseElementsAttr(*broadcastedValue));
   if (!denseAttr)
     return failure();
 
@@ -185,10 +175,45 @@ struct DivToSpatialCompute : OpConversionPattern<ONNXDivOp> {
   }
 };
 
+struct AddToSpatialCompute : OpConversionPattern<ONNXAddOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ONNXAddOp op, ONNXAddOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override {
+    auto resultType = dyn_cast<RankedTensorType>(op.getResult().getType());
+    if (!resultType || !resultType.hasStaticShape())
+      return failure();
+
+    FailureOr<BiasAddPlanCandidate> candidate =
+      classifyBiasAddPlanCandidate(adaptor.getA(), adaptor.getB(), resultType);
+    if (succeeded(candidate)) {
+      auto plan = spatial::SpatBiasAddPlanOp::create(
+        rewriter, op.getLoc(), resultType, candidate->data, candidate->bias, rewriter.getStringAttr("nchw"));
+      rewriter.replaceOp(op, plan.getResult());
+      return success();
+    }
+
+    auto lhs = prepareElementwiseOperand(adaptor.getA(), resultType, rewriter, op.getLoc());
+    if (failed(lhs))
+      return failure();
+    auto rhs = prepareElementwiseOperand(adaptor.getB(), resultType, rewriter, op.getLoc());
+    if (failed(rhs))
+      return failure();
+
+    auto computeOp =
+      createSpatCompute<2>(rewriter, op.getLoc(), resultType, {}, ValueRange {*lhs, *rhs}, [&](Value x, Value y) {
+        auto loweredOp = spatial::SpatVAddOp::create(rewriter, op.getLoc(), resultType, x, y);
+        spatial::SpatYieldOp::create(rewriter, op.getLoc(), loweredOp.getResult());
+      });
+    rewriter.replaceOp(op, computeOp);
+    return success();
+  }
+};
+
 } // namespace
 
 void populateElementwisePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
-  patterns.add<BinaryElementwiseToSpatialCompute<ONNXAddOp, spatial::SpatVAddOp>>(ctx);
+  patterns.add<AddToSpatialCompute>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXSubOp, spatial::SpatVSubOp>>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXMulOp, spatial::SpatVMulOp>>(ctx);
   patterns.add<DivToSpatialCompute>(ctx);

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -6,10 +6,11 @@
 
 #include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Accelerators/PIM/Pass/PIMPasses.h"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
 
 using namespace mlir;
 
@@ -19,7 +20,6 @@ namespace {
 static constexpr StringLiteral kLogicalLayout = "nchw";
 static constexpr StringLiteral kDenseLayout = "dense_nchw";
 static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";
-static constexpr StringLiteral kRowStripIndexMap = "packed_hwc_rows_to_nchw";
 
 enum class SelectedLayout {
   DenseNchw,
@@ -34,6 +34,8 @@ static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& l
 static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, SelectedLayout>& layouts) {
   if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(user))
     return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::NchwRowStrip;
+  if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user))
+    return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::NchwRowStrip;
   if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
     return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::NchwRowStrip;
   return false;
@@ -49,21 +51,26 @@ static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap<Value, Selecte
   return true;
 }
 
-static std::pair<SmallVector<int64_t>, SmallVector<int64_t>> buildRowStripMetadata(RankedTensorType type) {
+static bool canConsumeRowStripAsUser(Operation* user) {
-  SmallVector<int64_t> offsets;
+  if (isa<spatial::SpatReluPlanOp>(user))
-  SmallVector<int64_t> sizes;
+    return true;
-  const int64_t channels = type.getDimSize(1);
+  if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user)) {
-  const int64_t height = type.getDimSize(2);
+    auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
-  const int64_t width = type.getDimSize(3);
+    return resultType && isSupportedBiasAddValue(biasAddPlan.getBias(), resultType);
-  offsets.reserve(height * 4);
-  sizes.reserve(height * 4);
-  for (int64_t row = 0; row < height; ++row) {
-    offsets.append({0, 0, row, 0});
-    sizes.append({1, channels, 1, width});
   }
-  return {offsets, sizes};
+  if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
+    return succeeded(canConsumeAndProduceRowStrip(convPlan));
+  return false;
 }
 
+static bool hasRowStripConsumer(Value value) {
+  for (Operation* user : value.getUsers())
+    if (canConsumeRowStripAsUser(user))
+      return true;
+  return false;
+}
+
+
 static bool canSelectConvRowStrip(spatial::SpatConv2DPlanOp convPlan,
                                   llvm::DenseMap<Value, SelectedLayout>& layouts) {
   SelectedLayout inputLayout = getSelectedLayout(layouts, convPlan.getInput());
@@ -76,6 +83,9 @@ static SelectedLayout chooseConvLayout(spatial::SpatConv2DPlanOp convPlan,
                                        llvm::DenseMap<Value, SelectedLayout>& layouts) {
   if (!canSelectConvRowStrip(convPlan, layouts))
     return SelectedLayout::DenseNchw;
+  if (getSelectedLayout(layouts, convPlan.getInput()) != SelectedLayout::NchwRowStrip
+      && !hasRowStripConsumer(convPlan.getResult()))
+    return SelectedLayout::DenseNchw;
   if (!allUsersCanHandleRowStrip(convPlan.getResult(), layouts))
     return SelectedLayout::DenseNchw;
   return SelectedLayout::NchwRowStrip;
@@ -85,11 +95,27 @@ static SelectedLayout chooseReluLayout(spatial::SpatReluPlanOp reluPlan,
                                        llvm::DenseMap<Value, SelectedLayout>& layouts) {
   if (getSelectedLayout(layouts, reluPlan.getInput()) != SelectedLayout::NchwRowStrip)
     return SelectedLayout::DenseNchw;
+  if (!hasRowStripConsumer(reluPlan.getResult()))
+    return SelectedLayout::DenseNchw;
   if (!allUsersCanHandleRowStrip(reluPlan.getResult(), layouts))
     return SelectedLayout::DenseNchw;
   return SelectedLayout::NchwRowStrip;
 }
 
+static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan,
+                                          llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::NchwRowStrip)
+    return SelectedLayout::DenseNchw;
+  auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
+  if (!resultType || !isSupportedBiasAddValue(biasAddPlan.getBias(), resultType))
+    return SelectedLayout::DenseNchw;
+  if (!hasRowStripConsumer(biasAddPlan.getResult()))
+    return SelectedLayout::DenseNchw;
+  if (!allUsersCanHandleRowStrip(biasAddPlan.getResult(), layouts))
+    return SelectedLayout::DenseNchw;
+  return SelectedLayout::NchwRowStrip;
+}
+
 static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Value value) {
   auto outputType = cast<RankedTensorType>(value.getType());
   auto [offsets, sizes] = buildRowStripMetadata(outputType);
@@ -173,6 +199,14 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
           }
           continue;
         }
+        if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
+          SelectedLayout selected = chooseBiasAddLayout(biasAddPlan, layouts);
+          if (layouts[biasAddPlan.getResult()] != selected) {
+            layouts[biasAddPlan.getResult()] = selected;
+            changed = true;
+          }
+          continue;
+        }
       }
     }
 
@@ -180,6 +214,8 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
       Value producedValue;
       if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(&op))
         producedValue = convPlan.getResult();
+      else if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op))
+        producedValue = biasAddPlan.getResult();
       else if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op))
         producedValue = reluPlan.getResult();
       else

src/PIM/Conversion/SpatialToGraphviz/CMakeLists.txt

@@ -1,17 +0,0 @@
-add_onnx_mlir_rewriter(SpatialToGraphviz)
-
-add_pim_library(OMSpatialToGraphviz
-  SpatialToGraphviz.cpp
-
-  EXCLUDE_FROM_OM_LIBS
-
-  LINK_LIBS PUBLIC
-  MLIRTosaDialect
-  OMCompilerOptions
-  OMPimCommon
-  OMONNXOps
-  SpatialOps
-
-  ACCEL_INCLUDE_DIRS PRIVATE
-  ${PIM_GENERATED_INCLUDE_DIRS}
-)

src/PIM/Conversion/SpatialToGraphviz/SpatialToGraphviz.cpp

@@ -1,259 +0,0 @@
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Dialect/Tosa/IR/TosaOps.h"
-#include "mlir/IR/Block.h"
-#include "mlir/IR/Diagnostics.h"
-#include "mlir/IR/Value.h"
-#include "mlir/Pass/Pass.h"
-#include "mlir/Support/LLVM.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-
-#include "llvm/Support/Casting.h"
-#include "llvm/Support/Format.h"
-
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
-#include "src/Dialect/ONNX/ONNXOps.hpp"
-
-#define FORMAT_OPERATION(op) 'x' << llvm::format_hex_no_prefix(reinterpret_cast<size_t>(op), 0)
-#define FORMAT_ARGUMENT(computeOpPointer, argumentNum) llvm::format("Arg_%p_%u", computeOpPointer, argumentNum)
-
-using namespace mlir;
-
-namespace onnx_mlir {
-
-namespace {
-
-struct SpatialToGraphvizPass : public PassWrapper<SpatialToGraphvizPass, OperationPass<ModuleOp>> {
-
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialToGraphvizPass)
-
-  StringRef getArgument() const override { return "convert-spatial-to-graphviz"; }
-
-  StringRef getDescription() const override { return "Lower ONNX ops to Spatial ops."; }
-
-  SpatialToGraphvizPass(raw_ostream& os = llvm::errs())
-  : os(os) {}
-  SpatialToGraphvizPass(const SpatialToGraphvizPass& pass)
-  : SpatialToGraphvizPass(pass.os) {}
-  void runOnOperation() final;
-
-private:
-  raw_ostream& os;
-
-  /**
-   * Draws the subgraph for a given spatial::SpatCompute, including:
-   * 1. Input nodes (block arguments)
-   * 2. Operations
-   * 3. Edges between yield (output) and its users
-   *
-   * @param op The spatial::SpatCompute to draw the subgraph for.
-   * @param computeNum The number of the compute operation.
-   */
-  void drawComputeOpSubgraph(spatial::SpatCompute op, size_t computeNum) {
-    os << "\tsubgraph cluster" << computeNum << " {\n\t\tlabel=\"Compute" << computeNum << "\";\n"
-       << "\t\tstyle=filled;\n"
-       << "\t\tcolor=lightblue;\n";
-
-    Block& block = op.getBody().front();
-
-    // Inputs
-    size_t inputNum = 0;
-    for (BlockArgument& input : block.getArguments()) {
-
-      auto fromOp = FORMAT_ARGUMENT(op.getOperation(), inputNum);
-
-      os << "\t\t" << fromOp << " [label=\"Arg" << inputNum << "\",shape=box];\n";
-      for (auto userOp : input.getUsers())
-        os << "\t\t" << fromOp << " -> " << FORMAT_OPERATION(userOp) << ";\n";
-      inputNum++;
-    }
-
-    // Iterate operations
-    for (auto& childOp : block.getOperations()) {
-      os << "\t\t" << FORMAT_OPERATION(&childOp) << " [label=\"" << childOp.getName() << "\"];\n";
-
-      drawEdgesFromOpToItsUsers(&childOp);
-    }
-
-    os << "\t}\n";
-
-    // Draw edges from the yield to the users of this computeOp
-    Operation* yieldOp = block.getTerminator();
-    if (!isa<spatial::SpatYieldOp>(yieldOp)) {
-      yieldOp->emitError("Terminator of block must be YieldOp ???");
-      signalPassFailure();
-      return;
-    }
-
-    for (auto computeOpResult : op->getResults()) {
-      for (auto& computeOpUse : computeOpResult.getUses()) {
-        auto toOp = FORMAT_ARGUMENT(computeOpUse.getOwner(), computeOpUse.getOperandNumber());
-        os << "\t" << FORMAT_OPERATION(yieldOp) << " -> " << toOp << ";\n";
-      }
-    }
-  }
-
-  /**
-   * @brief Draws the subgraph for a concatOp.
-   *
-   * This function draws a subgraph for a concatOp. The subgraph consists of a
-   * node for each input of the concatOp, as well as an output node. Edges are
-   * created from the output node to each user of the concatOp.
-   *
-   * @param concatOp The concatOp for which the subgraph is drawn.
-   * @param concatOpNum The number of the concatOp.
-   */
-  void drawConcatOpSubgraph(Operation* concatOp, size_t concatOpNum) {
-    os << "\tsubgraph clusterconcat" << concatOpNum << " {\n\t\tlabel=\"ConcatOp" << concatOpNum << "\";\n"
-       << "\t\tstyle=filled;\n"
-       << "\t\tcolor=orange;\n";
-
-    // Inputs
-    size_t inputNum = 0;
-    for (Value input : concatOp->getOperands()) {
-      auto fromOp = FORMAT_ARGUMENT(concatOp, inputNum);
-
-      os << "\t\t" << fromOp << " [label=\"Input" << inputNum << "\"];\n";
-      for (auto userOp : input.getUsers())
-        os << "\t\t" << fromOp << " -> " << FORMAT_OPERATION(userOp) << ";\n";
-      inputNum++;
-    }
-
-    // Output
-    os << "\t\t" << FORMAT_OPERATION(concatOp) << " [label=Out];\n";
-
-    os << "\t}\n";
-
-    // Edges from output to users
-
-    for (auto& computeOpUse : concatOp->getResult(0).getUses()) {
-      os << "\t" << FORMAT_OPERATION(concatOp) << " -> "
-         << FORMAT_ARGUMENT(computeOpUse.getOwner(), computeOpUse.getOperandNumber()) << ";\n";
-    }
-  }
-
-  /**
-   * Draws the ExtractSliceOp in the graph visualization.
-   *
-   * This function takes a tensor::ExtractSliceOp and adds the corresponding
-   * node and edges to the graph visualization. It creates a node with the
-   * label as the static offsets attribute of the sliceOp, and connects it to
-   * the compute operations that use the result of the sliceOp.
-   *
-   * @param sliceOp The tensor::ExtractSliceOp to be drawn in the graph
-   * visualization.
-   */
-  void drawExtractSliceOp(tensor::ExtractSliceOp sliceOp) {
-    auto nodeId = FORMAT_ARGUMENT(sliceOp.getOperation(), 0);
-    os << "\t" << nodeId << " [label=\"Slice: ";
-    sliceOp.getStaticOffsetsAttr().print(os);
-    os << "\",color=lawngreen];\n";
-
-    for (auto& computeOpUse : sliceOp.getResult().getUses()) {
-      os << "\t" << nodeId << " -> " << FORMAT_ARGUMENT(computeOpUse.getOwner(), computeOpUse.getOperandNumber())
-         << ";\n";
-    }
-  }
-
-  void drawBiasTileOp(tensor::ExtractSliceOp sliceOp) {
-    auto nodeId = FORMAT_ARGUMENT(sliceOp.getOperation(), 0);
-    os << "\t" << nodeId << " [label=\"Bias: ";
-    sliceOp.getStaticOffsetsAttr().print(os);
-    os << "\",color=lightpink];\n";
-
-    for (auto user : sliceOp.getResult().getUsers())
-      os << "\t" << nodeId << " -> " << FORMAT_OPERATION(user) << ";\n";
-  }
-
-  /**
-   * Draws edges from the given operation to its users.
-   *
-   * @param fromOp The operation from which the edges are drawn.
-   */
-  void drawEdgesFromOpToItsUsers(mlir::Operation* fromOp) {
-    for (auto result : fromOp->getResults())
-      for (auto userOp : result.getUsers())
-        os << "\t\t" << FORMAT_OPERATION(fromOp) << " -> " << FORMAT_OPERATION(userOp) << ";\n";
-  }
-
-  /**
-   * Draws input node and edges for the given `funcOp`.
-   *
-   * @param funcOp The `funcOp` for which to draw input nodes and edges.
-   */
-  void drawInputNodesAndEdges(func::FuncOp& funcOp) {
-    os << "\tinput [label=\"Module Input\",color=green];\n";
-
-    size_t funcOpArgNum = 0;
-    for (BlockArgument& arg : funcOp.getArguments()) {
-
-      for (auto& useOp : arg.getUses()) {
-        os << "\tinput -> " << FORMAT_ARGUMENT(useOp.getOwner(), useOp.getOperandNumber()) << "[label=" << funcOpArgNum
-           << "];\n";
-      }
-      funcOpArgNum++;
-    }
-  }
-};
-
-void SpatialToGraphvizPass::runOnOperation() {
-  ModuleOp module = getOperation();
-
-  auto entryFunc = getPimEntryFunc(module);
-  if (failed(entryFunc)) {
-    module.emitError("failed to locate the PIM entry function for Spatial graph visualization");
-    signalPassFailure();
-    return;
-  }
-  func::FuncOp func = *entryFunc;
-
-  os << "digraph G {\n"
-     << "\tnode [style=filled,color=white];\n";
-
-  size_t computeNum = 0;
-  size_t concatNum = 0;
-
-  // Iterate over the ComputeOps within FuncOp:
-  // 1. Print their subgraph
-  // 2. Print the edges from its inputs to its outputs
-  for (Operation& op : func.getOps()) {
-    if (auto computeOp = dyn_cast<spatial::SpatCompute>(op)) {
-      drawComputeOpSubgraph(computeOp, computeNum++);
-    }
-    else if (auto concatOp = dyn_cast<tensor::ConcatOp>(op)) {
-      drawConcatOpSubgraph(concatOp, concatNum++);
-    }
-    else if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(op)) {
-      auto producerOp = extractSliceOp->getOperand(0).getDefiningOp();
-      if (producerOp) {
-        // Skip extractSliceOp if producer is constant weights (ONNXConstantOp)
-        if (llvm::isa<ONNXConstantOp>(producerOp))
-          continue;
-        // If produced by tosa::ReshapeOp (i.e. it is a bias tile) connect
-        // directly to its user, which is not a ComputeOp argument.
-        if (llvm::isa<tosa::ReshapeOp>(producerOp)) {
-          drawBiasTileOp(extractSliceOp);
-          continue;
-        }
-      }
-
-      drawExtractSliceOp(extractSliceOp);
-    }
-  }
-
-  // Draw input node, and edges to it users
-  drawInputNodesAndEdges(func);
-
-  // Draw output node (use the return Operation - argument number=0 - as nodeId)
-  auto returnOp = func.getBody().front().getTerminator();
-  os << '\t' << FORMAT_ARGUMENT(returnOp, 0) << " [label=\"Module Output\",color=green];\n";
-
-  os << "}\n";
-}
-
-} // namespace
-
-std::unique_ptr<Pass> createSpatialToGraphvizPass() { return std::make_unique<SpatialToGraphvizPass>(); }
-
-} // namespace onnx_mlir

src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp

@@ -302,76 +302,87 @@ void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncO
 
 LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp moduleOp) const {
   bool hasFailure = false;
-  moduleOp.walk([&](Operation* op) {
-    auto verifyOperand = [&](Value operand, unsigned operandIndex) {
-      if (!isa<BaseMemRefType>(operand.getType()))
-        return;
-      if (succeeded(resolveContiguousAddress(operand)) || succeeded(compileContiguousAddressExpr(operand)))
-        return;
-      op->emitOpError() << "operand #" << operandIndex
-                        << " is not backed by contiguous addressable storage after PIM bufferization";
-      hasFailure = true;
-    };
 
-    if (auto memCopyOp = dyn_cast<PimMemCopyOp>(op)) {
+  auto verifyWithKnowledge = [&](auto coreLikeOp, const StaticValueKnowledge& initialKnowledge) {
-      if (!pim::isNormalizedCopyOp(memCopyOp)) {
+    (void) walkPimCoreBlockStructurally(
-        memCopyOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
+      coreLikeOp.getBody().front(), initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) {
-        hasFailure = true;
+        auto verifyOperand = [&](Value operand, unsigned operandIndex) {
-      }
+          if (!isa<BaseMemRefType>(operand.getType()))
-      verifyOperand(memCopyOp.getTarget(), 0);
+            return;
-      verifyOperand(memCopyOp.getSource(), 1);
+          if (succeeded(resolveContiguousAddress(operand, knowledge)) || succeeded(compileContiguousAddressExpr(operand)))
-      return;
+            return;
-    }
+          op.emitOpError() << "operand #" << operandIndex
-    if (auto loadOp = dyn_cast<PimMemCopyHostToDevOp>(op)) {
+                           << " is not backed by contiguous addressable storage after PIM bufferization";
-      if (!pim::isNormalizedCopyOp(loadOp)) {
+          hasFailure = true;
-        loadOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
+        };
-        hasFailure = true;
+
-      }
+        if (auto memCopyOp = dyn_cast<PimMemCopyOp>(&op)) {
-      verifyOperand(loadOp.getDeviceTarget(), 2);
+          if (!pim::isNormalizedCopyOp(memCopyOp)) {
-      verifyOperand(loadOp.getHostSource(), 3);
+            memCopyOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
-      return;
+            hasFailure = true;
-    }
+          }
-    if (auto storeOp = dyn_cast<PimMemCopyDevToHostOp>(op)) {
+          verifyOperand(memCopyOp.getTarget(), 0);
-      if (!pim::isNormalizedCopyOp(storeOp)) {
+          verifyOperand(memCopyOp.getSource(), 1);
-        storeOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
+          return success();
-        hasFailure = true;
+        }
-      }
+        if (auto loadOp = dyn_cast<PimMemCopyHostToDevOp>(&op)) {
-      verifyOperand(storeOp.getHostTarget(), 2);
+          if (!pim::isNormalizedCopyOp(loadOp)) {
-      verifyOperand(storeOp.getDeviceSource(), 3);
+            loadOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
-      return;
+            hasFailure = true;
-    }
+          }
-    if (auto sendOp = dyn_cast<PimSendOp>(op)) {
+          verifyOperand(loadOp.getDeviceTarget(), 2);
-      verifyOperand(sendOp.getInput(), 0);
+          verifyOperand(loadOp.getHostSource(), 3);
-      return;
+          return success();
-    }
+        }
-    if (auto receiveOp = dyn_cast<PimReceiveOp>(op)) {
+        if (auto storeOp = dyn_cast<PimMemCopyDevToHostOp>(&op)) {
-      verifyOperand(receiveOp.getOutputBuffer(), 0);
+          if (!pim::isNormalizedCopyOp(storeOp)) {
-      return;
+            storeOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
-    }
+            hasFailure = true;
-    if (auto concatOp = dyn_cast<PimConcatOp>(op)) {
+          }
-      verifyOperand(concatOp.getOutputBuffer(), 0);
+          verifyOperand(storeOp.getHostTarget(), 2);
-      for (auto inputAndIndex : llvm::enumerate(concatOp.getInputs()))
+          verifyOperand(storeOp.getDeviceSource(), 3);
-        verifyOperand(inputAndIndex.value(), inputAndIndex.index() + 1);
+          return success();
-      return;
+        }
-    }
+        if (auto sendOp = dyn_cast<PimSendOp>(&op)) {
-    if (isa<PimTransposeOp,
+          verifyOperand(sendOp.getInput(), 0);
-            PimVMMOp,
+          return success();
-            PimVVAddOp,
+        }
-            PimVVSubOp,
+        if (auto receiveOp = dyn_cast<PimReceiveOp>(&op)) {
-            PimVVMulOp,
+          verifyOperand(receiveOp.getOutputBuffer(), 0);
-            PimVVMaxOp,
+          return success();
-            PimVVDMulOp,
+        }
-            PimVAvgOp,
+        if (auto concatOp = dyn_cast<PimConcatOp>(&op)) {
-            PimVReluOp,
+          verifyOperand(concatOp.getOutputBuffer(), 0);
-            PimVTanhOp,
+          for (auto inputAndIndex : llvm::enumerate(concatOp.getInputs()))
-            PimVSigmOp,
+            verifyOperand(inputAndIndex.value(), inputAndIndex.index() + 1);
-            PimVSoftmaxOp>(op)) {
+          return success();
-      for (auto operandAndIndex : llvm::enumerate(op->getOperands())) {
+        }
-        if (auto vmmOp = dyn_cast<PimVMMOp>(op); vmmOp && operandAndIndex.index() == 0)
+        if (isa<PimTransposeOp,
-          continue;
+                PimVMMOp,
-        verifyOperand(operandAndIndex.value(), operandAndIndex.index());
+                PimVVAddOp,
-      }
+                PimVVSubOp,
-    }
+                PimVVMulOp,
+                PimVVMaxOp,
+                PimVVDMulOp,
+                PimVAvgOp,
+                PimVReluOp,
+                PimVTanhOp,
+                PimVSigmOp,
+                PimVSoftmaxOp>(&op)) {
+          for (auto operandAndIndex : llvm::enumerate(op.getOperands())) {
+            if (auto vmmOp = dyn_cast<PimVMMOp>(&op); vmmOp && operandAndIndex.index() == 0)
+              continue;
+            verifyOperand(operandAndIndex.value(), operandAndIndex.index());
+          }
+        }
+        return success();
+      });
+  };
+
+  moduleOp.walk([&](pim::PimCoreOp coreOp) { verifyWithKnowledge(coreOp, seedCoreKnowledge(coreOp)); });
+  moduleOp.walk([&](pim::PimCoreBatchOp coreBatchOp) {
+    StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, 0);
+    verifyWithKnowledge(coreBatchOp, knowledge);
   });
 
   if (hasFailure) {

src/PIM/Dialect/Spatial/Spatial.td

@@ -232,6 +232,22 @@ def SpatReluPlanOp : SpatOp<"relu_plan", []> {
   let hasVerifier = 1;
 }
 
+def SpatBiasAddPlanOp : SpatOp<"bias_add_plan", []> {
+  let summary = "Layout-aware Conv-style bias add planning op";
+
+  let arguments = (ins
+    SpatTensor:$input,
+    SpatTensor:$bias,
+    StrAttr:$logicalLayout
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+}
+
 def SpatBlueprintOp : SpatOp<"blueprint", []> {
   let summary = "Blueprint for assembling logical tensors from published fragments";
 

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -1,5 +1,6 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Block.h"
@@ -59,6 +60,21 @@ static LogicalResult verifyStaticWeights(ComputeOpTy computeOp, StringRef kind)
   return success();
 }
 
+static bool isStaticScfForInductionVar(Value value) {
+  auto blockArg = dyn_cast<BlockArgument>(value);
+  if (!blockArg)
+    return false;
+
+  auto loop = dyn_cast_or_null<scf::ForOp>(blockArg.getOwner()->getParentOp());
+  if (!loop || loop.getInductionVar() != value)
+    return false;
+
+  std::optional<int64_t> lowerBound = matchConstantIndexValue(loop.getLowerBound());
+  std::optional<int64_t> upperBound = matchConstantIndexValue(loop.getUpperBound());
+  std::optional<int64_t> step = matchConstantIndexValue(loop.getStep());
+  return lowerBound && upperBound && step && *step > 0 && *upperBound >= *lowerBound;
+}
+
 static bool isStaticIndexExpr(Value value) {
   if (matchConstantIndexValue(value))
     return true;
@@ -80,7 +96,7 @@ static bool isStaticIndexExpr(Value value) {
 }
 
 static bool isSupportedLaneOffsetExpr(Value value, BlockArgument laneArg) {
-  if (value == laneArg || isStaticIndexExpr(value))
+  if (value == laneArg || isStaticIndexExpr(value) || isStaticScfForInductionVar(value))
     return true;
 
   auto affineApply = value.getDefiningOp<affine::AffineApplyOp>();
@@ -436,6 +452,39 @@ LogicalResult SpatReluPlanOp::verify() {
   return success();
 }
 
+LogicalResult SpatBiasAddPlanOp::verify() {
+  if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.bias_add_plan")))
+    return failure();
+  if (!isKnownLogicalLayout(getLogicalLayout()))
+    return emitError("requires a known logical layout");
+
+  auto inputType = dyn_cast<RankedTensorType>(getInput().getType());
+  auto biasType = dyn_cast<RankedTensorType>(getBias().getType());
+  auto outputType = dyn_cast<RankedTensorType>(getOutput().getType());
+  if (!inputType || !biasType || !outputType)
+    return emitError("requires ranked tensor input, bias, and output");
+  if (!inputType.hasStaticShape() || !biasType.hasStaticShape() || !outputType.hasStaticShape())
+    return emitError("requires static tensor input, bias, and output");
+  if (inputType != outputType)
+    return emitError("requires matching input and output tensor types");
+  if (outputType.getRank() != 4)
+    return emitError("requires rank-4 input/output tensors");
+  if (getLogicalLayout() != "nchw")
+    return emitError("requires logical layout \"nchw\"");
+  if (biasType.getElementType() != outputType.getElementType())
+    return emitError("requires bias element type to match the output element type");
+
+  ArrayRef<int64_t> biasShape = biasType.getShape();
+  const int64_t channels = outputType.getDimSize(1);
+  const bool supported = biasShape.empty() || (biasShape.size() == 1 && biasShape[0] == channels)
+                         || (biasShape.size() == 2 && biasShape[0] == 1 && biasShape[1] == channels)
+                         || (biasShape.size() == 4 && biasShape[0] == 1 && biasShape[1] == channels
+                             && biasShape[2] == 1 && biasShape[3] == 1);
+  if (!supported)
+    return emitError("requires scalar or per-channel bias broadcastable over NCHW");
+  return success();
+}
+
 LogicalResult SpatBlueprintOp::verify() {
   auto modeAttr = getModeAttr();
   bool isFragmentAssembly = modeAttr && modeAttr.getValue() == "fragment_assembly";

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp

@@ -121,6 +121,51 @@ struct CompactRunPlan {
   llvm::SmallVector<RunOutputDemand, 4> outputs;
 };
 
+struct ScalarPeerEdgeKey {
+  int64_t sourceCore = 0;
+  int64_t targetCore = 0;
+  mlir::Type payloadType;
+};
+
+struct ScalarPeerReceiveKey {
+  int64_t sourceCore = 0;
+  int64_t targetCore = 0;
+  mlir::Type payloadType;
+  std::optional<int64_t> channelId;
+};
+
+struct PendingScalarSend {
+  ClassId sourceClass = 0;
+  int64_t sourceCore = 0;
+  int64_t targetCore = 0;
+  mlir::Type payloadType;
+  ScalarPeerEdgeKey waitForReceive;
+  mlir::Operation* payloadAnchor = nullptr;
+  mlir::Value payload;
+  MessageVector messages;
+  mlir::Location loc;
+};
+
+struct PendingProjectedScalarSend {
+  ClassId sourceClass = 0;
+  int64_t sourceCore = 0;
+  int64_t targetCore = 0;
+  mlir::Type payloadType;
+  ScalarPeerEdgeKey waitForReceive;
+  mlir::Operation* payloadAnchor = nullptr;
+  mlir::Value payload;
+  MessageVector messages;
+  ProjectedTransferDescriptor descriptor;
+  mlir::Location loc;
+};
+
+struct PendingProjectedInputSend {
+  ClassId sourceClass = 0;
+  mlir::Value payload;
+  llvm::SmallVector<ProjectedInputTransferFragment*, 4> fragments;
+  mlir::Location loc;
+};
+
 enum class BatchInputDemandKind {
   LaneFragment,
   ProjectedFragment,
@@ -235,10 +280,19 @@ struct MaterializerState {
     projectedTransfers;
   llvm::DenseMap<mlir::Operation*, llvm::DenseMap<ClassId, ProjectedExtractReplacement>>
     projectedExtractReplacements;
+  llvm::DenseMap<mlir::Operation*, llvm::DenseMap<ClassId, ProjectedInputTransferPlan>>
+    projectedInputTransferPlans;
   AvailableValueStore availableValues;
   llvm::DenseMap<mlir::Value, mlir::Value> hostReplacements;
   llvm::DenseMap<mlir::Value, ClassId> hostOutputOwners;
   llvm::SmallVector<PendingProjectedHostOutputFragment, 32> pendingProjectedHostOutputFragments;
+  llvm::SmallVector<ScalarPeerEdgeKey, 16> plannedScalarPeerReceives;
+  llvm::SmallVector<ScalarPeerEdgeKey, 8> materializedScalarPeerReceives;
+  llvm::SmallVector<PendingScalarSend, 8> pendingScalarSends;
+  llvm::SmallVector<PendingProjectedScalarSend, 8> pendingProjectedScalarSends;
+  llvm::SmallVector<PendingProjectedInputSend, 16> pendingProjectedInputSends;
+  llvm::DenseSet<ClassId> projectedInputPhaseBarrierClasses;
+  llvm::DenseMap<ClassId, unsigned> pendingProjectedHighToLowReceives;
   llvm::DenseSet<mlir::Operation*> oldComputeOps;
 
   MaterializerState(mlir::func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId)

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ProjectedFragments.hpp

@@ -47,6 +47,24 @@ struct ProjectedExtractReplacement {
   ProjectedFragmentLayout layout;
 };
 
+struct ProjectedInputTransferFragment {
+  ProducerKey producer;
+  llvm::SmallVector<int64_t, 4> fragmentOffsets;
+  unsigned targetLane = 0;
+  unsigned ordinal = 0;
+  int64_t channelId = 0;
+  int32_t sourceCoreId = 0;
+  int32_t targetCoreId = 0;
+  bool sendEmitted = false;
+};
+
+struct ProjectedInputTransferPlan {
+  ProjectedBatchInputKey inputKey;
+  mlir::Operation* extractOp = nullptr;
+  ProjectedFragmentLayout layout;
+  llvm::SmallVector<ProjectedInputTransferFragment, 8> fragments;
+};
+
 struct PendingProjectedHostOutputFragment {
   mlir::Value originalOutput;
   ClassId sourceClass = 0;

src/PIM/Pass/PIMPasses.h

@@ -11,8 +11,6 @@ std::unique_ptr<mlir::Pass> createONNXToSpatialPass();
 std::unique_ptr<mlir::Pass> createSpatialLayoutPlanningPass();
 std::unique_ptr<mlir::Pass> createLowerSpatialPlansPass();
 
-std::unique_ptr<mlir::Pass> createSpatialToGraphvizPass();
-
 std::unique_ptr<mlir::Pass> createSpatialToPimPass();
 
 std::unique_ptr<mlir::Pass> createPimBufferizationPass();

src/PIM/PimAccelerator.cpp

@@ -74,7 +74,6 @@ void PimAccelerator::registerPasses(int optLevel) const {
   registerPass(createONNXToSpatialPass);
   registerPass(createSpatialLayoutPlanningPass);
   registerPass(createLowerSpatialPlansPass);
-  registerPass(createSpatialToGraphvizPass);
   registerPass(createSpatialToPimPass);
   registerPass(createPimBufferizationPass);
   registerPass(createPimMemoryCoalescingPass);