second temp commit: i will soft-reset and recommit after next changes

README.md

@@ -52,12 +52,22 @@ ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts
    `Patterns/{Math,NN,Tensor}` and currently cover Conv, Gemm, MatMul,
    elementwise Add/Mul/Div, ReduceMean, pooling, Relu, Sigmoid, Softmax,
    Concat, Gather, Reshape, Resize, and Split.
+   The compiler-layer target adapter supplies the target-neutral
+   `SpatialTargetInfo`. Layout-aware plan ops advertise typed alternatives
+   through the Spatial layout interface; the layout planner records the
+   selected layout and explicit materialization edges. `LowerSpatialPlans`
+   then pattern-lowers those selected plans. Contraction and Conv lowering
+   keep semantic problems, target-dependent plans, and IR materializers in
+   separate layers.
 
-2. **Merge compute nodes**
+2. **Merge, schedule, and realize Spatial communication**
    (`src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes`).
-   Builds a compute graph, schedules it with the PEFT scheduler, and materializes
+   `TrivialGraphComputeMerge` performs local graph merging, then
-   the merge schedule into Spatial IR. Supporting scheduling code lives under
+   `ScheduleSpatialGraph` materializes scheduled computes and explicit deferred
-   `MergeComputeNodes/Scheduling`.
+   communication. `VerifyScheduledSpatial` checks that intermediate contract;
+   `RealizeSpatialCommunication` resolves transfers and forwarding; and
+   `VerifyRealizedSpatial` checks the final scheduled graph. Supporting
+   scheduling code lives under `MergeComputeNodes/Scheduling`.
 
 3. **Spatial -> Pim** (`src/PIM/Conversion/SpatialToPim`).
    Lowers Spatial operations to the `pim` dialect (`src/PIM/Dialect/Pim`),
@@ -65,8 +75,12 @@ ONNX-MLIR -> Spatial -> Pim (tensor) -> Pim (bufferized) -> PIM artifacts
    tensor materialization, and return-path normalization.
 
 4. **Bufferization** (`src/PIM/Dialect/Pim/Transforms/Bufferization`).
-   Converts tensor-semantics PIM IR into memref-semantics PIM IR using MLIR's
+   `PimBufferizationPreparation` establishes writable destinations without
-   bufferization interfaces.
+   duplicating the one-shot copy analysis, `PimOneShotBufferization` runs
+   MLIR's one-shot analysis,
+   `PimMemoryNormalization` forwards/removes redundant copies and normalizes
+   addressable accesses, and `PimBufferizationVerification` checks tensor
+   absence, contiguity, and copy address spaces.
 
 5. **PIM local-memory planning**
    (`src/PIM/Dialect/Pim/Transforms/LocalMemoryPlanning`).

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -15,6 +15,8 @@
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp"
 #include "src/Accelerators/PIM/Pass/PIMPasses.h"
 #include "src/Compiler/CompilerPasses.hpp"
@@ -80,6 +82,34 @@ spatial::SchedulingTarget getDefaultPimSchedulingTarget() {
   return target;
 }
 
+spatial::ConvLoweringStrategy getSpatialConvLoweringStrategy(PimConvLoweringType strategy) {
+  switch (strategy) {
+  case PimConvLoweringAuto: return spatial::ConvLoweringStrategy::Auto;
+  case PimConvLoweringLegacy: return spatial::ConvLoweringStrategy::Legacy;
+  case PimConvLoweringDepthwise: return spatial::ConvLoweringStrategy::Depthwise;
+  case PimConvLoweringPackedIm2Col: return spatial::ConvLoweringStrategy::PackedIm2Col;
+  case PimConvLoweringStreamedPatch: return spatial::ConvLoweringStrategy::StreamedPatch;
+  case PimConvLoweringStreamedPacked: return spatial::ConvLoweringStrategy::StreamedPacked;
+  case PimConvLoweringOutputChannelTiled: return spatial::ConvLoweringStrategy::OutputChannelTiled;
+  case PimConvLoweringInputKTiled: return spatial::ConvLoweringStrategy::InputKTiled;
+  case PimConvLoweringTiled2D: return spatial::ConvLoweringStrategy::Tiled2D;
+  }
+  llvm_unreachable("unknown PIM Conv lowering strategy");
+}
+
+spatial::SpatialTargetInfo getPimSpatialTargetInfo(const spatial::SchedulingTarget& target) {
+  spatial::SpatialTargetInfo info;
+  info.matrixShape = {target.matrixRows, target.matrixColumns};
+  info.matrixUnitsPerProcessor = target.residentWeightCapacity;
+  info.processorCount = target.processorCount;
+  info.vectorWidth = target.vectorWidth;
+  info.convIm2colMaxElements = pimConvIm2colMaxElements.getValue();
+  info.convStreamChunkPositions = pimConvStreamChunkPositions.getValue();
+  info.convLoweringStrategy = getSpatialConvLoweringStrategy(pimConvLowering.getValue());
+  info.useExperimentalConvImplementation = useExperimentalConvImpl.getValue();
+  return info;
+}
+
 const llvm::json::Object& requireObject(const llvm::json::Object& object,
                                         llvm::StringRef key,
                                         llvm::StringRef path) {
@@ -293,12 +323,17 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
 
   if (pimEmissionTarget >= EmitSpatial) {
     spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget();
-    pm.addPass(createONNXToSpatialPass());
+    spatial::SpatialTargetInfo targetInfo = getPimSpatialTargetInfo(schedulingTarget);
-    pm.addPass(createSpatialLayoutPlanningPass());
+    pm.addPass(createONNXToSpatialPass(targetInfo));
-    pm.addPass(createLowerSpatialPlansPass());
+    pm.addPass(createSpatialLayoutPlanningPass(targetInfo));
+    pm.addPass(createLowerSpatialPlansPass(targetInfo));
     pm.addPass(createTrivialGraphComputeMergePass(
       schedulingTarget.residentWeightCapacity));
-    pm.addPass(createMergeComputeNodesPass(schedulingTarget));
+    auto scheduledState = std::make_shared<spatial::ScheduledSpatialState>();
+    pm.addPass(spatial::createScheduleSpatialGraphPass(schedulingTarget, scheduledState));
+    pm.addPass(spatial::createVerifyScheduledSpatialPass(scheduledState));
+    pm.addPass(spatial::createRealizeSpatialCommunicationPass(schedulingTarget, scheduledState));
+    pm.addPass(spatial::createVerifyRealizedSpatialPass(scheduledState));
     pm.addPass(createMessagePass("Onnx lowered to Spatial"));
   }
 
@@ -308,7 +343,10 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
   }
 
   if (pimEmissionTarget >= EmitPimBufferized) {
-    pm.addPass(createPimBufferizationPass());
+    pm.addPass(createPimBufferizationPreparationPass());
+    pm.addPass(createPimOneShotBufferizationPass());
+    pm.addPass(createPimMemoryNormalizationPass());
+    pm.addPass(createPimBufferizationVerificationPass());
     pm.addPass(createMessagePass("Pim bufferized"));
   }
 

src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt

@@ -27,11 +27,14 @@ add_pim_library(OMONNXToSpatial
   Patterns/Tensor/Split.cpp
   Patterns/Tensor/Transpose.cpp
   ONNXToSpatialPass.cpp
+  SpatialLayoutCapabilities.cpp
   SpatialLayoutPlanningPass.cpp
   LowerSpatialPlansPass.cpp
   Common/AttributeUtils.cpp
   Common/BiasAddUtils.cpp
   Common/ComputeRegionBuilder.cpp
+  Common/ContractionMaterialization.cpp
+  Common/ContractionPlanning.cpp
   Common/MatrixProductLowering.cpp
   Common/RowStripLayoutUtils.cpp
   Common/ShapeTilingUtils.cpp
@@ -46,8 +49,6 @@ add_pim_library(OMONNXToSpatial
   MLIRLinalgDialect
   MLIRSCFDialect
   MLIRTosaDialect
-  OMCompilerOptions
-  OMPimCompilerOptions
   OMONNXOps
   SpatialOps
   OMPimCommon

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

@@ -25,6 +25,9 @@ FailureOr<Value> createFragmentAssemblyBlueprint(Value physicalBatch,
   const int64_t laneCount = physicalType.getDimSize(0);
   if (laneCount <= 0)
     return emitError(loc, "fragment assembly requires at least one physical source slot"), failure();
+  auto physicalLayoutValue = spatial::symbolizePhysicalLayout(physicalLayout);
+  if (!physicalLayoutValue)
+    return emitError(loc, "unknown physical layout for fragment assembly"), failure();
   const int64_t fragmentElements = physicalType.getNumElements() / laneCount;
   SmallVector<int64_t> operandIndices(entries.size(), 0), sourceSlots, sourceOffsets, offsets, sizes,
     strides(entries.size() * rank, 1);
@@ -48,9 +51,10 @@ FailureOr<Value> createFragmentAssemblyBlueprint(Value physicalBatch,
     llvm::append_range(sizes, entry.sizes);
   }
   auto blueprint = spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, physicalBatch, ValueRange {},
-    rewriter.getStringAttr("nchw"), rewriter.getStringAttr(physicalLayout),
+    spatial::getNCHWLayout(rewriter.getContext()),
+    spatial::PhysicalLayoutAttr::get(rewriter.getContext(), *physicalLayoutValue),
     rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes),
-    rewriter.getStringAttr(indexMap), rewriter.getStringAttr("fragment_assembly"),
+    rewriter.getStringAttr(indexMap), spatial::getFragmentAssemblyMode(rewriter.getContext()),
     rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots),
     rewriter.getDenseI64ArrayAttr(sourceOffsets), rewriter.getDenseI64ArrayAttr(strides),
     rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete"));

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

@@ -0,0 +1,39 @@
+#include "ContractionMaterialization.hpp"
+
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
+#include "MatrixProductLowering.hpp"
+
+namespace onnx_mlir {
+
+mlir::Value materializePaddedContractionInput(
+    mlir::Value input,
+    mlir::RankedTensorType paddedType,
+    mlir::PatternRewriter& rewriter,
+    mlir::Location loc) {
+  return createPaddedInputCompute(input, paddedType, rewriter, loc);
+}
+
+mlir::FailureOr<mlir::Value> materializeTransposedContractionConstant(
+    mlir::Value input,
+    mlir::RankedTensorType resultType,
+    llvm::ArrayRef<int64_t> permutation,
+    mlir::PatternRewriter& rewriter,
+    mlir::Location loc) {
+  auto denseAttr = getHostConstDenseElementsAttr(input);
+  auto inputType = denseAttr ? mlir::dyn_cast<mlir::RankedTensorType>(denseAttr.getType()) : nullptr;
+  if (!inputType || !inputType.hasStaticShape() || !resultType || !resultType.hasStaticShape()
+      || inputType.getRank() != resultType.getRank())
+    return mlir::failure();
+
+  auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation);
+  if (mlir::failed(transposedAttr) || transposedAttr->getType() != resultType)
+    return mlir::failure();
+
+  return getOrCreateConstant(rewriter,
+                             rewriter.getInsertionBlock()->getParentOp(),
+                             *transposedAttr,
+                             resultType);
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,23 @@
+#pragma once
+
+#include "llvm/ADT/ArrayRef.h"
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+
+namespace onnx_mlir {
+
+mlir::Value materializePaddedContractionInput(
+    mlir::Value input,
+    mlir::RankedTensorType paddedType,
+    mlir::PatternRewriter& rewriter,
+    mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> materializeTransposedContractionConstant(
+    mlir::Value input,
+    mlir::RankedTensorType resultType,
+    llvm::ArrayRef<int64_t> permutation,
+    mlir::PatternRewriter& rewriter,
+    mlir::Location loc);
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,68 @@
+#include "ContractionPlanning.hpp"
+
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+
+#include <algorithm>
+
+namespace onnx_mlir {
+
+namespace {
+
+static int64_t ceilDivide(int64_t value, int64_t divisor) {
+  return divisor == 0 ? 0 : (value + divisor - 1) / divisor;
+}
+
+static llvm::SmallVector<int64_t> buildBatchMap(
+    llvm::ArrayRef<int64_t> sourceShape,
+    llvm::ArrayRef<int64_t> outputShape) {
+  llvm::SmallVector<int64_t> map(outputShape.size(), -1);
+  const int64_t offset = outputShape.size() - sourceShape.size();
+  for (int64_t source = 0; source < static_cast<int64_t>(sourceShape.size()); ++source) {
+    const int64_t output = source + offset;
+    if (sourceShape[source] != 1)
+      map[output] = source;
+  }
+  return map;
+}
+
+} // namespace
+
+ContractionPlan makeContractionPlan(
+    const ContractionProblem& problem,
+    const spatial::SpatialTargetInfo& target,
+    ContractionPlanKind kind,
+    int64_t laneCount,
+    int64_t fragmentRows) {
+  ContractionPlan plan;
+  plan.problem = problem;
+  plan.kind = kind;
+  plan.tileM = std::max<int64_t>(1, target.matrixShape.rows);
+  plan.tileK = std::max<int64_t>(1, target.matrixShape.rows);
+  plan.tileN = std::max<int64_t>(1, target.matrixShape.columns);
+  plan.reductionSlices = std::max<int64_t>(1, ceilDivide(problem.k, plan.tileK));
+  plan.outputTiles = std::max<int64_t>(1, ceilDivide(problem.n, plan.tileN));
+  plan.rowTiles = std::max<int64_t>(1, ceilDivide(problem.m, plan.tileM));
+  plan.fragmentRows = std::max<int64_t>(
+      1, fragmentRows != 0 ? fragmentRows : plan.tileM);
+  plan.lhsBatchMap = buildBatchMap(problem.lhsBatchShape, problem.outputBatchShape);
+  plan.rhsBatchMap = buildBatchMap(problem.rhsBatchShape, problem.outputBatchShape);
+
+  if (laneCount != 0)
+    plan.laneCount = laneCount;
+  else if (kind == ContractionPlanKind::StaticTiled)
+    plan.laneCount = problem.batch * problem.m * plan.reductionSlices * plan.outputTiles;
+  else if (kind == ContractionPlanKind::GroupedRowDynamicVVD)
+    plan.laneCount = problem.batch * ceilDivide(problem.m, plan.fragmentRows);
+  else
+    plan.laneCount = problem.batch * problem.m * problem.n;
+
+  plan.expectedMvmCount = kind == ContractionPlanKind::StaticTiled ? plan.laneCount : 0;
+  plan.expectedVvdCount = kind == ContractionPlanKind::StaticTiled ? 0 : plan.laneCount;
+  plan.expectedVectorCount = plan.laneCount * plan.reductionSlices;
+  if (problem.resultElementType && problem.n > 0)
+    plan.physicalFragmentType = mlir::RankedTensorType::get(
+        {plan.fragmentRows, problem.n}, problem.resultElementType);
+  return plan;
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,41 @@
+#pragma once
+
+#include "ContractionProblem.hpp"
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
+
+namespace onnx_mlir {
+
+enum class ContractionPlanKind {
+  StaticTiled,
+  BatchedDynamicVVD,
+  GroupedRowDynamicVVD,
+};
+
+struct ContractionPlan {
+  ContractionProblem problem;
+  ContractionPlanKind kind = ContractionPlanKind::StaticTiled;
+  int64_t tileM = 1;
+  int64_t tileK = 1;
+  int64_t tileN = 1;
+  int64_t fragmentRows = 1;
+  int64_t reductionSlices = 1;
+  int64_t outputTiles = 1;
+  int64_t rowTiles = 1;
+  int64_t laneCount = 0;
+  int64_t expectedMvmCount = 0;
+  int64_t expectedVvdCount = 0;
+  int64_t expectedVectorCount = 0;
+  llvm::SmallVector<int64_t> lhsBatchMap;
+  llvm::SmallVector<int64_t> rhsBatchMap;
+  mlir::RankedTensorType physicalFragmentType;
+};
+
+ContractionPlan makeContractionPlan(
+    const ContractionProblem& problem,
+    const spatial::SpatialTargetInfo& target,
+    ContractionPlanKind kind,
+    int64_t laneCount = 0,
+    int64_t fragmentRows = 0);
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,35 @@
+#pragma once
+
+#include "mlir/IR/BuiltinTypes.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include <cstdint>
+
+namespace onnx_mlir {
+
+enum class ContractionOrigin { Gemm, MatMul };
+
+struct ContractionProblem {
+  llvm::SmallVector<int64_t> lhsBatchShape;
+  llvm::SmallVector<int64_t> rhsBatchShape;
+  llvm::SmallVector<int64_t> outputBatchShape;
+  int64_t lhsBatch = 1;
+  int64_t rhsBatch = 1;
+  int64_t batch = 1;
+  int64_t m = 0;
+  int64_t k = 0;
+  int64_t n = 0;
+  ContractionOrigin origin = ContractionOrigin::MatMul;
+  mlir::Type lhsElementType;
+  mlir::Type rhsElementType;
+  mlir::Type resultElementType;
+  bool lhsTransposed = false;
+  bool rhsTransposed = false;
+  bool lhsWasVector = false;
+  bool rhsWasVector = false;
+  float alpha = 1.0f;
+  float beta = 1.0f;
+};
+
+} // namespace onnx_mlir

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

@@ -1,15 +1,70 @@
 #include "MatrixProductLowering.hpp"
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
 
 namespace onnx_mlir {
 
+static bool isInsideSpatialCompute(Operation* op) {
+  for (Operation* parent = op; parent; parent = parent->getParentOp())
+    if (spatial::isAnySpatialComputeLike(parent))
+      return true;
+  return false;
+}
+
+static Value buildLinalgTranspose(Value value,
+                                  RankedTensorType resultType,
+                                  ArrayRef<int64_t> permutation,
+                                  PatternRewriter& rewriter,
+                                  Location loc) {
+  Value init = tensor::EmptyOp::create(
+      rewriter, loc, resultType.getShape(), resultType.getElementType());
+  return linalg::TransposeOp::create(
+    rewriter, loc, value, init, permutation).getResult()[0];
+}
+
+static Value materializeConstantTranspose(Value value,
+                                          RankedTensorType resultType,
+                                          ArrayRef<int64_t> permutation,
+                                          PatternRewriter& rewriter) {
+  auto denseAttr = getHostConstDenseElementsAttr(value);
+  if (!denseAttr)
+    return {};
+  auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation);
+  if (failed(transposedAttr) || transposedAttr->getType() != resultType)
+    return {};
+  return getOrCreateConstant(
+      rewriter, rewriter.getInsertionBlock()->getParentOp(), *transposedAttr, resultType);
+}
+
+Value createLinalgTranspose(Value value,
+                            RankedTensorType resultType,
+                            ArrayRef<int64_t> permutation,
+                            PatternRewriter& rewriter,
+                            Location loc) {
+  if (Value constant = materializeConstantTranspose(value, resultType, permutation, rewriter))
+    return constant;
+
+  if (isInsideSpatialCompute(rewriter.getInsertionBlock()->getParentOp()))
+    return buildLinalgTranspose(value, resultType, permutation, rewriter, loc);
+
+  auto compute = createSpatCompute<1>(
+    rewriter, loc, TypeRange {resultType}, {}, ValueRange {value},
+    [&](Value input) {
+      spatial::SpatYieldOp::create(
+        rewriter, loc, buildLinalgTranspose(input, resultType, permutation, rewriter, loc));
+    });
+  return compute.getResult(0);
+}
+
 Value createZeroPaddedTensor(Value value, RankedTensorType resultType, PatternRewriter& rewriter, Location loc) {
   auto sourceType = cast<RankedTensorType>(value.getType());
   SmallVector<OpFoldResult> lowPads(sourceType.getRank(), rewriter.getIndexAttr(0));

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

@@ -5,8 +5,16 @@
 #include "mlir/IR/Value.h"
 #include "mlir/Transforms/DialectConversion.h"
 
+#include "llvm/ADT/ArrayRef.h"
+
 namespace onnx_mlir {
 
+mlir::Value createLinalgTranspose(mlir::Value value,
+                                  mlir::RankedTensorType resultType,
+                                  llvm::ArrayRef<int64_t> permutation,
+                                  mlir::PatternRewriter& rewriter,
+                                  mlir::Location loc);
+
 mlir::Value createZeroPaddedTensor(mlir::Value value,
                                    mlir::RankedTensorType resultType,
                                    mlir::PatternRewriter& rewriter,

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

@@ -5,9 +5,9 @@
 #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/ComputeRegionBuilder.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Dialect/ONNX/ONNXOps.hpp"
 
 #include <numeric>
 
@@ -33,6 +33,16 @@ FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, Ra
     tilesPerRow};
 }
 
+FailureOr<RowStripPhysicalValue> getRowStripPhysicalValue(Value value) {
+  auto blueprint = value.getDefiningOp<spatial::SpatBlueprintOp>();
+  auto logicalType = dyn_cast<RankedTensorType>(value.getType());
+  if (!blueprint || !logicalType || blueprint.getOutput() != value
+      || blueprint.getPhysicalLayout() != spatial::PhysicalLayout::NHWCRowStrip
+      || !spatial::isPhysicalView(blueprint.getMode()))
+    return failure();
+  return describeRowStripPhysicalValue(blueprint.getInput(), logicalType);
+}
+
 RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
   return RankedTensorType::get({logicalType.getDimSize(0), 1, logicalType.getDimSize(3),
                                 logicalType.getDimSize(1)},
@@ -144,6 +154,35 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
   return batchOp->getResult(0);
 }
 
+FailureOr<Value> createRowStripStorageBlueprint(Value storage,
+                                                RankedTensorType logicalType,
+                                                PatternRewriter& rewriter,
+                                                Location loc) {
+  FailureOr<RowStripPhysicalValue> value = describeRowStripPhysicalValue(storage, logicalType);
+  if (failed(value))
+    return failure();
+
+  auto blueprint = spatial::SpatBlueprintOp::create(
+      rewriter,
+      loc,
+      logicalType,
+      storage,
+      ValueRange {},
+      spatial::getNCHWLayout(rewriter.getContext()),
+      spatial::getNHWCRowStripLayout(rewriter.getContext()),
+      rewriter.getDenseI64ArrayAttr({}),
+      rewriter.getDenseI64ArrayAttr({}),
+      rewriter.getStringAttr(kRowStripIndexMap),
+      spatial::getPhysicalViewMode(rewriter.getContext()),
+      nullptr,
+      nullptr,
+      nullptr,
+      nullptr,
+      nullptr,
+      nullptr);
+  return blueprint.getOutput();
+}
+
 FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
                                                  PatternRewriter& rewriter,
                                                  Location loc) {
@@ -160,8 +199,8 @@ FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& va
         rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
       if (failed(fragment))
         return failure();
-      Value nchw = ONNXTransposeOp::create(
+      Value nchw = createLinalgTranspose(
-        rewriter, loc, nchwFragmentType, *fragment, rewriter.getI64ArrayAttr({0, 3, 1, 2}));
+        *fragment, nchwFragmentType, {0, 3, 1, 2}, rewriter, loc);
       publishGraphBatchPhysicalFragment(rewriter, loc, nchw, args.outputs.front(), args.lane);
       return success();
     });
@@ -176,7 +215,7 @@ FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& va
         {1, std::min(tileChannels, value.logicalType.getDimSize(1) - channelOffset), 1,
          value.logicalType.getDimSize(3)}});
     }
-  return createFragmentAssemblyBlueprint(transposed->getResult(0), value.logicalType, entries, "nhwc_row_strip",
+  return createFragmentAssemblyBlueprint(transposed->getResult(0), value.logicalType, entries, "dense_nchw",
                                          kRowStripIndexMap, rewriter, loc);
 }
 

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

@@ -6,6 +6,12 @@
 
 namespace onnx_mlir {
 
+namespace spatial {
+class SpatBlueprintOp;
+class SpatGraphCompute;
+struct SpatialTargetInfo;
+} // namespace spatial
+
 inline constexpr llvm::StringLiteral kRowStripIndexMap = "nhwc_row_strip_fragments";
 
 struct RowStripPhysicalValue {
@@ -18,6 +24,8 @@ struct RowStripPhysicalValue {
 mlir::FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(mlir::Value storage,
                                                                      mlir::RankedTensorType logicalType);
 
+mlir::FailureOr<RowStripPhysicalValue> getRowStripPhysicalValue(mlir::Value value);
+
 std::pair<llvm::SmallVector<int64_t>, llvm::SmallVector<int64_t>>
 buildRowStripMetadata(mlir::RankedTensorType type);
 
@@ -53,6 +61,11 @@ mlir::FailureOr<mlir::Value> createRowStripStorageFromRows(mlir::Value rows,
                                                            mlir::PatternRewriter& rewriter,
                                                            mlir::Location loc);
 
+mlir::FailureOr<mlir::Value> createRowStripStorageBlueprint(mlir::Value storage,
+                                                            mlir::RankedTensorType logicalType,
+                                                            mlir::PatternRewriter& rewriter,
+                                                            mlir::Location loc);
+
 mlir::FailureOr<mlir::Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
                                                                mlir::PatternRewriter& rewriter,
                                                                mlir::Location loc);
@@ -80,4 +93,14 @@ mlir::FailureOr<mlir::Value> applyRowStripConcat(llvm::ArrayRef<RowStripPhysical
                                                  mlir::PatternRewriter& rewriter,
                                                  mlir::Location loc);
 
+mlir::LogicalResult canLowerFlattenFromRowStrip(
+    spatial::SpatGraphCompute flattenOp,
+    const spatial::SpatialTargetInfo& target);
+
+mlir::LogicalResult lowerFlattenFromRowStrip(
+    const RowStripPhysicalValue& input,
+    spatial::SpatGraphCompute flattenOp,
+    const spatial::SpatialTargetInfo& target,
+    mlir::PatternRewriter& rewriter);
+
 } // namespace onnx_mlir

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

@@ -5,7 +5,6 @@
 
 #include "ShapeTilingUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 
 using namespace mlir;
@@ -67,11 +66,15 @@ sliceVector(const Value& vectorToSlice, int64_t sliceSize, PatternRewriter& rewr
 }
 
 DenseMap<CoreId, SmallVector<Value>>
-sliceVectorPerCrossbarPerCore(const Value& vectorToSlice, PatternRewriter& rewriter, Location loc) {
+sliceVectorPerCrossbarPerCore(const Value& vectorToSlice,
-  SmallVector<Value> slices = sliceVector(vectorToSlice, crossbarSize, rewriter, loc);
+                              PatternRewriter& rewriter,
+                              Location loc,
+                              const spatial::SpatialTargetInfo& target) {
+  SmallVector<Value> slices = sliceVector(
+    vectorToSlice, static_cast<int64_t>(target.matrixShape.rows), rewriter, loc);
   DenseMap<CoreId, SmallVector<Value>> slicesPerCore;
   for (size_t sliceId = 0; sliceId < slices.size(); sliceId++) {
-    size_t coreId = sliceId / crossbarCountInCore;
+    size_t coreId = sliceId / target.matrixUnitsPerProcessor;
     slicesPerCore[coreId].push_back(slices[sliceId]);
   }
   return slicesPerCore;

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

@@ -7,6 +7,7 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
 
 namespace onnx_mlir {
 
@@ -26,6 +27,9 @@ llvm::SmallVector<mlir::Value> sliceVector(const mlir::Value& vectorToSlice,
 /// Partitions one logical vector into per-core crossbar-sized slices using the
 /// current PIM target geometry.
 llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>> sliceVectorPerCrossbarPerCore(
-  const mlir::Value& vectorToSlice, mlir::PatternRewriter& rewriter, mlir::Location loc);
+  const mlir::Value& vectorToSlice,
+  mlir::PatternRewriter& rewriter,
+  mlir::Location loc,
+  const spatial::SpatialTargetInfo& target);
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -34,9 +34,15 @@ struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp
   StringRef getDescription() const override { return "Lower ONNX ops to Spatial ops."; }
 
   ONNXToSpatialPass() = default;
-  ONNXToSpatialPass(const ONNXToSpatialPass& pass) {}
+  explicit ONNXToSpatialPass(const spatial::SpatialTargetInfo& target)
+      : target(target), hasTarget(true) {}
+  ONNXToSpatialPass(const ONNXToSpatialPass& pass)
+      : target(pass.target), hasTarget(pass.hasTarget) {}
 
   void runOnOperation() override;
+
+  spatial::SpatialTargetInfo target;
+  bool hasTarget = false;
 };
 
 } // namespace
@@ -106,6 +112,11 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
 
 void ONNXToSpatialPass::runOnOperation() {
   ModuleOp moduleOp = getOperation();
+  if (!hasTarget) {
+    moduleOp.emitError("ONNX-to-Spatial lowering requires an injected SpatialTargetInfo");
+    signalPassFailure();
+    return;
+  }
   MLIRContext* ctx = &getContext();
 
   ConversionTarget preTarget(*ctx);
@@ -127,7 +138,7 @@ void ONNXToSpatialPass::runOnOperation() {
   }
 
   RewritePatternSet matmulPatterns(ctx);
-  populateMatMulFusionPatterns(matmulPatterns, ctx);
+  populateMatMulFusionPatterns(matmulPatterns, ctx, target);
   if (failed(applyPatternsGreedily(moduleOp, std::move(matmulPatterns)))) {
     moduleOp.emitError("failed to lower MatMul before producer conversion");
     signalPassFailure();
@@ -182,7 +193,7 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXSplitOp>();
 
   RewritePatternSet conversionPatterns(ctx);
-  populateConversionPatterns(conversionPatterns, ctx);
+  populateConversionPatterns(conversionPatterns, ctx, this->target);
   if (failed(applyPartialConversion(moduleOp, target, std::move(conversionPatterns)))) {
     moduleOp.emitError("failed to convert required ONNX ops to Spatial ops");
     signalPassFailure();
@@ -258,4 +269,8 @@ void ONNXToSpatialPass::runOnOperation() {
 
 std::unique_ptr<Pass> createONNXToSpatialPass() { return std::make_unique<ONNXToSpatialPass>(); }
 
+std::unique_ptr<Pass> createONNXToSpatialPass(const spatial::SpatialTargetInfo& target) {
+  return std::make_unique<ONNXToSpatialPass>(target);
+}
+
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp

@@ -130,8 +130,7 @@ template <typename ComputeOpTy>
 void verifyNoNestedFragmentAssemblyBlueprints(ComputeOpTy compute,
                                                   pim::CappedDiagnosticReporter& diagnostics) {
   compute.getBody().walk([&](spatial::SpatBlueprintOp blueprint) {
-    std::optional<StringRef> mode = blueprint.getMode();
+    if (!spatial::isFragmentAssembly(blueprint.getMode()))
-    if (!mode || *mode != "fragment_assembly")
       return;
     diagnostics.report(blueprint.getOperation(), [&](Operation* illegalOp) {
       illegalOp->emitOpError("fragment assembly blueprint must be host-level after merge materialization");

src/PIM/Conversion/ONNXToSpatial/Patterns.cpp

@@ -7,12 +7,14 @@ namespace onnx_mlir {
 
 void populatePrePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { populateGeneratedPrePatterns(patterns, ctx); }
 
-void populateConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+void populateConversionPatterns(RewritePatternSet& patterns,
+                                MLIRContext* ctx,
+                                const spatial::SpatialTargetInfo& target) {
   populateElementwisePatterns(patterns, ctx);
-  populateMatMulRewritePatterns(patterns, ctx);
+  populateMatMulRewritePatterns(patterns, ctx, target);
-  populateGemmPatterns(patterns, ctx);
+  populateGemmPatterns(patterns, ctx, target);
-  populateConvPatterns(patterns, ctx);
+  populateConvPatterns(patterns, ctx, target);
-  populatePoolPatterns(patterns, ctx);
+  populatePoolPatterns(patterns, ctx, target);
   populateReduceMeanPatterns(patterns, ctx);
   populateReluPatterns(patterns, ctx);
   populateSigmoidPatterns(patterns, ctx);

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -8,20 +8,36 @@
 
 namespace onnx_mlir {
 
+namespace spatial {
+struct SpatialTargetInfo;
+}
+
 void populatePrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-void populateConversionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateConversionPatterns(mlir::RewritePatternSet& patterns,
+                                mlir::MLIRContext* ctx,
+                                const spatial::SpatialTargetInfo& target);
 void populatePostPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
 void populateGeneratedPrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateWeightPromotionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
-void populateConvPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateConvPatterns(mlir::RewritePatternSet& patterns,
+                          mlir::MLIRContext* ctx,
+                          const spatial::SpatialTargetInfo& target);
 void populateElementwisePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateElementwiseFusionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
-void populateGemmPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateGemmPatterns(mlir::RewritePatternSet& patterns,
-void populateMatMulRewritePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+                          mlir::MLIRContext* ctx,
-void populateMatMulFusionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+                          const spatial::SpatialTargetInfo& target);
-void populatePoolPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateMatMulRewritePatterns(mlir::RewritePatternSet& patterns,
+                                   mlir::MLIRContext* ctx,
+                                   const spatial::SpatialTargetInfo& target);
+void populateMatMulFusionPatterns(mlir::RewritePatternSet& patterns,
+                                  mlir::MLIRContext* ctx,
+                                  const spatial::SpatialTargetInfo& target);
+void populatePoolPatterns(mlir::RewritePatternSet& patterns,
+                          mlir::MLIRContext* ctx,
+                          const spatial::SpatialTargetInfo& target);
 void populateReduceMeanPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateReluPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateSigmoidPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);

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

@@ -14,10 +14,11 @@
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
-#include "src/Accelerators/PIM/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/MatrixProductLowering.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.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"
@@ -30,11 +31,14 @@ namespace onnx_mlir {
 namespace {
 
 struct ConvToGemm : OpConversionPattern<ONNXConvOp> {
-  using OpConversionPattern::OpConversionPattern;
+  explicit ConvToGemm(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
+      : OpConversionPattern<ONNXConvOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(ONNXConvOp convOp,
                                 ONNXConvOpAdaptor convOpAdaptor,
                                 ConversionPatternRewriter& rewriter) const override;
+
+  const spatial::SpatialTargetInfo& target;
 };
 
 struct PreparedConvInput {
@@ -43,45 +47,21 @@ struct PreparedConvInput {
 };
 
 static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter& rewriter);
-static StringRef stringifyConvLoweringStrategy(PimConvLoweringType strategy) {
+static StringRef stringifyConvLoweringStrategy(spatial::ConvLoweringStrategy strategy) {
   switch (strategy) {
-  case PimConvLoweringAuto: return "auto";
+  case spatial::ConvLoweringStrategy::Auto: return "auto";
-  case PimConvLoweringLegacy: return "legacy";
+  case spatial::ConvLoweringStrategy::Legacy: return "legacy";
-  case PimConvLoweringDepthwise: return "depthwise";
+  case spatial::ConvLoweringStrategy::Depthwise: return "depthwise";
-  case PimConvLoweringPackedIm2Col: return "packed-im2col";
+  case spatial::ConvLoweringStrategy::PackedIm2Col: return "packed-im2col";
-  case PimConvLoweringStreamedPatch: return "streamed-patch";
+  case spatial::ConvLoweringStrategy::StreamedPatch: return "streamed-patch";
-  case PimConvLoweringStreamedPacked: return "streamed-packed";
+  case spatial::ConvLoweringStrategy::StreamedPacked: return "streamed-packed";
-  case PimConvLoweringOutputChannelTiled: return "output-channel-tiled";
+  case spatial::ConvLoweringStrategy::OutputChannelTiled: return "output-channel-tiled";
-  case PimConvLoweringInputKTiled: return "input-k-tiled";
+  case spatial::ConvLoweringStrategy::InputKTiled: return "input-k-tiled";
-  case PimConvLoweringTiled2D: return "tiled-2d";
+  case spatial::ConvLoweringStrategy::Tiled2D: return "tiled-2d";
   }
   llvm_unreachable("unknown conv lowering strategy");
 }
 
-static PimConvLoweringType chooseConvLoweringStrategy(const ConvGeometry& geo,
-                                                      PimConvLoweringType requested) {
-  if (requested != PimConvLoweringAuto)
-    return requested;
-
-  // Transform-based convolution is intentionally not selected for this ISA:
-  // it would require explicit transform sequences and staging traffic on top of
-  // the same crossbar MVM primitive, which is not attractive here.
-  if (geo.isDepthwise)
-    return PimConvLoweringDepthwise;
-  if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize && geo.pack >= 2 && geo.im2colElements <= pimConvIm2colMaxElements)
-    return PimConvLoweringPackedIm2Col;
-  if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize && geo.pack >= 2 && geo.im2colElements > pimConvIm2colMaxElements)
-    return PimConvLoweringStreamedPacked;
-  if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize)
-    return PimConvLoweringStreamedPatch;
-  if (geo.k <= geo.xbarSize && geo.c > geo.xbarSize)
-    return PimConvLoweringOutputChannelTiled;
-  if (geo.k > geo.xbarSize && geo.c <= geo.xbarSize)
-    return PimConvLoweringLegacy;
-  return PimConvLoweringTiled2D;
-}
-
-
 static Value expandBiasIfNeeded(Value bias, PatternRewriter& rewriter, Location loc) {
   auto biasType = cast<RankedTensorType>(bias.getType());
   if (biasType.getRank() != 1)
@@ -194,7 +174,11 @@ static Value createCollectedConvOutput(ValueRange gemmRows,
                                        int64_t packFactor,
                                        PatternRewriter& rewriter,
                                        Location loc);
-static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp, Value x, Value w, Value b);
+static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp,
+                                                             Value x,
+                                                             Value w,
+                                                             Value b,
+                                                             const spatial::SpatialTargetInfo& target);
 
 namespace depthwise {
 
@@ -215,10 +199,10 @@ static std::optional<Tiling> computeTiling(int64_t batchSize,
                                            int64_t wHeight,
                                            int64_t wWidth,
                                            int64_t outHeight,
-                                           int64_t outWidth) {
+                                           int64_t outWidth,
+                                           int64_t xbarDim) {
   const int64_t kernelElements = wHeight * wWidth;
   const int64_t outputMultiplier = numChannelsOut / numChannelsIn;
-  const int64_t xbarDim = static_cast<int64_t>(crossbarSize.getValue());
   if (kernelElements <= 0 || outputMultiplier <= 0 || kernelElements > xbarDim || outputMultiplier > xbarDim)
     return std::nullopt;
 
@@ -249,8 +233,9 @@ static Value buildPackedWeights(DenseElementsAttr wDenseAttr,
                                 const Tiling& tiling,
                                 PatternRewriter& rewriter,
                                 Location loc,
+                                int64_t xbarDim,
                                 int64_t paddedInputRows = -1) {
-  const int64_t paddedOutputChannels = static_cast<int64_t>(crossbarSize.getValue());
+  const int64_t paddedOutputChannels = xbarDim;
   const int64_t packedInputRows = paddedInputRows > 0 ? paddedInputRows : tiling.tileInputRows;
   auto packedWeightType = RankedTensorType::get(
     {tiling.numChannelTiles, packedInputRows, paddedOutputChannels}, wType.getElementType());
@@ -396,7 +381,7 @@ static Value createWeightTile(Value packedWeights,
                               PatternRewriter& rewriter,
                               Location loc) {
   SmallVector<OpFoldResult> offsets {channelTileIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-  const int64_t paddedOutputChannels = static_cast<int64_t>(crossbarSize.getValue());
+  const int64_t paddedOutputChannels = packedWeightType.getDimSize(2);
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
                                    rewriter.getIndexAttr(tiling.tileInputRows),
                                    rewriter.getIndexAttr(paddedOutputChannels)};
@@ -528,7 +513,8 @@ static bool canUseStructuredRewrite(const ConvLoweringState& state) {
                               state.wHeight,
                               state.wWidth,
                               state.outHeight,
-                              state.outWidth);
+                              state.outWidth,
+                              state.targetInfo().matrixShape.rows);
   if (!tiling)
     return false;
 
@@ -559,7 +545,8 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
                               state.wType.getDimSize(2),
                               state.wType.getDimSize(3),
                               state.outType.getDimSize(2),
-                              state.outType.getDimSize(3));
+                              state.outType.getDimSize(3),
+                              state.targetInfo().matrixShape.rows);
   if (!tiling) {
     convOp->emitOpError("failed to derive a structured depthwise tiling that fits Spatial weighted VMM lowering");
     return failure();
@@ -579,9 +566,10 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
                                                       paddedInputType.getDimSize(3),
                                                       paddedInputType.getDimSize(1)},
                                                      paddedInputType.getElementType());
-  Value channelLastInput = ONNXTransposeOp::create(
+  Value channelLastInput = createLinalgTranspose(
-    rewriter, loc, channelLastInputType, paddedInput, rewriter.getI64ArrayAttr({0, 2, 3, 1}));
+    paddedInput, channelLastInputType, {0, 2, 3, 1}, rewriter, loc);
-  Value packedWeights = buildPackedWeights(wDenseAttr, state.wType, *tiling, rewriter, loc);
+  Value packedWeights = buildPackedWeights(
+    wDenseAttr, state.wType, *tiling, rewriter, loc, state.targetInfo().matrixShape.rows);
 
   Value expandedBias;
   SmallVector<Value> batchInputs {channelLastInput};
@@ -600,7 +588,7 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
     RankedTensorType::get({tiling->totalPatches, state.outType.getDimSize(1)}, state.outType.getElementType());
   auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.outType.getElementType());
   auto paddedRowTileType = RankedTensorType::get(
-    {1, static_cast<int64_t>(crossbarSize.getValue())}, state.outType.getElementType());
+    {1, static_cast<int64_t>(state.targetInfo().matrixShape.rows)}, state.outType.getElementType());
   auto piecesType = spatial::getGraphBatchPhysicalResultType(
     tiling->totalPatches * tiling->numChannelTiles, rowTileType);
   auto inputTileType =
@@ -826,8 +814,7 @@ static Value createWeightMatrix(
     });
     if (!transpose)
       return flattened;
-    return ONNXTransposeOp::create(rewriter, loc, plan.wTransType, flattened, rewriter.getI64ArrayAttr({1, 0}))
+    return createLinalgTranspose(flattened, plan.wTransType, {1, 0}, rewriter, loc);
-      .getResult();
   };
 
   if (isCompileTimeComputable(weights))
@@ -963,17 +950,17 @@ static FailureOr<Value> rewriteInputKTiledConv(const ConvLoweringState& state,
                                                PatternRewriter& rewriter,
                                                Location loc) {
   PreparedConvInput preparedInput = prepareInputForIm2Col(state, rewriter, loc);
-  ConvGeometry geo = buildConvGeometry(state);
+  ConvGeometry geo = buildConvGeometry(state, state.targetInfo());
   const int64_t xbarDim = geo.xbarSize;
   const int64_t numKSlices = ceilIntegerDivide(geo.k, xbarDim);
   const int64_t paddedK = numKSlices * xbarDim;
   const uint64_t maxLanesPerBatch =
     std::max<uint64_t>(1,
-                       static_cast<uint64_t>(crossbarCountInCore.getValue())
+                       static_cast<uint64_t>(state.targetInfo().matrixUnitsPerProcessor)
                          / static_cast<uint64_t>(std::max<int64_t>(1, numKSlices * 4)));
   const uint64_t rowChunkWidth = std::max<uint64_t>(
     1,
-    std::min<uint64_t>({chooseStreamChunkPositions(geo, /*packFactor=*/1),
+    std::min<uint64_t>({chooseStreamChunkPositions(geo, /*packFactor=*/1, state.targetInfo()),
                         maxLanesPerBatch,
                         static_cast<uint64_t>(state.outWidth)}));
   const auto elementType = state.outType.getElementType();
@@ -1227,7 +1214,7 @@ buildConvGemmPlan(const ConvLoweringState& state,
   const int64_t wMaxDim = std::max(plan.patchSize, state.numChannelsOut);
   plan.maxParallelPixels = forcedPackFactor
                            ? *forcedPackFactor
-                           : std::max<int64_t>(1, static_cast<int64_t>(crossbarSize.getValue()) / wMaxDim);
+                           : std::max<int64_t>(1, static_cast<int64_t>(state.targetInfo().matrixShape.rows) / wMaxDim);
   plan.effectiveMaxParallelPixels =
     (canPackWeightsAsConstants && canPackBiasAsConstants) ? plan.maxParallelPixels : 1;
   plan.packedNumRows = ceilIntegerDivide(plan.chunkNumPatches, plan.effectiveMaxParallelPixels);
@@ -1251,7 +1238,8 @@ static Value createIm2colRows(const ConvLoweringState& state,
                               const ConvGemmPlan& plan,
                               PatternRewriter& rewriter,
                               Location loc) {
-  if (plan.gemmInputRowsType.getDimSize(1) > crossbarSize.getValue()) {
+  if (plan.gemmInputRowsType.getDimSize(1)
+      > static_cast<int64_t>(state.targetInfo().matrixShape.rows)) {
     assert(plan.effectiveMaxParallelPixels == 1 && "multi-crossbar im2col rows cannot pack pixels");
     auto compute = createSpatCompute<1>(
       rewriter, loc, TypeRange {plan.gemmInputRowsType}, {}, preparedInput.value, [&](Value input) {
@@ -1419,15 +1407,15 @@ static Value maybeUnpackChunkRows(Value gemmRows,
   return unpackCompute.getResult(0);
 }
 
-static Value createStreamedConvRows(const ConvLoweringState& state,
+static FailureOr<Value> createStreamedConvRows(const ConvLoweringState& state,
-                                    const PreparedConvInput& preparedInput,
+                                               const PreparedConvInput& preparedInput,
-                                    Value weightMatrix,
+                                               Value weightMatrix,
-                                    Value biasMatrix,
+                                               Value biasMatrix,
-                                    DenseElementsAttr wDenseAttr,
+                                               DenseElementsAttr wDenseAttr,
-                                    DenseElementsAttr biasDenseAttr,
+                                               DenseElementsAttr biasDenseAttr,
-                                    int64_t forcedPackFactor,
+                                               int64_t forcedPackFactor,
-                                    PatternRewriter& rewriter,
+                                               PatternRewriter& rewriter,
-                                    Location loc) {
+                                               Location loc) {
   const int64_t totalPatches = state.batchSize * state.outHeight * state.outWidth;
   ConvGemmPlan plan = buildConvGemmPlan(state, static_cast<bool>(wDenseAttr),
     !state.hasBias || static_cast<bool>(biasDenseAttr), 0, totalPatches, forcedPackFactor);
@@ -1435,14 +1423,18 @@ static Value createStreamedConvRows(const ConvLoweringState& state,
   Value packedWeights = buildPackedWeights(wDenseAttr, weightMatrix, state, plan, rewriter, loc);
   Value gemmBias = state.hasBias ? state.b : createZeroGemmBias(plan.gemmOutputRowsType, rewriter);
   Value packedBias = buildPackedBias(gemmBias, biasMatrix, biasDenseAttr, state, plan, rewriter, loc);
-  Value gemmRows = ONNXGemmOp::create(rewriter, loc, plan.gemmOutputRowsType, inputRows,
+  FailureOr<Value> gemmRows = lowerGemmToSpatial(
-    packedWeights, packedBias, APFloat(1.0f), APFloat(1.0f), 0, !wDenseAttr).getY();
+    state.diagnosticAnchor, inputRows, packedWeights, packedBias,
-  return maybeUnpackChunkRows(gemmRows, plan, rewriter, loc);
+    plan.gemmOutputRowsType, /*transA=*/false, /*transB=*/!wDenseAttr,
+    /*alpha=*/1.0f, /*beta=*/1.0f, state.targetInfo(), rewriter, loc);
+  if (failed(gemmRows))
+    return failure();
+  return maybeUnpackChunkRows(*gemmRows, plan, rewriter, loc);
 }
 
-static Value rewritePackedIm2ColConv(const ConvLoweringState& state,
+static FailureOr<Value> rewritePackedIm2ColConv(const ConvLoweringState& state,
-                                     PatternRewriter& rewriter,
+                                                PatternRewriter& rewriter,
-                                     Location loc) {
+                                                Location loc) {
   auto wDenseAttr = getHostConstDenseElementsAttr(state.w);
   PreparedConvInput preparedInput = prepareInputForIm2Col(state, rewriter, loc);
   Value biasMatrix;
@@ -1466,19 +1458,14 @@ static Value rewritePackedIm2ColConv(const ConvLoweringState& state,
     gemmBias = state.b;
   Value gemmC = buildPackedBias(gemmBias, biasMatrix, biasDenseAttr, state, plan, rewriter, loc);
 
-  Value gemmRows = ONNXGemmOp::create(rewriter,
+  FailureOr<Value> gemmRows = lowerGemmToSpatial(
-                                      loc,
+    state.diagnosticAnchor, gemmInputRows, gemmB, gemmC,
-                                      plan.gemmOutputRowsType,
+    plan.gemmOutputRowsType, /*transA=*/false, /*transB=*/!wDenseAttr,
-                                      gemmInputRows,
+    /*alpha=*/1.0f, /*beta=*/1.0f, state.targetInfo(), rewriter, loc);
-                                      gemmB,
+  if (failed(gemmRows))
-                                      gemmC,
+    return failure();
-                                      APFloat(1.0f),
-                                      APFloat(1.0f),
-                                      /*transA=*/0,
-                                      /*transB=*/!wDenseAttr)
-                     .getY();
 
-  return createCollectedConvOutput(ValueRange {gemmRows},
+  return createCollectedConvOutput(ValueRange {*gemmRows},
                                    state.outType,
                                    plan.gemmOutType,
                                    plan.nhwcType,
@@ -1490,10 +1477,10 @@ static Value rewritePackedIm2ColConv(const ConvLoweringState& state,
                                    loc);
 }
 
-static Value rewriteStreamedConv(const ConvLoweringState& state,
+static FailureOr<Value> rewriteStreamedConv(const ConvLoweringState& state,
-                                 PatternRewriter& rewriter,
+                                            PatternRewriter& rewriter,
-                                 Location loc,
+                                            Location loc,
-                                 int64_t forcedPackFactor) {
+                                            int64_t forcedPackFactor) {
   auto wDenseAttr = getHostConstDenseElementsAttr(state.w);
   PreparedConvInput preparedInput = prepareInputForIm2Col(state, rewriter, loc);
   Value biasMatrix;
@@ -1506,20 +1493,22 @@ static Value rewriteStreamedConv(const ConvLoweringState& state,
   ConvGemmPlan seedPlan = buildConvGemmPlan(
     state, static_cast<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(biasDenseAttr), 0, 1, forcedPackFactor);
   Value weightMatrix = createWeightMatrix(state.w, seedPlan, static_cast<bool>(wDenseAttr), rewriter, loc);
-  Value collectedRows = createStreamedConvRows(state,
+  FailureOr<Value> collectedRows = createStreamedConvRows(state,
-                                              preparedInput,
+                                                          preparedInput,
-                                              weightMatrix,
+                                                          weightMatrix,
-                                              biasMatrix,
+                                                          biasMatrix,
-                                              wDenseAttr,
+                                                          wDenseAttr,
-                                              biasDenseAttr,
+                                                          biasDenseAttr,
-                                              forcedPackFactor,
+                                                          forcedPackFactor,
-                                              rewriter,
+                                                          rewriter,
-                                              loc);
+                                                          loc);
-  auto gemmOutType = cast<RankedTensorType>(collectedRows.getType());
+  if (failed(collectedRows))
+    return failure();
+  auto gemmOutType = cast<RankedTensorType>(collectedRows->getType());
   auto nhwcType = RankedTensorType::get({state.batchSize, state.outHeight, state.outWidth, state.numChannelsOut},
                                         state.outType.getElementType());
   return createCollectedConvOutput(
-    ValueRange {collectedRows}, state.outType, gemmOutType, nhwcType, state.outType, gemmOutType.getDimSize(0),
+    ValueRange {*collectedRows}, state.outType, gemmOutType, nhwcType, state.outType, gemmOutType.getDimSize(0),
     state.numChannelsOut, /*packFactor=*/1, rewriter, loc);
 }
 
@@ -1534,12 +1523,12 @@ static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter
 static bool rowStripOutputTileFitsOneCore(const ConvGeometry& geometry) {
   return ceilIntegerDivide(geometry.k, geometry.xbarSize)
            * ceilIntegerDivide(geometry.c, geometry.xbarSize)
-         <= static_cast<int64_t>(crossbarCountInCore.getValue());
+         <= geometry.matrixUnitsPerProcessor;
 }
 
 static bool rowStripOutputChannelTileFitsOneCore(const ConvGeometry& geometry) {
   return ceilIntegerDivide(geometry.k, geometry.xbarSize)
-         <= static_cast<int64_t>(crossbarCountInCore.getValue());
+         <= geometry.matrixUnitsPerProcessor;
 }
 
 static int64_t chooseRowStripPixelPackFactor(const ConvLoweringState& state, int64_t xbarDim) {
@@ -1581,7 +1570,7 @@ static bool canConsumePixelMajorRowStripFragments(const ConvLoweringState& state
     failureReason = "non_constant_weight";
     return false;
   }
-  if (!rowStripOutputChannelTileFitsOneCore(buildConvGeometry(state))) {
+  if (!rowStripOutputChannelTileFitsOneCore(buildConvGeometry(state, state.targetInfo()))) {
     failureReason = "output_channel_tile_does_not_fit_one_core";
     return false;
   }
@@ -1790,8 +1779,7 @@ static Value extractDenseConvWindowRow(Value denseInput,
                                    rewriter.getIndexAttr(state.xWidth)};
   Value nchw = tensor::ExtractSliceOp::create(
     rewriter, loc, nchwType, denseInput, offsets, sizes, getUnitStrides(rewriter, 4));
-  return ONNXTransposeOp::create(
+  return createLinalgTranspose(nchw, fragmentType, {0, 2, 3, 1}, rewriter, loc);
-    rewriter, loc, fragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1}));
 }
 
 static Value createRowStripWindowMaskTable(const ConvLoweringState& state, PatternRewriter& rewriter) {
@@ -2429,7 +2417,7 @@ static FailureOr<Value> createOutputChannelTiledRowStripConvOutput(const ConvLow
 
 static FailureOr<Value>
   createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRewriter& rewriter, Location loc) {
-  ConvGeometry geometry = buildConvGeometry(state);
+  ConvGeometry geometry = buildConvGeometry(state, state.targetInfo());
   if (state.group != 1 || state.batchSize != 1 || !rowStripOutputChannelTileFitsOneCore(geometry))
     return failure();
 
@@ -2481,7 +2469,7 @@ static FailureOr<Value> createConvOutputFromPixelMajorRowStripFragments(Value ro
   if (!canConsumePixelMajorRowStripFragments(state, failureReason))
     return failure();
 
-  ConvGeometry geometry = buildConvGeometry(state);
+  ConvGeometry geometry = buildConvGeometry(state, state.targetInfo());
   const int64_t xbarDim = geometry.xbarSize;
   const int64_t basePatchSize = state.numChannelsIn * state.wHeight * state.wWidth;
   const int64_t baseNumKSlices = ceilIntegerDivide(basePatchSize, xbarDim);
@@ -2521,7 +2509,7 @@ static FailureOr<Value> createPointwiseOutputFromRowStripFragments(Value rowStri
                                                                    Location loc) {
   FailureOr<RowStripPhysicalValue> input = describeRowStripPhysicalValue(rowStripStorage, state.xType);
   if (failed(input)) return failure();
-  ConvGeometry geometry = buildConvGeometry(state);
+  ConvGeometry geometry = buildConvGeometry(state, state.targetInfo());
   const int64_t xbarDim = geometry.xbarSize;
   const int64_t inputFragmentChannels = input->fragmentType.getDimSize(3);
   if (inputFragmentChannels % xbarDim != 0 || state.numChannelsIn % xbarDim != 0)
@@ -2621,8 +2609,10 @@ static bool canConsumeDepthwiseRowStrip(const ConvLoweringState& state) {
                                          state.wHeight,
                                          state.wWidth,
                                          state.outHeight,
-                                         state.outWidth);
+                                         state.outWidth,
-  return tiling && tiling->numChannelTiles <= static_cast<int64_t>(crossbarCountInCore.getValue());
+                                         state.targetInfo().matrixShape.rows);
+  return tiling && tiling->numChannelTiles
+                  <= static_cast<int64_t>(state.targetInfo().matrixUnitsPerProcessor);
 }
 
 static Value insertDepthwiseInputSegment(Value inputWindow,
@@ -2726,16 +2716,23 @@ static FailureOr<Value> createDepthwiseOutputFromRowStripFragments(Value rowStri
                                          state.wHeight,
                                          state.wWidth,
                                          state.outHeight,
-                                         state.outWidth);
+                                         state.outWidth,
+                                         state.targetInfo().matrixShape.rows);
   auto weight = getHostConstDenseElementsAttr(state.w);
   if (!tiling || !weight)
     return failure();
 
   Value packedWeights = depthwise::buildPackedWeights(
-    weight, state.wType, *tiling, rewriter, loc, static_cast<int64_t>(crossbarSize.getValue()));
+    weight,
+    state.wType,
+    *tiling,
+    rewriter,
+    loc,
+    static_cast<int64_t>(state.targetInfo().matrixShape.rows),
+    static_cast<int64_t>(state.targetInfo().matrixShape.rows));
   Value bias = state.hasBias ? expandBiasIfNeeded(state.b, rewriter, loc) : Value();
   auto paddedOutputType = RankedTensorType::get(
-    {1, static_cast<int64_t>(crossbarSize.getValue())}, state.outType.getElementType());
+    {1, static_cast<int64_t>(state.targetInfo().matrixShape.rows)}, state.outType.getElementType());
   auto outputTileType = RankedTensorType::get(
     {1, tiling->tileOutputChannels}, state.outType.getElementType());
   auto outputPixelType = RankedTensorType::get(
@@ -2759,7 +2756,7 @@ static FailureOr<Value> createDepthwiseOutputFromRowStripFragments(Value rowStri
       Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0);
       Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1);
       Value cOutWidth = getOrCreateIndexConstant(rewriter, anchor, state.outWidth);
-      const int64_t xbarDim = static_cast<int64_t>(crossbarSize.getValue());
+      const int64_t xbarDim = static_cast<int64_t>(state.targetInfo().matrixShape.rows);
       auto paddedInputScratchType = RankedTensorType::get(
         {tiling->numChannelTiles, 1, 1, xbarDim}, state.xType.getElementType(), state.xType.getEncoding());
       auto tileScratchType = RankedTensorType::get(
@@ -2773,7 +2770,8 @@ static FailureOr<Value> createDepthwiseOutputFromRowStripFragments(Value rowStri
       SmallVector<Value> biasTiles;
       SmallVector<Value> tileIndices;
       SmallVector<OpFoldResult> weightTileSizes {
-        rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)};
+        rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim),
+        rewriter.getIndexAttr(xbarDim)};
       for (int64_t tile = 0; tile < tiling->numChannelTiles; ++tile) {
         Value tileIndex = getOrCreateIndexConstant(rewriter, anchor, tile);
         tileIndices.push_back(tileIndex);
@@ -2861,10 +2859,10 @@ static FailureOr<Value> createDepthwiseOutputFromRowStripFragments(Value rowStri
 
 static FailureOr<Value> createConvOutputFromRowStripInput(const ConvLoweringState& state,
                                                           Value rowStripInput,
-                                                          PimConvLoweringType strategy,
+                                                          spatial::ConvLoweringStrategy strategy,
                                                           PatternRewriter& rewriter,
                                                           Location loc) {
-  if (strategy == PimConvLoweringDepthwise)
+  if (strategy == spatial::ConvLoweringStrategy::Depthwise)
     return createDepthwiseOutputFromRowStripFragments(rowStripInput, state, rewriter, loc);
   if (state.xHeight == 1 && state.xWidth == 1 && state.wHeight == 1 && state.wWidth == 1)
     return createPointwiseOutputFromRowStripFragments(rowStripInput, state, rewriter, loc);
@@ -2905,17 +2903,23 @@ static Value createCollectedConvOutput(ValueRange gemmRows,
                                                     {0, 1, 2},
                                                     {3}
     });
-    Value nchwOut = ONNXTransposeOp::create(rewriter, loc, outType, nhwcOut, rewriter.getI64ArrayAttr({0, 3, 1, 2}));
+    Value nchwOut = createLinalgTranspose(nhwcOut, outType, {0, 3, 1, 2}, rewriter, loc);
     spatial::SpatYieldOp::create(rewriter, loc, nchwOut);
   });
   return collectComputeOp.getResult(0);
 }
 
-static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp, Value x, Value w, Value b) {
+static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp,
+                                                             Value x,
+                                                             Value w,
+                                                             Value b,
+                                                             const spatial::SpatialTargetInfo& target) {
   ConvLoweringState state;
+  state.diagnosticAnchor = convOp.getOperation();
   state.x = x;
   state.w = w;
   state.b = b;
+  state.target = &target;
   state.xType = cast<RankedTensorType>(state.x.getType());
   state.wType = cast<RankedTensorType>(state.w.getType());
   state.outType = cast<RankedTensorType>(convOp.getY().getType());
@@ -3019,6 +3023,7 @@ static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp,
     state.padWidthBegin = getI64Attr(*padsAttr, 1);
     state.padHeightEnd = getI64Attr(*padsAttr, 2);
     state.padWidthEnd = getI64Attr(*padsAttr, 3);
+    classifyConvProblem(state);
     return state;
   }
 
@@ -3043,6 +3048,7 @@ static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp,
       state.padWidthEnd = totalPadW / 2;
       state.padWidthBegin = totalPadW - state.padWidthEnd;
     }
+    classifyConvProblem(state);
     return state;
   }
 
@@ -3051,18 +3057,25 @@ static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp,
     return failure();
   }
 
+  classifyConvProblem(state);
   return state;
 }
 
-static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp, ONNXConvOpAdaptor convOpAdaptor) {
+static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp,
-  return analyzeConvLoweringState(convOp, convOpAdaptor.getX(), convOpAdaptor.getW(), convOpAdaptor.getB());
+                                                             ONNXConvOpAdaptor convOpAdaptor,
+                                                             const spatial::SpatialTargetInfo& target) {
+  return analyzeConvLoweringState(
+    convOp, convOpAdaptor.getX(), convOpAdaptor.getW(), convOpAdaptor.getB(), target);
 }
 
-static FailureOr<ConvLoweringState> analyzeConvLoweringState(spatial::SpatConv2DPlanOp planOp) {
+static FailureOr<ConvLoweringState> analyzeConvLoweringState(
+  spatial::SpatConv2DPlanOp planOp, const spatial::SpatialTargetInfo& target) {
   ConvLoweringState state;
+  state.diagnosticAnchor = planOp.getOperation();
   state.x = planOp.getInput();
   state.w = planOp.getWeight();
   state.b = planOp.getBias() ? planOp.getBias() : Value();
+  state.target = &target;
   state.xType = dyn_cast<RankedTensorType>(state.x.getType());
   state.wType = dyn_cast<RankedTensorType>(state.w.getType());
   state.outType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
@@ -3111,79 +3124,56 @@ static FailureOr<ConvLoweringState> analyzeConvLoweringState(spatial::SpatConv2D
   state.strideWidth = strides[1];
   state.dilationHeight = dilations[0];
   state.dilationWidth = dilations[1];
+  classifyConvProblem(state);
   return state;
 }
 
-static FailureOr<PimConvLoweringType> resolveRequestedConvLoweringStrategy(Operation* op) {
+static FailureOr<spatial::ConvLoweringStrategy>
-  if (!useExperimentalConvImpl)
+resolveRequestedConvLoweringStrategy(Operation* op, const spatial::SpatialTargetInfo& target) {
-    return pimConvLowering.getValue();
+  if (!target.useExperimentalConvImplementation)
+    return target.convLoweringStrategy;
 
-  if (pimConvLowering != PimConvLoweringAuto && pimConvLowering != PimConvLoweringPackedIm2Col) {
+  if (target.convLoweringStrategy != spatial::ConvLoweringStrategy::Auto
+      && target.convLoweringStrategy != spatial::ConvLoweringStrategy::PackedIm2Col) {
     op->emitOpError() << "--use-experimental-conv-impl conflicts with --pim-conv-lowering="
-                      << stringifyConvLoweringStrategy(pimConvLowering);
+                      << stringifyConvLoweringStrategy(target.convLoweringStrategy);
     return failure();
   }
-  return PimConvLoweringPackedIm2Col;
+  return spatial::ConvLoweringStrategy::PackedIm2Col;
 }
 
-static LogicalResult verifyForcedConvLoweringStrategy(Operation* op,
+static FailureOr<ConvPlan> selectConvLoweringPlan(
-                                                      const ConvGeometry& geo,
+  Operation* op, const ConvLoweringState& state) {
-                                                      PimConvLoweringType strategy) {
+  FailureOr<spatial::ConvLoweringStrategy> requested =
-  switch (strategy) {
+    resolveRequestedConvLoweringStrategy(op, state.targetInfo());
-  case PimConvLoweringAuto:
-  case PimConvLoweringLegacy:
-    return success();
-  case PimConvLoweringDepthwise:
-    if (geo.isDepthwise)
-      return success();
-    return op->emitOpError("forced depthwise Conv lowering requires a depthwise convolution");
-  case PimConvLoweringPackedIm2Col:
-    if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize && geo.pack >= 2 && geo.im2colElements <= pimConvIm2colMaxElements)
-      return success();
-    return op->emitOpError("forced packed-im2col Conv lowering requires K/C to fit, pack >= 2, and im2col within budget");
-  case PimConvLoweringStreamedPatch:
-    if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize)
-      return success();
-    return op->emitOpError("forced streamed-patch Conv lowering requires K and C to each fit one crossbar");
-  case PimConvLoweringStreamedPacked:
-    if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize && geo.pack >= 2)
-      return success();
-    return op->emitOpError("forced streamed-packed Conv lowering requires K/C to fit and pack >= 2");
-  case PimConvLoweringOutputChannelTiled:
-    if (geo.k <= geo.xbarSize && geo.c > geo.xbarSize)
-      return success();
-    return op->emitOpError("forced output-channel-tiled Conv lowering requires K <= X and C > X");
-  case PimConvLoweringInputKTiled:
-    if (geo.k > geo.xbarSize && geo.c <= geo.xbarSize)
-      return success();
-    return op->emitOpError("forced input-k-tiled Conv lowering requires K > X and C <= X");
-  case PimConvLoweringTiled2D:
-    if (geo.k > geo.xbarSize && geo.c > geo.xbarSize)
-      return success();
-    return op->emitOpError("forced tiled-2d Conv lowering requires K > X and C > X");
-  }
-  llvm_unreachable("unknown conv lowering strategy");
-}
-
-static FailureOr<PimConvLoweringType> selectConvLoweringStrategy(Operation* op,
-                                                                 const ConvLoweringState& state) {
-  FailureOr<PimConvLoweringType> requested = resolveRequestedConvLoweringStrategy(op);
   if (failed(requested))
     return failure();
 
-  ConvGeometry geometry = buildConvGeometry(state);
+  if (*requested == spatial::ConvLoweringStrategy::Auto) {
-  PimConvLoweringType strategy = chooseConvLoweringStrategy(geometry, *requested);
+    for (const ConvPlan& candidate : buildConvPlanCandidates(state, state.targetInfo())) {
-  if (strategy == PimConvLoweringDepthwise && !depthwise::canUseStructuredRewrite(state)
+      if (candidate.strategy == spatial::ConvLoweringStrategy::Depthwise
-      && *requested == PimConvLoweringAuto)
+          && !depthwise::canUseStructuredRewrite(state)) {
-    strategy = PimConvLoweringLegacy;
+        continue;
-  if (failed(verifyForcedConvLoweringStrategy(op, geometry, strategy)))
+      }
+      return candidate;
+    }
+    op->emitOpError("has no applicable Conv lowering candidate for the injected Spatial target");
     return failure();
-  return strategy;
+  }
+
+  FailureOr<ConvPlan> candidate = makeConvPlan(state, *requested, state.targetInfo());
+  if (failed(candidate)) {
+    op->emitOpError() << "forced Conv lowering `"
+                      << stringifyConvLoweringStrategy(*requested)
+                      << "` is not applicable to this Conv problem";
+    return failure();
+  }
+  return *candidate;
 }
 
 static FailureOr<Value> lowerDenseSelectedConvPlan(Operation* op,
                                                    const ConvLoweringState& state,
-                                                   PimConvLoweringType strategy,
+                                                   spatial::ConvLoweringStrategy strategy,
                                                    PatternRewriter& rewriter,
                                                    Location loc);
 
@@ -3196,18 +3186,18 @@ static ConvLoweringState makeGroupedConvLoweringState(const ConvLoweringState& p
 static FailureOr<Value> buildConvValueForStrategy(Operation* op,
                                                   Location loc,
                                                   const ConvLoweringState& state,
-                                                  PimConvLoweringType strategy,
+                                                  spatial::ConvLoweringStrategy strategy,
                                                   PatternRewriter& rewriter);
 
 static FailureOr<Value> buildGroupedConvValue(Operation* op,
                                               Location loc,
                                               const ConvLoweringState& state,
-                                              PimConvLoweringType strategy,
+                                              spatial::ConvLoweringStrategy strategy,
                                               PatternRewriter& rewriter);
 
 static FailureOr<Value> lowerGroupedSelectedConvPlan(Operation* op,
                                                      const ConvLoweringState& state,
-                                                     PimConvLoweringType strategy,
+                                                     spatial::ConvLoweringStrategy strategy,
                                                      PatternRewriter& rewriter,
                                                      Location loc) {
   return buildGroupedConvValue(op, loc, state, strategy, rewriter);
@@ -3215,7 +3205,7 @@ static FailureOr<Value> lowerGroupedSelectedConvPlan(Operation* op,
 
 static FailureOr<Value> lowerDenseSelectedConvPlan(Operation* op,
                                                    const ConvLoweringState& state,
-                                                   PimConvLoweringType strategy,
+                                                   spatial::ConvLoweringStrategy strategy,
                                                    PatternRewriter& rewriter,
                                                    Location loc) {
   return buildConvValueForStrategy(op, loc, state, strategy, rewriter);
@@ -3224,29 +3214,29 @@ static FailureOr<Value> lowerDenseSelectedConvPlan(Operation* op,
 static FailureOr<Value> buildConvValueForStrategy(Operation* op,
                                                   Location loc,
                                                   const ConvLoweringState& state,
-                                                  PimConvLoweringType strategy,
+                                                  spatial::ConvLoweringStrategy strategy,
                                                   PatternRewriter& rewriter) {
-  const ConvGeometry geo = buildConvGeometry(state);
+  const ConvGeometry geo = buildConvGeometry(state, state.targetInfo());
   switch (strategy) {
-  case PimConvLoweringDepthwise: {
+  case spatial::ConvLoweringStrategy::Depthwise: {
     return depthwise::rewriteConv(op, state, rewriter, loc);
   }
-  case PimConvLoweringLegacy:
+  case spatial::ConvLoweringStrategy::Legacy:
-  case PimConvLoweringPackedIm2Col: {
+  case spatial::ConvLoweringStrategy::PackedIm2Col: {
     return standard::rewritePackedIm2ColConv(state, rewriter, loc);
   }
-  case PimConvLoweringStreamedPatch:
+  case spatial::ConvLoweringStrategy::StreamedPatch:
-  case PimConvLoweringOutputChannelTiled:
+  case spatial::ConvLoweringStrategy::OutputChannelTiled:
-  case PimConvLoweringTiled2D: {
+  case spatial::ConvLoweringStrategy::Tiled2D: {
     return standard::rewriteStreamedConv(state, rewriter, loc, /*forcedPackFactor=*/1);
   }
-  case PimConvLoweringInputKTiled: {
+  case spatial::ConvLoweringStrategy::InputKTiled: {
     return standard::rewriteInputKTiledConv(state, rewriter, loc);
   }
-  case PimConvLoweringStreamedPacked: {
+  case spatial::ConvLoweringStrategy::StreamedPacked: {
     return standard::rewriteStreamedConv(state, rewriter, loc, geo.pack);
   }
-  case PimConvLoweringAuto:
+  case spatial::ConvLoweringStrategy::Auto:
     break;
   }
   op->emitOpError("unexpected auto strategy at Conv lowering dispatch");
@@ -3282,13 +3272,14 @@ static ConvLoweringState makeGroupedConvLoweringState(
   state.numChannelsInPerGroup = state.numChannelsIn;
   state.numChannelsOutPerGroup = state.numChannelsOut;
   state.hasBias = static_cast<bool>(groupB);
+  classifyConvProblem(state);
   return state;
 }
 
 static FailureOr<Value> buildGroupedConvValue(Operation* op,
                                               Location loc,
                                               const ConvLoweringState& state,
-                                              PimConvLoweringType strategy,
+                                              spatial::ConvLoweringStrategy strategy,
                                               PatternRewriter& rewriter) {
   SmallVector<Value> xSlices = sliceTensor(state.x, /*axis=*/1, state.numChannelsInPerGroup, rewriter, loc);
   SmallVector<Value> wSlices = sliceTensor(state.w, /*axis=*/0, state.numChannelsOutPerGroup, rewriter, loc);
@@ -3344,7 +3335,7 @@ static FailureOr<Value> buildGroupedConvValue(Operation* op,
 LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp,
                                           ONNXConvOpAdaptor convOpAdaptor,
                                           ConversionPatternRewriter& rewriter) const {
-  FailureOr<ConvLoweringState> state = analyzeConvLoweringState(convOp, convOpAdaptor);
+  FailureOr<ConvLoweringState> state = analyzeConvLoweringState(convOp, convOpAdaptor, target);
   if (failed(state))
     return failure();
   SmallVector<int64_t> pads {
@@ -3362,15 +3353,20 @@ LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp,
                                                     rewriter.getDenseI64ArrayAttr(strides),
                                                     rewriter.getDenseI64ArrayAttr(dilations),
                                                     rewriter.getI64IntegerAttr(state->group),
-                                                    rewriter.getStringAttr("nchw"));
+                                                    spatial::getNCHWLayout(rewriter.getContext()));
   rewriter.replaceOp(convOp, convPlan.getResult());
   return success();
 }
 
-void populateConvPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.insert<ConvToGemm>(ctx); }
+void populateConvPatterns(RewritePatternSet& patterns,
+                          MLIRContext* ctx,
+                          const spatial::SpatialTargetInfo& target) {
+  patterns.insert<ConvToGemm>(ctx, target);
+}
 
-LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) {
+LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp,
-  FailureOr<ConvLoweringState> state = analyzeConvLoweringState(planOp);
+                                         const spatial::SpatialTargetInfo& target) {
+  FailureOr<ConvLoweringState> state = analyzeConvLoweringState(planOp, target);
   if (failed(state))
     return failure();
 
@@ -3383,38 +3379,41 @@ LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) {
   if (state->hasBias && !isSupportedBiasAddValue(state->b, state->outType))
     return failure();
 
-  ConvGeometry geometry = buildConvGeometry(*state);
+  ConvGeometry geometry = buildConvGeometry(*state, state->targetInfo());
   if (!rowStripOutputChannelTileFitsOneCore(geometry))
     return failure();
-  FailureOr<PimConvLoweringType> strategy = selectConvLoweringStrategy(planOp.getOperation(), *state);
+  FailureOr<ConvPlan> plan =
-  if (failed(strategy))
+    selectConvLoweringPlan(planOp.getOperation(), *state);
+  if (failed(plan))
     return failure();
 
-  switch (*strategy) {
+  switch (plan->strategy) {
-  case PimConvLoweringLegacy:
+  case spatial::ConvLoweringStrategy::Legacy:
-  case PimConvLoweringDepthwise:
+  case spatial::ConvLoweringStrategy::Depthwise:
-  case PimConvLoweringPackedIm2Col:
+  case spatial::ConvLoweringStrategy::PackedIm2Col:
-  case PimConvLoweringStreamedPatch:
+  case spatial::ConvLoweringStrategy::StreamedPatch:
-  case PimConvLoweringOutputChannelTiled:
+  case spatial::ConvLoweringStrategy::OutputChannelTiled:
-  case PimConvLoweringTiled2D:
+  case spatial::ConvLoweringStrategy::Tiled2D:
-  case PimConvLoweringStreamedPacked:
+  case spatial::ConvLoweringStrategy::StreamedPacked:
     return success();
-  case PimConvLoweringAuto:
+  case spatial::ConvLoweringStrategy::Auto:
-  case PimConvLoweringInputKTiled:
+  case spatial::ConvLoweringStrategy::InputKTiled:
     return failure();
   }
   llvm_unreachable("unknown conv lowering strategy");
 }
 
-LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp) {
+LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp,
-  FailureOr<ConvLoweringState> state = analyzeConvLoweringState(planOp);
+                                           const spatial::SpatialTargetInfo& target) {
+  FailureOr<ConvLoweringState> state = analyzeConvLoweringState(planOp, target);
   if (failed(state))
     return failure();
 
-  FailureOr<PimConvLoweringType> strategy = selectConvLoweringStrategy(planOp.getOperation(), *state);
+  FailureOr<ConvPlan> plan =
-  if (failed(strategy))
+    selectConvLoweringPlan(planOp.getOperation(), *state);
+  if (failed(plan))
     return failure();
-  if (*strategy == PimConvLoweringDepthwise)
+  if (plan->strategy == spatial::ConvLoweringStrategy::Depthwise)
     return canConsumeDepthwiseRowStrip(*state) ? success() : failure();
   StringRef failureReason;
   return canConsumePixelMajorRowStripFragments(*state, failureReason) ? success() : failure();
@@ -3424,23 +3423,25 @@ FailureOr<Value>
 lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
                         std::optional<Value> rowStripInput,
                         bool emitRowStripLayout,
+                        const spatial::SpatialTargetInfo& target,
                         PatternRewriter& rewriter) {
-  FailureOr<ConvLoweringState> state = analyzeConvLoweringState(planOp);
+  FailureOr<ConvLoweringState> state = analyzeConvLoweringState(planOp, target);
   if (failed(state))
     return failure();
 
-  FailureOr<PimConvLoweringType> strategy = selectConvLoweringStrategy(planOp.getOperation(), *state);
+  FailureOr<ConvPlan> plan =
-  if (failed(strategy))
+    selectConvLoweringPlan(planOp.getOperation(), *state);
+  if (failed(plan))
     return failure();
 
   if (emitRowStripLayout) {
     if (rowStripInput) {
-      if (failed(canConsumeAndProduceRowStrip(planOp)))
+      if (failed(canConsumeAndProduceRowStrip(planOp, target)))
         return planOp.emitOpError("selected row-strip input/output layout is not supported for this Conv plan"), failure();
       return createConvOutputFromRowStripInput(
-        *state, *rowStripInput, *strategy, rewriter, planOp.getLoc());
+        *state, *rowStripInput, plan->strategy, rewriter, planOp.getLoc());
     }
-    if (failed(canLowerConvPlanToRowStrip(planOp)))
+    if (failed(canLowerConvPlanToRowStrip(planOp, target)))
       return planOp.emitOpError("selected row-strip layout is not supported for this Conv plan"), failure();
     FailureOr<Value> rowStripStorage = createRowStripConvOutputFromDenseInput(*state, rewriter, planOp.getLoc());
     if (failed(rowStripStorage))
@@ -3448,11 +3449,11 @@ lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
     return *rowStripStorage;
   }
 
-  if (*strategy == PimConvLoweringDepthwise)
+  if (plan->strategy == spatial::ConvLoweringStrategy::Depthwise)
-    return lowerDenseSelectedConvPlan(planOp.getOperation(), *state, *strategy, rewriter, planOp.getLoc());
+    return lowerDenseSelectedConvPlan(planOp.getOperation(), *state, plan->strategy, rewriter, planOp.getLoc());
   if (state->group != 1)
-    return lowerGroupedSelectedConvPlan(planOp.getOperation(), *state, *strategy, rewriter, planOp.getLoc());
+    return lowerGroupedSelectedConvPlan(planOp.getOperation(), *state, plan->strategy, rewriter, planOp.getLoc());
-  return lowerDenseSelectedConvPlan(planOp.getOperation(), *state, *strategy, rewriter, planOp.getLoc());
+  return lowerDenseSelectedConvPlan(planOp.getOperation(), *state, plan->strategy, rewriter, planOp.getLoc());
 }
 
 } // namespace onnx_mlir

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

@@ -3,47 +3,277 @@
 #include <algorithm>
 
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 
 namespace onnx_mlir {
 
+namespace {
+
+static int64_t ceilDivide(int64_t value, int64_t divisor) {
+  return divisor == 0 ? 0 : (value + divisor - 1) / divisor;
+}
+
+} // namespace
+
 bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup) {
   return group == numChannelsIn && numChannelsInPerGroup == 1 && numChannelsOut % group == 0;
 }
 
-ConvGeometry buildConvGeometry(const ConvLoweringState& state) {
+void classifyConvProblem(ConvProblem& problem) {
+  problem.isDepthwise = isDepthwiseConv(
+    problem.group, problem.numChannelsIn, problem.numChannelsOut,
+    problem.numChannelsInPerGroup);
+  problem.isGrouped = problem.group > 1;
+  problem.isPointwise = problem.wHeight == 1 && problem.wWidth == 1
+    && problem.strideHeight == 1 && problem.strideWidth == 1
+    && problem.dilationHeight == 1 && problem.dilationWidth == 1
+    && problem.padHeightBegin == 0 && problem.padHeightEnd == 0
+    && problem.padWidthBegin == 0 && problem.padWidthEnd == 0;
+}
+
+ConvGeometry buildConvGeometry(const ConvProblem& problem,
+                               const spatial::SpatialTargetInfo& target) {
   ConvGeometry geo {
-    state.batchSize,
+    problem.batchSize,
-    state.numChannelsIn,
+    problem.numChannelsIn,
-    state.xHeight,
+    problem.xHeight,
-    state.xWidth,
+    problem.xWidth,
-    state.numChannelsOut,
+    problem.numChannelsOut,
-    state.wHeight,
+    problem.wHeight,
-    state.wWidth,
+    problem.wWidth,
-    state.outHeight,
+    problem.outHeight,
-    state.outWidth,
+    problem.outWidth,
-    state.group,
+    problem.group,
-    state.numChannelsInPerGroup,
+    problem.numChannelsInPerGroup,
-    state.numChannelsOutPerGroup,
+    problem.numChannelsOutPerGroup,
-    state.numChannelsInPerGroup * state.wHeight * state.wWidth,
+    problem.numChannelsInPerGroup * problem.wHeight * problem.wWidth,
-    state.numChannelsOutPerGroup,
+    problem.numChannelsOutPerGroup,
-    state.batchSize * state.outHeight * state.outWidth,
+    problem.batchSize * problem.outHeight * problem.outWidth,
-    static_cast<int64_t>(crossbarSize.getValue()),
+    static_cast<int64_t>(target.matrixShape.rows),
+    static_cast<int64_t>(target.matrixUnitsPerProcessor),
     1,
     0,
-    state.hasBias,
+    problem.hasBias,
-    isDepthwiseConv(state.group, state.numChannelsIn, state.numChannelsOut, state.numChannelsInPerGroup),
+    isDepthwiseConv(problem.group,
+                    problem.numChannelsIn,
+                    problem.numChannelsOut,
+                    problem.numChannelsInPerGroup),
   };
   geo.pack = std::max<int64_t>(1, geo.xbarSize / std::max<int64_t>(geo.k, geo.c));
   geo.im2colElements = static_cast<uint64_t>(std::max<int64_t>(0, geo.p)) * static_cast<uint64_t>(std::max<int64_t>(0, geo.k));
   return geo;
 }
 
-uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor) {
+static ConvMaterializationKind getMaterializationKind(
+    const ConvProblem& problem, spatial::ConvLoweringStrategy strategy) {
+  if (strategy == spatial::ConvLoweringStrategy::Depthwise)
+    return ConvMaterializationKind::StructuredDepthwise;
+  if (problem.isPointwise)
+    return ConvMaterializationKind::PointwiseContraction;
+  switch (strategy) {
+  case spatial::ConvLoweringStrategy::Depthwise:
+    return ConvMaterializationKind::StructuredDepthwise;
+  case spatial::ConvLoweringStrategy::PackedIm2Col:
+  case spatial::ConvLoweringStrategy::Legacy:
+    return ConvMaterializationKind::PackedIm2Col;
+  case spatial::ConvLoweringStrategy::StreamedPatch:
+    return ConvMaterializationKind::StreamedPatch;
+  case spatial::ConvLoweringStrategy::StreamedPacked:
+    return ConvMaterializationKind::StreamedPacked;
+  case spatial::ConvLoweringStrategy::OutputChannelTiled:
+    return ConvMaterializationKind::OutputChannelTiled;
+  case spatial::ConvLoweringStrategy::InputKTiled:
+    return ConvMaterializationKind::InputKTiled;
+  case spatial::ConvLoweringStrategy::Tiled2D:
+    return ConvMaterializationKind::Tiled2D;
+  case spatial::ConvLoweringStrategy::Auto:
+    break;
+  }
+  llvm_unreachable("auto is not a Conv materialization kind");
+}
+
+static ConvPlan makeCandidatePlan(const ConvProblem& problem,
+                                  spatial::ConvLoweringStrategy strategy,
+                                  const spatial::SpatialTargetInfo& target) {
+  ConvPlan plan;
+  plan.geometry = buildConvGeometry(problem, target);
+  plan.strategy = strategy;
+  plan.materializationKind = getMaterializationKind(problem, strategy);
+  plan.laneCount = plan.geometry.p;
+  plan.reductionCount = std::max<int64_t>(
+    1, (plan.geometry.k + plan.geometry.xbarSize - 1) / plan.geometry.xbarSize);
+  plan.mvmCount = plan.laneCount * plan.reductionCount;
+  plan.vectorCount = plan.mvmCount;
+  plan.weightElements = static_cast<uint64_t>(std::max<int64_t>(0, problem.numChannelsOut))
+    * static_cast<uint64_t>(std::max<int64_t>(0, plan.geometry.k));
+  plan.scratchElements = plan.geometry.im2colElements;
+  plan.materializationElements = strategy == spatial::ConvLoweringStrategy::Depthwise
+    ? 0
+    : std::min<uint64_t>(plan.geometry.im2colElements, target.convIm2colMaxElements);
+  plan.requiresInputMaterialization = strategy != spatial::ConvLoweringStrategy::Depthwise;
+  plan.producesRowStrip = strategy != spatial::ConvLoweringStrategy::InputKTiled
+    && ceilDivide(plan.geometry.k, plan.geometry.xbarSize) <= plan.geometry.matrixUnitsPerProcessor;
+  plan.consumesRowStrip = plan.producesRowStrip;
+  // Conv materializers emit local compute and leave inter-core communication
+  // to Spatial scheduling; zero is an explicit ownership statement here.
+  plan.communicationElements = 0;
+  plan.usesContraction = problem.isPointwise || strategy != spatial::ConvLoweringStrategy::Depthwise;
+  if (problem.isPointwise) {
+    ContractionProblem contraction;
+    contraction.origin = ContractionOrigin::Gemm;
+    contraction.batch = 1;
+    contraction.m = plan.geometry.p;
+    contraction.k = plan.geometry.c;
+    contraction.n = problem.numChannelsOutPerGroup;
+    contraction.lhsElementType = problem.xType.getElementType();
+    contraction.rhsElementType = problem.wType.getElementType();
+    contraction.resultElementType = problem.outType.getElementType();
+    plan.contraction = makeContractionPlan(
+      contraction, target, ContractionPlanKind::StaticTiled);
+    plan.hasContractionPlan = true;
+    plan.laneCount = plan.contraction.laneCount;
+    plan.mvmCount = plan.contraction.expectedMvmCount;
+    plan.vectorCount = plan.contraction.expectedVectorCount;
+    plan.reductionCount = plan.contraction.reductionSlices;
+  }
+  return plan;
+}
+
+static bool fitsSingleCrossbar(const ConvGeometry& geo) {
+  return geo.k <= geo.xbarSize && geo.c <= geo.xbarSize;
+}
+
+static bool fitsPackedIm2Col(const ConvGeometry& geo,
+                             const spatial::SpatialTargetInfo& target) {
+  return fitsSingleCrossbar(geo) && geo.pack >= 2
+    && geo.im2colElements <= target.convIm2colMaxElements;
+}
+
+static mlir::FailureOr<ConvPlan> buildDepthwiseCandidate(
+    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
+  if (!problem.isDepthwise)
+    return mlir::failure();
+  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::Depthwise, target);
+}
+
+static mlir::FailureOr<ConvPlan> buildPackedIm2ColCandidate(
+    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
+  ConvGeometry geo = buildConvGeometry(problem, target);
+  if (!fitsPackedIm2Col(geo, target))
+    return mlir::failure();
+  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::PackedIm2Col, target);
+}
+
+static mlir::FailureOr<ConvPlan> buildStreamedPatchCandidate(
+    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
+  if (!fitsSingleCrossbar(buildConvGeometry(problem, target)))
+    return mlir::failure();
+  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::StreamedPatch, target);
+}
+
+static mlir::FailureOr<ConvPlan> buildStreamedPackedCandidate(
+    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
+  ConvGeometry geo = buildConvGeometry(problem, target);
+  if (!fitsSingleCrossbar(geo) || geo.pack < 2)
+    return mlir::failure();
+  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::StreamedPacked, target);
+}
+
+static mlir::FailureOr<ConvPlan> buildOutputChannelTiledCandidate(
+    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
+  ConvGeometry geo = buildConvGeometry(problem, target);
+  if (geo.k > geo.xbarSize || geo.c <= geo.xbarSize)
+    return mlir::failure();
+  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::OutputChannelTiled, target);
+}
+
+static mlir::FailureOr<ConvPlan> buildInputKTiledCandidate(
+    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
+  ConvGeometry geo = buildConvGeometry(problem, target);
+  if (geo.k <= geo.xbarSize || geo.c > geo.xbarSize)
+    return mlir::failure();
+  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::InputKTiled, target);
+}
+
+static mlir::FailureOr<ConvPlan> buildTiled2DCandidate(
+    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
+  ConvGeometry geo = buildConvGeometry(problem, target);
+  if (geo.k <= geo.xbarSize || geo.c <= geo.xbarSize)
+    return mlir::failure();
+  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::Tiled2D, target);
+}
+
+static mlir::FailureOr<ConvPlan> buildLegacyCandidate(
+    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
+  // Legacy is retained as the explicit compatibility/debug materializer and
+  // as the safe fallback when structured depthwise lowering is unavailable.
+  return makeCandidatePlan(problem, spatial::ConvLoweringStrategy::Legacy, target);
+}
+
+mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
+                                       spatial::ConvLoweringStrategy strategy,
+                                       const spatial::SpatialTargetInfo& target) {
+  switch (strategy) {
+  case spatial::ConvLoweringStrategy::Auto:
+    return mlir::failure();
+  case spatial::ConvLoweringStrategy::Legacy:
+    return buildLegacyCandidate(problem, target);
+  case spatial::ConvLoweringStrategy::Depthwise:
+    return buildDepthwiseCandidate(problem, target);
+  case spatial::ConvLoweringStrategy::PackedIm2Col:
+    return buildPackedIm2ColCandidate(problem, target);
+  case spatial::ConvLoweringStrategy::StreamedPatch:
+    return buildStreamedPatchCandidate(problem, target);
+  case spatial::ConvLoweringStrategy::StreamedPacked:
+    return buildStreamedPackedCandidate(problem, target);
+  case spatial::ConvLoweringStrategy::OutputChannelTiled:
+    return buildOutputChannelTiledCandidate(problem, target);
+  case spatial::ConvLoweringStrategy::InputKTiled:
+    return buildInputKTiledCandidate(problem, target);
+  case spatial::ConvLoweringStrategy::Tiled2D:
+    return buildTiled2DCandidate(problem, target);
+  }
+  llvm_unreachable("unknown Conv lowering strategy");
+}
+
+llvm::SmallVector<ConvPlan, 8> buildConvPlanCandidates(
+    const ConvProblem& problem, const spatial::SpatialTargetInfo& target) {
+  ConvGeometry geo = buildConvGeometry(problem, target);
+  llvm::SmallVector<ConvPlan, 8> candidates;
+  auto append = [&](spatial::ConvLoweringStrategy strategy) {
+    mlir::FailureOr<ConvPlan> candidate = makeConvPlan(problem, strategy, target);
+    if (succeeded(candidate))
+      candidates.push_back(*candidate);
+  };
+
+  if (problem.isDepthwise) {
+    append(spatial::ConvLoweringStrategy::Depthwise);
+    append(spatial::ConvLoweringStrategy::Legacy);
+    return candidates;
+  }
+  if (fitsPackedIm2Col(geo, target))
+    append(spatial::ConvLoweringStrategy::PackedIm2Col);
+  if (fitsSingleCrossbar(geo) && geo.pack >= 2)
+    append(spatial::ConvLoweringStrategy::StreamedPacked);
+  if (fitsSingleCrossbar(geo))
+    append(spatial::ConvLoweringStrategy::StreamedPatch);
+  if (geo.k <= geo.xbarSize && geo.c > geo.xbarSize)
+    append(spatial::ConvLoweringStrategy::OutputChannelTiled);
+  if (geo.k > geo.xbarSize && geo.c <= geo.xbarSize)
+    append(spatial::ConvLoweringStrategy::Legacy);
+  if (geo.k > geo.xbarSize && geo.c <= geo.xbarSize)
+    append(spatial::ConvLoweringStrategy::InputKTiled);
+  if (geo.k > geo.xbarSize && geo.c > geo.xbarSize)
+    append(spatial::ConvLoweringStrategy::Tiled2D);
+  return candidates;
+}
+
+uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
+                                    int64_t packFactor,
+                                    const spatial::SpatialTargetInfo& target) {
   const uint64_t patchElements = static_cast<uint64_t>(std::max<int64_t>(1, geo.k));
-  uint64_t chunkPositions = std::max<uint64_t>(1, pimConvIm2colMaxElements / patchElements);
+  uint64_t chunkPositions = std::max<uint64_t>(1, target.convIm2colMaxElements / patchElements);
   chunkPositions = std::min<uint64_t>(chunkPositions, static_cast<uint64_t>(std::max<int64_t>(1, geo.p)));
-  chunkPositions = std::min<uint64_t>(chunkPositions, std::max<uint64_t>(1, pimConvStreamChunkPositions));
+  chunkPositions = std::min<uint64_t>(chunkPositions, std::max<uint64_t>(1, target.convStreamChunkPositions));
 
   if (packFactor > 1 && chunkPositions > static_cast<uint64_t>(packFactor)) {
     chunkPositions -= chunkPositions % static_cast<uint64_t>(packFactor);
@@ -52,24 +282,26 @@ uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor)
   return std::max<uint64_t>(1, chunkPositions);
 }
 
-RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvLoweringState& state) {
+RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem) {
-  const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin;
+  const int64_t rawBegin = outputRows.begin * problem.strideHeight - problem.padHeightBegin;
   const int64_t rawEnd =
-    (outputRows.end - 1) * state.strideHeight - state.padHeightBegin + state.dilationHeight * (state.wHeight - 1) + 1;
+    (outputRows.end - 1) * problem.strideHeight - problem.padHeightBegin
-  return {std::max<int64_t>(0, rawBegin), std::min<int64_t>(state.xHeight, rawEnd)};
+    + problem.dilationHeight * (problem.wHeight - 1) + 1;
+  return {std::max<int64_t>(0, rawBegin), std::min<int64_t>(problem.xHeight, rawEnd)};
 }
 
-ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state) {
+ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvProblem& problem) {
   ConvRowDemand demand;
   demand.outputRows = outputRows;
-  demand.neededInputRows = computeConvInputRowsForOutputRows(outputRows, state);
+  demand.neededInputRows = computeConvInputRowsForOutputRows(outputRows, problem);
   demand.acquiredInputRows = demand.neededInputRows;
 
-  const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin;
+  const int64_t rawBegin = outputRows.begin * problem.strideHeight - problem.padHeightBegin;
   const int64_t rawEnd =
-    (outputRows.end - 1) * state.strideHeight - state.padHeightBegin + state.dilationHeight * (state.wHeight - 1) + 1;
+    (outputRows.end - 1) * problem.strideHeight - problem.padHeightBegin
+    + problem.dilationHeight * (problem.wHeight - 1) + 1;
   demand.topHaloRows = std::max<int64_t>(0, -rawBegin);
-  demand.bottomHaloRows = std::max<int64_t>(0, rawEnd - state.xHeight);
+  demand.bottomHaloRows = std::max<int64_t>(0, rawEnd - problem.xHeight);
   demand.acquiredInputRows = demand.neededInputRows;
   return demand;
 }

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.hpp

@@ -3,14 +3,19 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Value.h"
 
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
 #include <cstdint>
 
+namespace mlir {
+class Operation;
+} // namespace mlir
+
 namespace onnx_mlir {
 
-struct ConvLoweringState {
+struct ConvProblem {
-  mlir::Value x;
-  mlir::Value w;
-  mlir::Value b;
   mlir::RankedTensorType xType;
   mlir::RankedTensorType wType;
   mlir::RankedTensorType outType;
@@ -35,6 +40,19 @@ struct ConvLoweringState {
   int64_t dilationHeight;
   int64_t dilationWidth;
   bool hasBias;
+  bool isDepthwise = false;
+  bool isGrouped = false;
+  bool isPointwise = false;
+};
+
+struct ConvLoweringState : ConvProblem {
+  mlir::Operation* diagnosticAnchor = nullptr;
+  mlir::Value x;
+  mlir::Value w;
+  mlir::Value b;
+  const spatial::SpatialTargetInfo* target = nullptr;
+
+  const spatial::SpatialTargetInfo& targetInfo() const { return *target; }
 };
 
 struct ConvGeometry {
@@ -54,6 +72,7 @@ struct ConvGeometry {
   int64_t c;
   int64_t p;
   int64_t xbarSize;
+  int64_t matrixUnitsPerProcessor;
   int64_t pack;
   uint64_t im2colElements;
   bool hasBias;
@@ -73,14 +92,59 @@ struct ConvRowDemand {
   int64_t bottomHaloRows = 0;
 };
 
+enum class ConvMaterializationKind : uint8_t {
+  StructuredDepthwise,
+  PointwiseContraction,
+  PackedIm2Col,
+  StreamedPatch,
+  StreamedPacked,
+  OutputChannelTiled,
+  InputKTiled,
+  Tiled2D,
+};
+
+struct ConvPlan {
+  ConvGeometry geometry;
+  spatial::ConvLoweringStrategy strategy = spatial::ConvLoweringStrategy::Auto;
+  ConvMaterializationKind materializationKind = ConvMaterializationKind::PackedIm2Col;
+  int64_t laneCount = 0;
+  int64_t mvmCount = 0;
+  int64_t vectorCount = 0;
+  int64_t reductionCount = 0;
+  uint64_t weightElements = 0;
+  uint64_t scratchElements = 0;
+  uint64_t materializationElements = 0;
+  uint64_t communicationElements = 0;
+  spatial::PhysicalLayout resultLayout = spatial::PhysicalLayout::DenseNCHW;
+  bool consumesRowStrip = false;
+  bool producesRowStrip = false;
+  bool requiresInputMaterialization = false;
+  bool requiresOutputMaterialization = false;
+  bool usesContraction = false;
+  bool hasContractionPlan = false;
+  ContractionPlan contraction;
+};
+
 bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup);
 
-ConvGeometry buildConvGeometry(const ConvLoweringState& state);
+void classifyConvProblem(ConvProblem& problem);
 
-uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor);
+ConvGeometry buildConvGeometry(const ConvProblem& problem,
+                               const spatial::SpatialTargetInfo& target);
 
-RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvLoweringState& state);
+mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
+                                       spatial::ConvLoweringStrategy strategy,
+                                       const spatial::SpatialTargetInfo& target);
 
-ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state);
+llvm::SmallVector<ConvPlan, 8> buildConvPlanCandidates(
+    const ConvProblem& problem, const spatial::SpatialTargetInfo& target);
+
+uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
+                                    int64_t packFactor,
+                                    const spatial::SpatialTargetInfo& target);
+
+RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem);
+
+ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvProblem& problem);
 
 } // namespace onnx_mlir

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

@@ -31,7 +31,7 @@ struct SiluToSpatialPlan : OpRewritePattern<ONNXMulOp> {
       return failure();
 
     auto plan = spatial::SpatSiluPlanOp::create(
-      rewriter, mulOp.getLoc(), mulOp.getResult().getType(), input, rewriter.getStringAttr("nchw"));
+      rewriter, mulOp.getLoc(), mulOp.getResult().getType(), input, spatial::getNCHWLayout(rewriter.getContext()));
     rewriter.replaceOp(mulOp, plan.getResult());
     rewriter.eraseOp(sigmoidOp);
     return success();
@@ -260,14 +260,16 @@ struct AddToSpatialCompute : OpConversionPattern<ONNXAddOp> {
       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, op.getLoc(), resultType, candidate->data, candidate->bias,
+        spatial::getNCHWLayout(rewriter.getContext()));
       rewriter.replaceOp(op, plan.getResult());
       return success();
     }
 
     if (resultType.getRank() == 4 && adaptor.getA().getType() == resultType && adaptor.getB().getType() == resultType) {
       auto plan = spatial::SpatAddPlanOp::create(
-        rewriter, op.getLoc(), resultType, adaptor.getA(), adaptor.getB(), rewriter.getStringAttr("nchw"));
+        rewriter, op.getLoc(), resultType, adaptor.getA(), adaptor.getB(),
+        spatial::getNCHWLayout(rewriter.getContext()));
       rewriter.replaceOp(op, plan.getResult());
       return success();
     }

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

@@ -1,5 +1,6 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -21,6 +22,10 @@
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionProblem.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionMaterialization.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -31,7 +36,7 @@ namespace onnx_mlir {
 namespace {
 
 static FailureOr<Value>
-materializeScaledConstantTensor(Value value, float factor, ConversionPatternRewriter& rewriter, Location loc) {
+materializeScaledConstantTensor(Value value, float factor, PatternRewriter& rewriter, Location loc) {
   if (factor == 1.0f)
     return value;
 
@@ -57,7 +62,12 @@ materializeScaledConstantTensor(Value value, float factor, ConversionPatternRewr
 }
 
 static Value createGemmBatchKOffset(
-  Value lane, int64_t numOutRows, int64_t numKSlices, ConversionPatternRewriter& rewriter, Location loc) {
+  Value lane,
+  int64_t numOutRows,
+  int64_t numKSlices,
+  int64_t xbarSize,
+  PatternRewriter& rewriter,
+  Location loc) {
   if (numKSlices == 1)
     return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
 
@@ -65,7 +75,7 @@ static Value createGemmBatchKOffset(
   AffineExpr d0 = getAffineDimExpr(0, context);
   return createOrFoldAffineApply(rewriter,
                                  loc,
-                                 (d0.floorDiv(numOutRows) % numKSlices) * crossbarSize.getValue(),
+                                 (d0.floorDiv(numOutRows) % numKSlices) * xbarSize,
                                  ValueRange {lane},
                                  rewriter.getInsertionBlock()->getParentOp());
 }
@@ -74,7 +84,8 @@ static Value createGemmBatchHOffset(Value lane,
                                     int64_t numOutRows,
                                     int64_t numKSlices,
                                     int64_t numOutHSlices,
-                                    ConversionPatternRewriter& rewriter,
+                                    int64_t xbarSize,
+                                    PatternRewriter& rewriter,
                                     Location loc) {
   if (numOutHSlices == 1)
     return getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
@@ -83,14 +94,14 @@ static Value createGemmBatchHOffset(Value lane,
   AffineExpr d0 = getAffineDimExpr(0, context);
   return createOrFoldAffineApply(rewriter,
                                  loc,
-                                 d0.floorDiv(numOutRows * numKSlices) * crossbarSize.getValue(),
+                                 d0.floorDiv(numOutRows * numKSlices) * xbarSize,
                                  ValueRange {lane},
                                  rewriter.getInsertionBlock()->getParentOp());
 }
 
 static FailureOr<Value> materializePaddedConstantMatrix(Value value,
                                                         RankedTensorType resultType,
-                                                        ConversionPatternRewriter& rewriter,
+                                                        PatternRewriter& rewriter,
                                                         Location loc) {
   auto sourceType = cast<RankedTensorType>(value.getType());
   if (sourceType == resultType)
@@ -121,7 +132,7 @@ static FailureOr<Value> materializePaddedConstantMatrix(Value value,
 static FailureOr<Value> materializePaddedBroadcastedConstantTensor(Value value,
                                                                    RankedTensorType resultType,
                                                                    int64_t unpaddedColumns,
-                                                                   ConversionPatternRewriter& rewriter,
+                                                                   PatternRewriter& rewriter,
                                                                    Location loc) {
   auto denseAttr = getHostConstDenseElementsAttr(value);
   if (!denseAttr)
@@ -187,7 +198,7 @@ static FailureOr<Value> materializePaddedBroadcastedConstantTensor(Value value,
 static FailureOr<Value> prepareBias(Value c,
                                     RankedTensorType outType,
                                     RankedTensorType paddedOutType,
-                                    ConversionPatternRewriter& rewriter,
+                                    PatternRewriter& rewriter,
                                     Location loc) {
   auto cType = cast<RankedTensorType>(c.getType());
   if (!cType.hasStaticShape())
@@ -203,9 +214,15 @@ static FailureOr<Value> prepareBias(Value c,
 }
 
 static Value extractATile(
-  Value a, Value row, Value kOffset, RankedTensorType aTileType, ConversionPatternRewriter& rewriter, Location loc) {
+  Value a,
+  Value row,
+  Value kOffset,
+  RankedTensorType aTileType,
+  int64_t xbarSize,
+  PatternRewriter& rewriter,
+  Location loc) {
   SmallVector<OpFoldResult> offsets {row, kOffset};
-  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
+  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarSize)};
   SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
 
   return tensor::ExtractSliceOp::create(rewriter, loc, aTileType, a, offsets, sizes, strides).getResult();
@@ -219,7 +236,8 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
                                                            int64_t numOutRows,
                                                            int64_t numKSlices,
                                                            int64_t numOutHSlices,
-                                                           ConversionPatternRewriter& rewriter,
+                                                           int64_t xbarSize,
+                                                           PatternRewriter& rewriter,
                                                            Location loc) {
   const int64_t laneCount = partialPiecesType.getDimSize(0);
   auto batchOp = createSpatComputeBatch(
@@ -232,21 +250,21 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
     [&](detail::SpatComputeBatchBodyArgs args) {
       Value row =
         onnx_mlir::affineModConst(rewriter, loc, args.lane, numOutRows, rewriter.getInsertionBlock()->getParentOp());
-      Value kOffset = createGemmBatchKOffset(args.lane, numOutRows, numKSlices, rewriter, loc);
+      Value kOffset = createGemmBatchKOffset(args.lane, numOutRows, numKSlices, xbarSize, rewriter, loc);
-      Value hOffset = createGemmBatchHOffset(args.lane, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
+      Value hOffset = createGemmBatchHOffset(
+        args.lane, numOutRows, numKSlices, numOutHSlices, xbarSize, rewriter, loc);
 
       auto aTileType =
-        RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, aType.getElementType());
+        RankedTensorType::get({1, xbarSize}, aType.getElementType());
       auto bTileType = RankedTensorType::get(
-        {static_cast<int64_t>(crossbarSize.getValue()), static_cast<int64_t>(crossbarSize.getValue())},
+        {xbarSize, xbarSize},
         paddedBType.getElementType());
       auto pieceType =
-        RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, partialPiecesType.getElementType());
+        RankedTensorType::get({1, xbarSize}, partialPiecesType.getElementType());
-      Value aTile = extractATile(args.inputs.front(), row, kOffset, aTileType, rewriter, loc);
+      Value aTile = extractATile(args.inputs.front(), row, kOffset, aTileType, xbarSize, rewriter, loc);
 
       SmallVector<OpFoldResult> bOffsets {kOffset, hOffset};
-      SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(crossbarSize.getValue()),
+      SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarSize), rewriter.getIndexAttr(xbarSize)};
-                                        rewriter.getIndexAttr(crossbarSize.getValue())};
       SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, 2);
       Value bTile = extractStaticSliceOrIdentity(
         rewriter, loc, args.weights.front(), bTileType, bOffsets, bSizes, unitStrides);
@@ -260,7 +278,7 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
 }
 
 static Value extractDynamicGemmBColumn(
-  Value matrix, Value column, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) {
+  Value matrix, Value column, RankedTensorType vectorType, PatternRewriter& rewriter, Location loc) {
   SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), column};
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(vectorType.getDimSize(1)), rewriter.getIndexAttr(1)};
   SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
@@ -280,7 +298,7 @@ static Value extractDynamicGemmBColumn(
 }
 
 static Value extractDynamicGemmRowVector(
-  Value matrix, Value row, RankedTensorType vectorType, ConversionPatternRewriter& rewriter, Location loc) {
+  Value matrix, Value row, RankedTensorType vectorType, PatternRewriter& rewriter, Location loc) {
   SmallVector<OpFoldResult> offsets {row, rewriter.getIndexAttr(0)};
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1))};
   SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
@@ -317,7 +335,7 @@ static bool hasGemmBias(Value c) {
 
 static Value createScalarTensorConstant(RankedTensorType scalarType,
                                         float value,
-                                        ConversionPatternRewriter& rewriter,
+                                        PatternRewriter& rewriter,
                                         Location loc) {
   auto elementType = scalarType.getElementType();
   auto scalarAttr = rewriter.getFloatAttr(elementType, value);
@@ -330,7 +348,7 @@ static Value createBroadcastedBiasScalar(Value bias,
                                          Value row,
                                          Value column,
                                          RankedTensorType scalarType,
-                                         ConversionPatternRewriter& rewriter,
+                                         PatternRewriter& rewriter,
                                          Location loc) {
   SmallVector<OpFoldResult> unitStrides(biasType.getRank(), rewriter.getIndexAttr(1));
   if (biasType.getRank() == 1) {
@@ -365,7 +383,7 @@ static FailureOr<spatial::SpatComputeBatch> createVvdmulBatch(Value a,
                                                               RankedTensorType columnPiecesType,
                                                               RankedTensorType outType,
                                                               bool transposeB,
-                                                              ConversionPatternRewriter& rewriter,
+                                                              PatternRewriter& rewriter,
                                                               Location loc) {
   const int64_t numOutRows = outType.getDimSize(0);
   const int64_t numOutCols = outType.getDimSize(1);
@@ -425,7 +443,7 @@ static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scal
                                                                       RankedTensorType outType,
                                                                       float alpha,
                                                                       float beta,
-                                                                      ConversionPatternRewriter& rewriter,
+                                                                      PatternRewriter& rewriter,
                                                                       Location loc) {
   const int64_t numOutRows = outType.getDimSize(0);
   const int64_t numOutCols = outType.getDimSize(1);
@@ -510,7 +528,7 @@ static Value createPartialGroupOffset(Value hSlice,
                                       int64_t kSlice,
                                       int64_t numKSlices,
                                       int64_t numOutRows,
-                                      ConversionPatternRewriter& rewriter,
+                                      PatternRewriter& rewriter,
                                       Location loc) {
   MLIRContext* context = rewriter.getContext();
   AffineExpr d0 = getAffineDimExpr(0, context);
@@ -527,10 +545,12 @@ static Value extractReductionPiece(Value partialPiecesArg,
                                    RankedTensorType pieceType,
                                    int64_t numKSlices,
                                    int64_t numOutRows,
-                                   ConversionPatternRewriter& rewriter,
+                                   int64_t xbarSize,
+                                   PatternRewriter& rewriter,
                                    Location loc) {
   SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-  SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
+  SmallVector<OpFoldResult> pieceSizes {
+    rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarSize)};
   SmallVector<OpFoldResult> pieceOffsets {
     createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc),
     rewriter.getIndexAttr(0),
@@ -545,13 +565,15 @@ static Value reducePartialPiecesForHSlice(Value partialPiecesArg,
                                           RankedTensorType pieceType,
                                           int64_t numKSlices,
                                           int64_t numOutRows,
-                                          ConversionPatternRewriter& rewriter,
+                                          int64_t xbarSize,
+                                          PatternRewriter& rewriter,
                                           Location loc) {
   SmallVector<Value> activePieces;
   activePieces.reserve(numKSlices);
   for (int64_t kSlice = 0; kSlice < numKSlices; ++kSlice)
     activePieces.push_back(
-      extractReductionPiece(partialPiecesArg, hSlice, kSlice, pieceType, numKSlices, numOutRows, rewriter, loc));
+      extractReductionPiece(
+        partialPiecesArg, hSlice, kSlice, pieceType, numKSlices, numOutRows, xbarSize, rewriter, loc));
 
   while (activePieces.size() > 1) {
     SmallVector<Value> nextPieces;
@@ -574,11 +596,12 @@ static FailureOr<Value> createReductionOutput(Value partialPieces,
                                               RankedTensorType outType,
                                               RankedTensorType paddedOutType,
                                               int64_t numKSlices,
-                                              ConversionPatternRewriter& rewriter,
+                                              int64_t xbarSize,
+                                              PatternRewriter& rewriter,
                                               Location loc) {
   const int64_t numOutRows = outType.getDimSize(0);
-  const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(1), crossbarSize.getValue());
+  const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(1), xbarSize);
-  auto pieceType = RankedTensorType::get({numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
+  auto pieceType = RankedTensorType::get({numOutRows, xbarSize},
                                          partialPiecesType.getElementType());
 
   if (bias && cast<RankedTensorType>(bias.getType()) != paddedOutType)
@@ -590,20 +613,20 @@ static FailureOr<Value> createReductionOutput(Value partialPieces,
   SmallVector<Value> outputSlices;
   outputSlices.reserve(numOutHSlices);
   for (int64_t hSlice = 0; hSlice < numOutHSlices; ++hSlice) {
-    const int64_t columnOffset = hSlice * crossbarSize.getValue();
+    const int64_t columnOffset = hSlice * xbarSize;
     const int64_t columns =
-      std::min(static_cast<int64_t>(crossbarSize.getValue()), outType.getDimSize(1) - columnOffset);
+      std::min(xbarSize, outType.getDimSize(1) - columnOffset);
     auto outputSliceType = RankedTensorType::get({numOutRows, columns}, outType.getElementType());
     auto computeOp = createSpatCompute(
       rewriter, loc, TypeRange {outputSliceType}, {}, inputs, [&](ValueRange blockArgs) -> LogicalResult {
       Value hSliceValue =
         getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), hSlice);
       Value reduced = reducePartialPiecesForHSlice(
-        blockArgs[0], hSliceValue, pieceType, numKSlices, numOutRows, rewriter, loc);
+        blockArgs[0], hSliceValue, pieceType, numKSlices, numOutRows, xbarSize, rewriter, loc);
       if (bias) {
         SmallVector<OpFoldResult> biasOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(columnOffset)};
         SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows),
-                                              rewriter.getIndexAttr(crossbarSize.getValue())};
+                                              rewriter.getIndexAttr(xbarSize)};
         SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
         Value biasSlice =
           tensor::ExtractSliceOp::create(rewriter, loc, pieceType, blockArgs[1], biasOffsets, pieceSizes, unitStrides)
@@ -637,79 +660,101 @@ static FailureOr<Value> createReductionOutput(Value partialPieces,
 }
 
 struct GemmToSpatialComputes : OpConversionPattern<ONNXGemmOp> {
-  using OpConversionPattern::OpConversionPattern;
+  explicit GemmToSpatialComputes(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
+      : OpConversionPattern<ONNXGemmOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(ONNXGemmOp gemmOp,
                                 ONNXGemmOpAdaptor gemmOpAdaptor,
                                 ConversionPatternRewriter& rewriter) const override;
+
+  const spatial::SpatialTargetInfo& target;
 };
 
 } // namespace
 
-LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
+FailureOr<Value> lowerGemmToSpatial(
-                                                     ONNXGemmOpAdaptor gemmOpAdaptor,
+    Operation* diagnosticAnchor,
-                                                     ConversionPatternRewriter& rewriter) const {
+    Value a,
-  Location loc = gemmOp.getLoc();
+    Value b,
-  Value a = gemmOpAdaptor.getA();
+    Value c,
-  Value b = gemmOpAdaptor.getB();
+    RankedTensorType outType,
-  Value c = gemmOpAdaptor.getC();
+    bool transA,
-
+    bool transB,
+    float alpha,
+    float beta,
+    const spatial::SpatialTargetInfo& target,
+    PatternRewriter& rewriter,
+    Location loc) {
   auto aType = dyn_cast<RankedTensorType>(a.getType());
   auto bType = dyn_cast<RankedTensorType>(b.getType());
-  auto outType = dyn_cast<RankedTensorType>(gemmOp.getY().getType());
+  if (!diagnosticAnchor || !aType || !bType || !outType)
-  if (!aType || !bType || !outType)
     return failure();
   if (!aType.hasStaticShape()) {
-    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm input A");
+    pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm input A");
     return failure();
   }
   if (!bType.hasStaticShape()) {
-    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm input B");
+    pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm input B");
     return failure();
   }
   if (!outType.hasStaticShape()) {
-    pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm result");
+    pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm result");
     return failure();
   }
   if (aType.getRank() != 2) {
-    pim::emitUnsupportedRankDiagnostic(gemmOp, "Gemm input A", aType.getRank(), {2});
+    pim::emitUnsupportedRankDiagnostic(diagnosticAnchor, "Gemm input A", aType.getRank(), {2});
     return failure();
   }
   if (bType.getRank() != 2) {
-    pim::emitUnsupportedRankDiagnostic(gemmOp, "Gemm input B", bType.getRank(), {2});
+    pim::emitUnsupportedRankDiagnostic(diagnosticAnchor, "Gemm input B", bType.getRank(), {2});
     return failure();
   }
   if (outType.getRank() != 2) {
-    pim::emitUnsupportedRankDiagnostic(gemmOp, "Gemm result", outType.getRank(), {2});
+    pim::emitUnsupportedRankDiagnostic(diagnosticAnchor, "Gemm result", outType.getRank(), {2});
     return failure();
   }
 
-  if (gemmOpAdaptor.getTransA()) {
+  if (transA) {
     auto aShape = aType.getShape();
     auto transposedType = RankedTensorType::get({aShape[1], aShape[0]}, aType.getElementType(), aType.getEncoding());
-    a = ONNXTransposeOp::create(rewriter, loc, transposedType, a, rewriter.getI64ArrayAttr({1, 0})).getResult();
+    a = createLinalgTranspose(a, transposedType, {1, 0}, rewriter, loc);
     aType = transposedType;
   }
 
-  const int64_t numOutRows = outType.getDimSize(0);
+  ContractionProblem problem;
-  const int64_t numOutCols = outType.getDimSize(1);
+  problem.lhsBatchShape = {};
-  const int64_t reductionSize = aType.getDimSize(1);
+  problem.rhsBatchShape = {};
-  const bool transposeB = gemmOpAdaptor.getTransB();
+  problem.outputBatchShape = {};
+  problem.lhsBatch = 1;
+  problem.rhsBatch = 1;
+  problem.batch = 1;
+  problem.m = outType.getDimSize(0);
+  problem.k = aType.getDimSize(1);
+  problem.n = outType.getDimSize(1);
+  problem.origin = ContractionOrigin::Gemm;
+  problem.lhsElementType = aType.getElementType();
+  problem.rhsElementType = bType.getElementType();
+  problem.resultElementType = outType.getElementType();
+  problem.lhsTransposed = transA;
+  problem.rhsTransposed = transB;
+  problem.alpha = alpha;
+  problem.beta = beta;
+  const bool transposeB = transB;
 
   if (!isCompileTimeComputable(b)) {
+    ContractionPlan plan = makeContractionPlan(
+        problem, target, ContractionPlanKind::BatchedDynamicVVD);
     bool hasC = hasGemmBias(c);
-    float alpha = gemmOpAdaptor.getAlpha().convertToFloat();
-    float beta = gemmOpAdaptor.getBeta().convertToFloat();
     RankedTensorType biasType;
     if (hasC) {
       auto cType = dyn_cast<RankedTensorType>(c.getType());
       if (!cType || !cType.hasStaticShape()) {
-        pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm bias");
+        pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm bias");
         return failure();
       }
       auto verifiedBiasType = verifyDynamicGemmBiasType(cType, outType);
       if (failed(verifiedBiasType)) {
-        gemmOp.emitOpError("requires Gemm bias C to be broadcastable to the output shape");
+        diagnosticAnchor->emitOpError("requires Gemm bias C to be broadcastable to the output shape");
         return failure();
       }
       biasType = *verifiedBiasType;
@@ -717,19 +762,19 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
 
     const int64_t bReductionSize = bType.getDimSize(transposeB ? 1 : 0);
     const int64_t bOutputColumns = bType.getDimSize(transposeB ? 0 : 1);
-    if (aType.getDimSize(0) != numOutRows || bReductionSize != reductionSize || bOutputColumns != numOutCols) {
+    if (aType.getDimSize(0) != problem.m || bReductionSize != problem.k || bOutputColumns != problem.n) {
-      gemmOp.emitOpError("has inconsistent A, B, and output shapes");
+      diagnosticAnchor->emitOpError("has inconsistent A, B, and output shapes");
       return failure();
     }
 
-    const int64_t laneCount64 = numOutRows * numOutCols;
+    const int64_t laneCount64 = plan.laneCount;
     if (laneCount64 > std::numeric_limits<int32_t>::max()) {
-      gemmOp.emitOpError("requires Gemm dynamic batch lane count to fit in i32");
+      diagnosticAnchor->emitOpError("requires Gemm dynamic batch lane count to fit in i32");
       return failure();
     }
 
-    auto columnType = RankedTensorType::get({numOutRows, 1}, outType.getElementType());
+    auto columnType = RankedTensorType::get({problem.m, 1}, outType.getElementType());
-    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(numOutCols, columnType);
+    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(problem.n, columnType);
     auto batchOp = createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, transposeB, rewriter, loc);
     if (failed(batchOp))
       return failure();
@@ -737,94 +782,122 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
       batchOp->getResult(0), hasC ? c : Value(), scalarPiecesType, biasType, outType, alpha, beta, rewriter, loc);
     if (failed(outputCompute))
       return failure();
-    rewriter.replaceOp(gemmOp, outputCompute->getResults());
+    return outputCompute->getResult(0);
-    return success();
   }
 
   if (transposeB) {
     auto bShape = bType.getShape();
     auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType(), bType.getEncoding());
-    b = ONNXTransposeOp::create(rewriter, loc, transposedType, b, rewriter.getI64ArrayAttr({1, 0})).getResult();
+    if (isCompileTimeComputable(b)) {
+      auto transposedConstant = materializeTransposedContractionConstant(
+          b, transposedType, {1, 0}, rewriter, loc);
+      if (failed(transposedConstant)) {
+        diagnosticAnchor->emitOpError("requires Gemm input B transpose to remain statically materializable");
+        return failure();
+      }
+      b = *transposedConstant;
+    } else {
+      b = createLinalgTranspose(b, transposedType, {1, 0}, rewriter, loc);
+    }
     bType = transposedType;
   }
 
-  auto scaledB = materializeScaledConstantTensor(b, gemmOpAdaptor.getAlpha().convertToFloat(), rewriter, loc);
+  auto scaledB = materializeScaledConstantTensor(b, alpha, rewriter, loc);
   if (failed(scaledB)) {
-    gemmOp.emitOpError("requires constant Gemm input B when alpha is not 1.0");
+    diagnosticAnchor->emitOpError("requires constant Gemm input B when alpha is not 1.0");
     return failure();
   }
   b = *scaledB;
   bType = cast<RankedTensorType>(b.getType());
 
-  if (aType.getDimSize(0) != numOutRows || bType.getDimSize(0) != reductionSize || bType.getDimSize(1) != numOutCols) {
+  if (aType.getDimSize(0) != problem.m || bType.getDimSize(0) != problem.k || bType.getDimSize(1) != problem.n) {
-    gemmOp.emitOpError("has inconsistent A, B, and output shapes after transpose handling");
+    diagnosticAnchor->emitOpError("has inconsistent A, B, and output shapes after transpose handling");
     return failure();
   }
 
-  const int64_t numKSlices = ceilIntegerDivide(reductionSize, crossbarSize.getValue());
+  ContractionPlan plan = makeContractionPlan(
-  const int64_t numOutHSlices = ceilIntegerDivide(numOutCols, crossbarSize.getValue());
+      problem, target, ContractionPlanKind::StaticTiled);
-  const int64_t paddedReductionSize = numKSlices * static_cast<int64_t>(crossbarSize.getValue());
+  const int64_t xbarSize = plan.tileK;
-  const int64_t paddedOutCols = numOutHSlices * static_cast<int64_t>(crossbarSize.getValue());
+  const int64_t numKSlices = plan.reductionSlices;
+  const int64_t numOutHSlices = plan.outputTiles;
+  const int64_t paddedReductionSize = numKSlices * plan.tileK;
+  const int64_t paddedOutCols = numOutHSlices * plan.tileN;
 
   auto paddedBType = RankedTensorType::get({paddedReductionSize, paddedOutCols}, bType.getElementType());
   auto paddedB = materializePaddedConstantMatrix(b, paddedBType, rewriter, loc);
   if (failed(paddedB)) {
-    gemmOp.emitOpError("requires constant Gemm input B so tiled weights can be padded statically");
+    diagnosticAnchor->emitOpError("requires constant Gemm input B so tiled weights can be padded statically");
     return failure();
   }
   b = *paddedB;
-  auto paddedAType = RankedTensorType::get({numOutRows, paddedReductionSize}, aType.getElementType());
+  auto paddedAType = RankedTensorType::get({problem.m, paddedReductionSize}, aType.getElementType());
-  a = createPaddedInputCompute(a, paddedAType, rewriter, loc);
+  a = materializePaddedContractionInput(a, paddedAType, rewriter, loc);
   aType = paddedAType;
 
   Value bias;
   bool hasC = hasGemmBias(c);
-  auto paddedOutType = RankedTensorType::get({numOutRows, paddedOutCols}, outType.getElementType());
+  auto paddedOutType = RankedTensorType::get({problem.m, paddedOutCols}, outType.getElementType());
   if (hasC) {
     auto cType = dyn_cast<RankedTensorType>(c.getType());
     if (!cType || !cType.hasStaticShape()) {
-      pim::emitUnsupportedStaticShapeDiagnostic(gemmOp, "Gemm bias");
+      pim::emitUnsupportedStaticShapeDiagnostic(diagnosticAnchor, "Gemm bias");
       return failure();
     }
 
-    auto scaledC = materializeScaledConstantTensor(c, gemmOpAdaptor.getBeta().convertToFloat(), rewriter, loc);
+    auto scaledC = materializeScaledConstantTensor(c, beta, rewriter, loc);
     if (failed(scaledC)) {
-      gemmOp.emitOpError("requires constant Gemm bias C when beta is not 1.0");
+      diagnosticAnchor->emitOpError("requires constant Gemm bias C when beta is not 1.0");
       return failure();
     }
     c = *scaledC;
 
     auto preparedBias = prepareBias(c, outType, paddedOutType, rewriter, loc);
     if (failed(preparedBias)) {
-      gemmOp.emitOpError("requires Gemm bias C to be broadcastable to the output shape");
+      diagnosticAnchor->emitOpError("requires Gemm bias C to be broadcastable to the output shape");
       return failure();
     }
     bias = *preparedBias;
   }
 
-  const int64_t laneCount64 = numOutHSlices * numKSlices * numOutRows;
+  const int64_t laneCount64 = plan.laneCount;
   if (laneCount64 > std::numeric_limits<int32_t>::max()) {
-    gemmOp.emitOpError("requires Gemm tiled batch lane count to fit in i32");
+    diagnosticAnchor->emitOpError("requires Gemm tiled batch lane count to fit in i32");
     return failure();
   }
 
   auto partialPiecesType = spatial::getGraphBatchPhysicalResultType(
-    laneCount64, RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, outType.getElementType()));
+    laneCount64, RankedTensorType::get({1, xbarSize}, outType.getElementType()));
   auto batchOp =
-    createVmmBatch(a, b, aType, paddedBType, partialPiecesType, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
+    createVmmBatch(
+      a, b, aType, paddedBType, partialPiecesType, problem.m, numKSlices, numOutHSlices, xbarSize, rewriter, loc);
   if (failed(batchOp))
     return failure();
   auto reductionOutput = createReductionOutput(
-    batchOp->getResult(0), bias, partialPiecesType, outType, paddedOutType, numKSlices, rewriter, loc);
+    batchOp->getResult(0), bias, partialPiecesType, outType, paddedOutType, numKSlices, xbarSize, rewriter, loc);
   if (failed(reductionOutput))
     return failure();
 
-  rewriter.replaceOp(gemmOp, *reductionOutput);
+  return *reductionOutput;
+}
+
+LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
+                                                     ONNXGemmOpAdaptor gemmOpAdaptor,
+                                                     ConversionPatternRewriter& rewriter) const {
+  FailureOr<Value> result = lowerGemmToSpatial(
+      gemmOp.getOperation(), gemmOpAdaptor.getA(), gemmOpAdaptor.getB(), gemmOpAdaptor.getC(),
+      cast<RankedTensorType>(gemmOp.getY().getType()), gemmOpAdaptor.getTransA(),
+      gemmOpAdaptor.getTransB(), gemmOpAdaptor.getAlpha().convertToFloat(),
+      gemmOpAdaptor.getBeta().convertToFloat(), target, rewriter, gemmOp.getLoc());
+  if (failed(result))
+    return failure();
+  rewriter.replaceOp(gemmOp, *result);
   return success();
 }
 
-void populateGemmPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+void populateGemmPatterns(RewritePatternSet& patterns,
-  patterns.insert<GemmToSpatialComputes>(ctx);
+                          MLIRContext* ctx,
+                          const spatial::SpatialTargetInfo& target) {
+  patterns.insert<GemmToSpatialComputes>(ctx, target);
 }
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.hpp

@@ -0,0 +1,27 @@
+#pragma once
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/Value.h"
+#include "mlir/IR/PatternMatch.h"
+
+namespace onnx_mlir {
+namespace spatial {
+struct SpatialTargetInfo;
+}
+
+mlir::FailureOr<mlir::Value> lowerGemmToSpatial(
+    mlir::Operation* diagnosticAnchor,
+    mlir::Value a,
+    mlir::Value b,
+    mlir::Value c,
+    mlir::RankedTensorType outputType,
+    bool transA,
+    bool transB,
+    float alpha,
+    float beta,
+    const spatial::SpatialTargetInfo& target,
+    mlir::PatternRewriter& rewriter,
+    mlir::Location loc);
+
+} // namespace onnx_mlir

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

@@ -11,6 +11,9 @@
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionProblem.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionMaterialization.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ContractionPlanning.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -118,7 +121,7 @@ static FailureOr<Value> collapseFragmentAssemblyBatchDims(Value value,
   auto fragmentStrides = blueprint ? blueprint.getFragmentStrides() : std::nullopt;
   if (!blueprint || !inputType || !storageType || !inputType.hasStaticShape() || !storageType.hasStaticShape()
       || inputType.getRank() <= 3 || resultType.getRank() != 3 || !blueprint.getFragments().empty()
-      || blueprint.getMode() != "fragment_assembly" || !operandIndices || !sourceOffsets || !fragmentStrides
+      || !spatial::isFragmentAssembly(blueprint.getMode()) || !operandIndices || !sourceOffsets || !fragmentStrides
       || storageType.getRank() != inputType.getRank() + 1)
     return failure();
   if (blueprint.getIndexMap() == spatial::kContiguousRowMajorFragments
@@ -176,7 +179,7 @@ static FailureOr<Value> collapseFragmentAssemblyBatchDims(Value value,
                                                               collapsedStorage,
                                                               ValueRange {},
                                                               blueprint.getLogicalLayoutAttr(),
-                                                              rewriter.getStringAttr("fragmented"),
+                                                              spatial::getFragmentedLayout(rewriter.getContext()),
                                                               rewriter.getDenseI64ArrayAttr(offsets),
                                                               rewriter.getDenseI64ArrayAttr(sizes),
                                                               rewriter.getStringAttr("collapsed_fragments"),
@@ -462,6 +465,7 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVmmBatch(Value a,
                                                                   int64_t numOutRows,
                                                                   int64_t numKSlices,
                                                                   int64_t numOutHSlices,
+                                                                  int64_t xbarSize,
                                                                   PatternRewriter& rewriter,
                                                                   Location loc) {
   const int64_t laneCount = partialPiecesType.getDimSize(0);
@@ -480,16 +484,16 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVmmBatch(Value a,
       Value sliceLane = affineModConst(rewriter, loc, outerLane, numKSlices * numOutHSlices, anchorOp);
       Value kSlice = affineModConst(rewriter, loc, sliceLane, numKSlices, anchorOp);
       Value hSlice = affineFloorDivConst(rewriter, loc, sliceLane, numKSlices, anchorOp);
-      Value kOffset = affineMulConst(rewriter, loc, kSlice, crossbarSize.getValue(), anchorOp);
+      Value kOffset = affineMulConst(rewriter, loc, kSlice, xbarSize, anchorOp);
-      Value hOffset = affineMulConst(rewriter, loc, hSlice, crossbarSize.getValue(), anchorOp);
+      Value hOffset = affineMulConst(rewriter, loc, hSlice, xbarSize, anchorOp);
 
       auto aTileType =
-        RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, aType.getElementType());
+        RankedTensorType::get({1, xbarSize}, aType.getElementType());
       auto bTileType = RankedTensorType::get(
-        {static_cast<int64_t>(crossbarSize.getValue()), static_cast<int64_t>(crossbarSize.getValue())},
+        {xbarSize, xbarSize},
         bType.getElementType());
       auto pieceType =
-        RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, partialPiecesType.getElementType());
+        RankedTensorType::get({1, xbarSize}, partialPiecesType.getElementType());
 
       Value aTile = extractBatchedATile(
         args.inputs.front(), aBatchShape, outputBatchShape, batch, row, kOffset, aTileType, rewriter, loc);
@@ -522,8 +526,11 @@ static Value extractDynamicBatchedRowVector(Value matrix,
     {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
-static int64_t chooseDynamicMatMulRowsPerLane(int64_t rows, int64_t reductionSize, int64_t columns) {
+static int64_t chooseDynamicMatMulRowsPerLane(int64_t rows,
-  const int64_t crossbarElements = static_cast<int64_t>(crossbarSize.getValue() * crossbarSize.getValue());
+                                              int64_t reductionSize,
+                                              int64_t columns,
+                                              int64_t xbarSize) {
+  const int64_t crossbarElements = xbarSize * xbarSize;
   const int64_t target = std::min(rows, ceilIntegerDivide(reductionSize * columns, crossbarElements));
   int64_t rowsPerLane = 1;
   for (int64_t candidate = 2; candidate <= target; ++candidate)
@@ -678,6 +685,7 @@ static Value extractBatchedReductionPiece(Value partialPiecesArg,
                                           int64_t numKSlices,
                                           int64_t numOutHSlices,
                                           int64_t numOutRows,
+                                          int64_t xbarSize,
                                           PatternRewriter& rewriter,
                                           Location loc) {
   Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
@@ -687,7 +695,8 @@ static Value extractBatchedReductionPiece(Value partialPiecesArg,
   Value batchAndHSlice = arith::AddIOp::create(rewriter, loc, batchOffset, hOffset);
   Value pieceOffset = arith::AddIOp::create(rewriter, loc, batchAndHSlice, kOffset);
   SmallVector<OpFoldResult> offsets {pieceOffset, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
+  SmallVector<OpFoldResult> sizes {
+    rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarSize)};
   return extractMixedSliceOrIdentity(
     rewriter, loc, partialPiecesArg, pieceType,
     {offsets, sizes, getUnitStrides(rewriter, 3)});
@@ -700,13 +709,24 @@ static Value reduceBatchedPartialPiecesForHSlice(Value partialPiecesArg,
                                                  int64_t numKSlices,
                                                  int64_t numOutHSlices,
                                                  int64_t numOutRows,
+                                                 int64_t xbarSize,
                                                  PatternRewriter& rewriter,
                                                  Location loc) {
   SmallVector<Value> activePieces;
   activePieces.reserve(numKSlices);
   for (int64_t kSlice = 0; kSlice < numKSlices; ++kSlice)
     activePieces.push_back(extractBatchedReductionPiece(
-      partialPiecesArg, batch, hSlice, kSlice, pieceType, numKSlices, numOutHSlices, numOutRows, rewriter, loc));
+      partialPiecesArg,
+      batch,
+      hSlice,
+      kSlice,
+      pieceType,
+      numKSlices,
+      numOutHSlices,
+      numOutRows,
+      xbarSize,
+      rewriter,
+      loc));
 
   while (activePieces.size() > 1) {
     SmallVector<Value> nextPieces;
@@ -729,13 +749,14 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
                                                       RankedTensorType paddedOutType,
                                                       int64_t numBatches,
                                                       int64_t numKSlices,
+                                                      int64_t xbarSize,
                                                       PatternRewriter& rewriter,
                                                       Location loc) {
   auto computeOp = createSpatCompute<1>(
     rewriter, loc, TypeRange {outType}, {}, ValueRange {partialPieces}, [&](Value partialPiecesArg) -> LogicalResult {
       const int64_t numOutRows = outType.getDimSize(1);
-      const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(2), crossbarSize.getValue());
+      const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(2), xbarSize);
-      auto pieceType = RankedTensorType::get({numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
+      auto pieceType = RankedTensorType::get({numOutRows, xbarSize},
                                              partialPiecesType.getElementType());
 
       Value outputInit =
@@ -765,13 +786,22 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
             [&](OpBuilder&, Location hLoc, Value hSlice, ValueRange hIterArgs, SmallVectorImpl<Value>& hYielded) {
               Value outputAcc = hIterArgs.front();
               Value reduced = reduceBatchedPartialPiecesForHSlice(
-                partialPiecesArg, batch, hSlice, pieceType, numKSlices, numOutHSlices, numOutRows, rewriter, hLoc);
+                partialPiecesArg,
+                batch,
+                hSlice,
+                pieceType,
+                numKSlices,
+                numOutHSlices,
+                numOutRows,
+                xbarSize,
+                rewriter,
+                hLoc);
               Value hOffset = affineMulConst(
-                rewriter, hLoc, hSlice, crossbarSize.getValue(), rewriter.getInsertionBlock()->getParentOp());
+                rewriter, hLoc, hSlice, xbarSize, rewriter.getInsertionBlock()->getParentOp());
               SmallVector<OpFoldResult> outputOffsets {batch, rewriter.getIndexAttr(0), hOffset};
               SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1),
                                                      rewriter.getIndexAttr(numOutRows),
-                                                     rewriter.getIndexAttr(crossbarSize.getValue())};
+                                                     rewriter.getIndexAttr(xbarSize)};
               Value next =
                 tensor::InsertSliceOp::create(
                   rewriter, hLoc, reduced, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
@@ -805,39 +835,45 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
   return computeOp->getResult(0);
 }
 
-struct NormalizedMatMulInfo {
+struct NormalizedMatMulInfo : ContractionProblem {
+  NormalizedMatMulInfo(RankedTensorType lhsType,
+                       RankedTensorType rhsType,
+                       RankedTensorType outType,
+                       RankedTensorType normalizedLhsType,
+                       RankedTensorType normalizedRhsType,
+                       ContractionProblem problem,
+                       bool lhsWasVector,
+                       bool rhsWasVector)
+      : ContractionProblem(std::move(problem)),
+        lhsType(lhsType),
+        rhsType(rhsType),
+        outType(outType),
+        normalizedLhsType(normalizedLhsType),
+        normalizedRhsType(normalizedRhsType),
+        lhsWasVector(lhsWasVector),
+        rhsWasVector(rhsWasVector) {}
+
   RankedTensorType lhsType;
   RankedTensorType rhsType;
   RankedTensorType outType;
   RankedTensorType normalizedLhsType;
   RankedTensorType normalizedRhsType;
-  SmallVector<int64_t> lhsBatchShape;
-  SmallVector<int64_t> rhsBatchShape;
-  SmallVector<int64_t> outputBatchShape;
   bool lhsWasVector;
   bool rhsWasVector;
-  int64_t lhsBatch;
-  int64_t rhsBatch;
-  int64_t batch;
-  int64_t m;
-  int64_t k;
-  int64_t n;
 };
 
-struct MatMulLoweringPlan {
+struct MatMulLoweringPlan : ContractionProblem {
+  MatMulLoweringPlan(Value lhs, Value rhs, const NormalizedMatMulInfo& info)
+      : ContractionProblem(info),
+        lhs(lhs),
+        rhs(rhs),
+        lhsType(cast<RankedTensorType>(lhs.getType())),
+        rhsType(cast<RankedTensorType>(rhs.getType())) {}
+
   Value lhs;
   Value rhs;
   RankedTensorType lhsType;
   RankedTensorType rhsType;
-  SmallVector<int64_t> lhsBatchShape;
-  SmallVector<int64_t> rhsBatchShape;
-  SmallVector<int64_t> outputBatchShape;
-  int64_t lhsBatch;
-  int64_t rhsBatch;
-  int64_t batch;
-  int64_t m;
-  int64_t k;
-  int64_t n;
   bool transposedResult;
 };
 
@@ -901,22 +937,31 @@ static FailureOr<NormalizedMatMulInfo> analyzeMatMulShape(ONNXMatMulOp matmulOp)
     return failure();
   }
 
-  return NormalizedMatMulInfo {lhsType,
+  return NormalizedMatMulInfo(
-                               rhsType,
+    lhsType,
-                               outType,
+    rhsType,
-                               normalizedLhsType,
+    outType,
-                               normalizedRhsType,
+    normalizedLhsType,
-                               lhsBatchShape,
+    normalizedRhsType,
-                               rhsBatchShape,
+    ContractionProblem {lhsBatchShape,
-                               *outputBatchShape,
+                        rhsBatchShape,
-                               lhsWasVector,
+                        *outputBatchShape,
-                               rhsWasVector,
+                        lhsBatch,
-                               lhsBatch,
+                        rhsBatch,
-                               rhsBatch,
+                        batch,
-                               batch,
+                        m,
-                               m,
+                        k,
-                               k,
+                        n,
-                               n};
+                        ContractionOrigin::MatMul,
+                        lhsType.getElementType(),
+                        rhsType.getElementType(),
+                        outType.getElementType(),
+                        false,
+                        false,
+                        lhsWasVector,
+                        rhsWasVector},
+    lhsWasVector,
+    rhsWasVector);
 }
 
 static MatMulLoweringPlan buildLoweringPlan(Value normalizedLhs,
@@ -925,20 +970,8 @@ static MatMulLoweringPlan buildLoweringPlan(Value normalizedLhs,
                                             bool useTransposedForm,
                                             PatternRewriter& rewriter,
                                             Location loc) {
-  MatMulLoweringPlan plan {normalizedLhs,
+  MatMulLoweringPlan plan(normalizedLhs, normalizedRhs, info);
-                           normalizedRhs,
+  plan.transposedResult = false;
-                           cast<RankedTensorType>(normalizedLhs.getType()),
-                           cast<RankedTensorType>(normalizedRhs.getType()),
-                           info.lhsBatchShape,
-                           info.rhsBatchShape,
-                           info.outputBatchShape,
-                           info.lhsBatch,
-                           info.rhsBatch,
-                           info.batch,
-                           info.m,
-                           info.k,
-                           info.n,
-                           false};
   if (!useTransposedForm)
     return plan;
 
@@ -1057,7 +1090,9 @@ struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
 };
 
 struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
-  using OpRewritePattern::OpRewritePattern;
+  explicit MatMulBatchedToSpatialComputes(MLIRContext* ctx,
+                                           const spatial::SpatialTargetInfo& target)
+      : OpRewritePattern<ONNXMatMulOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(ONNXMatMulOp matmulOp, PatternRewriter& rewriter) const override {
     auto shapeInfo = analyzeMatMulShape(matmulOp);
@@ -1067,6 +1102,7 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
       return failure();
 
     Location loc = matmulOp.getLoc();
+    const int64_t xbarSize = static_cast<int64_t>(target.matrixShape.rows);
     bool useTransposedForm = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector
                           && isCompileTimeComputable(matmulOp.getA()) && !isCompileTimeComputable(matmulOp.getB());
     Value rhsRows = getLastTwoTransposeInput(matmulOp.getB());
@@ -1088,7 +1124,8 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
                             rhsStoredAsRows ? shapeInfo->k : shapeInfo->n,
                             rewriter,
                             loc);
-    MatMulLoweringPlan plan = buildLoweringPlan(lhs, rhs, *shapeInfo, useTransposedForm, rewriter, loc);
+    MatMulLoweringPlan plan = buildLoweringPlan(
+      lhs, rhs, *shapeInfo, useTransposedForm, rewriter, loc);
 
     plan.lhs = ensureBatchedTensor(plan.lhs, plan.lhsBatch, plan.m, plan.k, rewriter, loc);
     plan.rhs = ensureBatchedTensor(plan.rhs,
@@ -1103,10 +1140,12 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
       {plan.batch, plan.m, plan.n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
 
     if (isCompileTimeComputable(plan.rhs)) {
-      const int64_t numKSlices = ceilIntegerDivide(plan.k, crossbarSize.getValue());
+      ContractionPlan contractionPlan = makeContractionPlan(
-      const int64_t numOutHSlices = ceilIntegerDivide(plan.n, crossbarSize.getValue());
+        plan, target, ContractionPlanKind::StaticTiled);
-      const int64_t paddedReductionSize = numKSlices * static_cast<int64_t>(crossbarSize.getValue());
+      const int64_t numKSlices = contractionPlan.reductionSlices;
-      const int64_t paddedOutCols = numOutHSlices * static_cast<int64_t>(crossbarSize.getValue());
+      const int64_t numOutHSlices = contractionPlan.outputTiles;
+      const int64_t paddedReductionSize = numKSlices * xbarSize;
+      const int64_t paddedOutCols = numOutHSlices * xbarSize;
       auto paddedLhsType = RankedTensorType::get(
         {plan.lhsBatch, plan.m, paddedReductionSize}, plan.lhsType.getElementType(), plan.lhsType.getEncoding());
       auto paddedRhsType = RankedTensorType::get(
@@ -1117,10 +1156,11 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
       auto paddedRhs =
         materializePaddedBatchedWeight(plan.rhs, plan.rhsBatchShape, plan.outputBatchShape, paddedRhsType, rewriter);
       if (succeeded(paddedRhs)) {
-        Value paddedLhs = createPaddedInputCompute(plan.lhs, paddedLhsType, rewriter, loc);
+        Value paddedLhs = materializePaddedContractionInput(
-        const int64_t laneCount = plan.batch * plan.m * numKSlices * numOutHSlices;
+          plan.lhs, paddedLhsType, rewriter, loc);
+        const int64_t laneCount = contractionPlan.laneCount;
         auto partialPiecesType = spatial::getGraphBatchPhysicalResultType(
-          laneCount, RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, shapeInfo->outType.getElementType()));
+          laneCount, RankedTensorType::get({1, xbarSize}, shapeInfo->outType.getElementType()));
         auto batchOp = createBatchedVmmBatch(paddedLhs,
                                              *paddedRhs,
                                              paddedLhsType,
@@ -1132,6 +1172,7 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
                                              plan.m,
                                              numKSlices,
                                              numOutHSlices,
+                                             xbarSize,
                                              rewriter,
                                              loc);
         if (failed(batchOp))
@@ -1142,6 +1183,7 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
                                                     paddedOutType,
                                                     plan.batch,
                                                     numKSlices,
+                                                    xbarSize,
                                                     rewriter,
                                                     loc);
         if (failed(result))
@@ -1164,8 +1206,11 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
                                    ? shapeInfo->outType : directOutType;
     SmallVector<int64_t> blueprintBatchShape = !shapeInfo->lhsWasVector && !shapeInfo->rhsWasVector
                                              ? shapeInfo->outputBatchShape : SmallVector<int64_t> {plan.batch};
-    const int64_t rowsPerLane = chooseDynamicMatMulRowsPerLane(plan.m, plan.k, plan.n);
+    const int64_t rowsPerLane = chooseDynamicMatMulRowsPerLane(plan.m, plan.k, plan.n, xbarSize);
-    const int64_t laneCount = plan.batch * plan.m / rowsPerLane;
+    ContractionPlan contractionPlan = makeContractionPlan(
+      plan, target, ContractionPlanKind::GroupedRowDynamicVVD,
+      /*laneCount=*/plan.batch * plan.m / rowsPerLane, rowsPerLane);
+    const int64_t laneCount = contractionPlan.laneCount;
     SmallVector<int64_t> fragmentShape(blueprintType.getRank(), 1);
     fragmentShape[fragmentShape.size() - 2] = rowsPerLane;
     fragmentShape.back() = plan.n;
@@ -1210,6 +1255,8 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
       rewriter.eraseOp(foldedMultiply);
     return success();
   }
+
+  const spatial::SpatialTargetInfo& target;
 };
 
 struct TransposedRhsMatMulToSpatial : MatMulBatchedToSpatialComputes {
@@ -1224,12 +1271,17 @@ struct TransposedRhsMatMulToSpatial : MatMulBatchedToSpatialComputes {
 
 } // namespace
 
-void populateMatMulFusionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+void populateMatMulFusionPatterns(RewritePatternSet& patterns,
-  patterns.add<TransposedRhsMatMulToSpatial>(ctx);
+                                  MLIRContext* ctx,
+                                  const spatial::SpatialTargetInfo& target) {
+  patterns.add<TransposedRhsMatMulToSpatial>(ctx, target);
 }
 
-void populateMatMulRewritePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+void populateMatMulRewritePatterns(RewritePatternSet& patterns,
-  patterns.insert<MatMulToGemm, MatMulBatchedToSpatialComputes>(ctx);
+                                   MLIRContext* ctx,
+                                   const spatial::SpatialTargetInfo& target) {
+  patterns.insert<MatMulToGemm>(ctx);
+  patterns.insert<MatMulBatchedToSpatialComputes>(ctx, target);
 }
 
 } // namespace onnx_mlir

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

@@ -280,12 +280,12 @@ static FailureOr<Value> buildReduceMeanKeepdimsBlueprint(
   SmallVector<int64_t> fragmentStrides(fragmentOffsets.size(), 1);
   return spatial::SpatBlueprintOp::create(
            rewriter, loc, keepdimsType, batchValue, ValueRange {},
-           rewriter.getStringAttr("nchw"),
+           spatial::getNCHWLayout(rewriter.getContext()),
-           rewriter.getStringAttr("fragmented"),
+           spatial::getFragmentedLayout(rewriter.getContext()),
            rewriter.getDenseI64ArrayAttr(fragmentOffsets),
            rewriter.getDenseI64ArrayAttr(fragmentSizes),
            rewriter.getStringAttr("reduce_mean_keepdims_fragments"),
-           rewriter.getStringAttr("fragment_assembly"),
+           spatial::getFragmentAssemblyMode(rewriter.getContext()),
            rewriter.getDenseI64ArrayAttr(operandIndices),
            rewriter.getDenseI64ArrayAttr(sourceSlots),
            rewriter.getDenseI64ArrayAttr(sourceOffsets),

src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp

@@ -14,8 +14,8 @@
 
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.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"
@@ -32,8 +32,10 @@ static Value materializeTileTensor(PatternRewriter& rewriter, Location loc, Valu
   return insertStaticSlice(rewriter, loc, tile, empty, getZeroOffsets(rewriter, tileType.getRank()));
 }
 
-static Value
+static Value createPoolFillElement(OpBuilder& rewriter,
-createPoolFillElement(ConversionPatternRewriter& rewriter, Location loc, Type elementType, bool useMinimumValue) {
+                                   Location loc,
+                                   Type elementType,
+                                   bool useMinimumValue) {
   Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
   if (!useMinimumValue)
     return getOrCreateConstant(rewriter, anchorOp, rewriter.getZeroAttr(elementType), elementType);
@@ -51,7 +53,7 @@ createPoolFillElement(ConversionPatternRewriter& rewriter, Location loc, Type el
   llvm_unreachable("unsupported pool element type");
 }
 
-static Value createPoolFillTensor(ConversionPatternRewriter& rewriter,
+static Value createPoolFillTensor(OpBuilder& rewriter,
                                   Location loc,
                                   RankedTensorType tensorType,
                                   bool useMinimumValue) {
@@ -59,16 +61,15 @@ static Value createPoolFillTensor(ConversionPatternRewriter& rewriter,
   return tensor::SplatOp::create(rewriter, loc, tensorType, fillElement);
 }
 
-template <typename PoolOp>
+static Value createPaddedPoolInput(OpBuilder& rewriter,
-static Value createPaddedPoolInput(ConversionPatternRewriter& rewriter,
                                    Location loc,
-                                   PoolOp poolOp,
                                    Value input,
                                    RankedTensorType inputType,
                                    int64_t padTop,
                                    int64_t padLeft,
                                    int64_t padBottom,
-                                   int64_t padRight) {
+                                   int64_t padRight,
+                                   bool useMinimumValue) {
   if (padTop == 0 && padLeft == 0 && padBottom == 0 && padRight == 0)
     return input;
 
@@ -90,8 +91,8 @@ static Value createPaddedPoolInput(ConversionPatternRewriter& rewriter,
     padBlock->addArgument(rewriter.getIndexType(), loc);
   padOp.getRegion().push_back(padBlock);
   rewriter.setInsertionPointToStart(padBlock);
-  Value padValue =
+  Value padValue = createPoolFillElement(
-    createPoolFillElement(rewriter, loc, inputType.getElementType(), std::is_same_v<PoolOp, ONNXMaxPoolSingleOutOp>);
+    rewriter, loc, inputType.getElementType(), useMinimumValue);
   tensor::YieldOp::create(rewriter, loc, padValue);
   rewriter.setInsertionPointAfter(padOp);
   return padOp.getResult();
@@ -160,7 +161,10 @@ struct PoolToSpatialCompute;
 
 template <typename PoolOp, typename PoolOpAdaptor, typename ReduceOp>
 struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
-  using OpConversionPattern<PoolOp>::OpConversionPattern;
+  PoolToSpatialComputeBase(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
+      : OpConversionPattern<PoolOp>(ctx), target(target) {}
+
+  const spatial::SpatialTargetInfo& target;
 
   LogicalResult matchAndRewrite(PoolOp poolOp, PoolOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const final {
     Location loc = poolOp.getLoc();
@@ -241,7 +245,7 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
           rewriter.getDenseI64ArrayAttr({padTop, padLeft, padBottom, padRight}),
           rewriter.getDenseI64ArrayAttr({strideHeight, strideWidth}),
           rewriter.getDenseI64ArrayAttr({dilationHeight, dilationWidth}),
-          rewriter.getStringAttr("nchw"));
+          spatial::getNCHWLayout(rewriter.getContext()));
         rewriter.replaceOp(poolOp, plan.getResult());
         return success();
       }
@@ -251,12 +255,12 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
              && dilationHeight == 1 && dilationWidth == 1 && padTop == 0
              && padLeft == 0 && padBottom == 0 && padRight == 0) {
       auto plan = spatial::SpatGlobalAveragePoolPlanOp::create(
-        rewriter, loc, outType, x, rewriter.getStringAttr("nchw"));
+        rewriter, loc, outType, x, spatial::getNCHWLayout(rewriter.getContext()));
       rewriter.replaceOp(poolOp, plan.getResult());
       return success();
     }
 
-    const int64_t xbarSize = static_cast<int64_t>(crossbarSize.getValue());
+    const int64_t xbarSize = static_cast<int64_t>(target.matrixShape.rows);
     const int64_t channelTileCount = (channels + xbarSize - 1) / xbarSize;
     const int64_t outputPatchCount = batchSize * outputHeight * outputWidth;
     const bool countIncludePad = [&]() {
@@ -292,7 +296,9 @@ struct PoolToSpatialComputeBase : public OpConversionPattern<PoolOp> {
     auto computeOp =
       createSpatCompute<numInputs>(rewriter, loc, outType, {}, ValueRange {x}, [&](Value xArg) -> LogicalResult {
         Value paddedInput =
-          createPaddedPoolInput(rewriter, loc, poolOp, xArg, xType, padTop, padLeft, padBottom, padRight);
+          createPaddedPoolInput(rewriter, loc, xArg, xType, padTop, padLeft,
+                                padBottom, padRight,
+                                std::is_same_v<PoolOp, ONNXMaxPoolSingleOutOp>);
         Value pooledOutputInit = tensor::EmptyOp::create(rewriter, loc, outType.getShape(), outType.getElementType());
 
         Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
@@ -424,7 +430,8 @@ struct PoolToSpatialCompute<ONNXAveragePoolOp>
 
 } // namespace
 
-LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp) {
+LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
+                                            const spatial::SpatialTargetInfo&) {
   auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape())
@@ -439,6 +446,118 @@ LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp)
   return success();
 }
 
+FailureOr<Value> lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                                         const spatial::SpatialTargetInfo& target,
+                                         PatternRewriter& rewriter) {
+  auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
+  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape()
+      || inputType.getRank() != 4 || outputType.getRank() != 4)
+    return planOp.emitOpError("dense MaxPool lowering requires static rank-4 tensors"), failure();
+
+  auto kernel = planOp.getKernelShape();
+  auto pads = planOp.getPads();
+  auto strides = planOp.getStrides();
+  auto dilations = planOp.getDilations();
+  if (kernel.size() != 2 || pads.size() != 4 || strides.size() != 2 || dilations.size() != 2
+      || llvm::any_of(kernel, [](int64_t value) { return value <= 0; })
+      || llvm::any_of(strides, [](int64_t value) { return value <= 0; })
+      || llvm::any_of(dilations, [](int64_t value) { return value <= 0; })
+      || llvm::any_of(pads, [](int64_t value) { return value < 0; }))
+    return planOp.emitOpError("dense MaxPool lowering requires valid kernel, padding, stride, and dilation attributes"),
+           failure();
+
+  const int64_t batchSize = inputType.getDimSize(0);
+  const int64_t channels = inputType.getDimSize(1);
+  const int64_t outputHeight = outputType.getDimSize(2);
+  const int64_t outputWidth = outputType.getDimSize(3);
+  const int64_t tileWidth = std::max<int64_t>(1, target.matrixShape.rows);
+  const int64_t channelTileCount = (channels + tileWidth - 1) / tileWidth;
+  const int64_t outputPatchCount = batchSize * outputHeight * outputWidth;
+
+  auto compute = createSpatCompute<1>(
+    rewriter, planOp.getLoc(), outputType, {}, planOp.getInput(),
+    [&](Value input) -> LogicalResult {
+      Value paddedInput = createPaddedPoolInput(
+        rewriter, planOp.getLoc(), input, inputType,
+        pads[0], pads[1], pads[2], pads[3], /*useMinimumValue=*/true);
+      Value outputInit = tensor::EmptyOp::create(
+        rewriter, planOp.getLoc(), outputType.getShape(), outputType.getElementType());
+      Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
+      Value zero = getOrCreateIndexConstant(rewriter, anchor, 0);
+      Value one = getOrCreateIndexConstant(rewriter, anchor, 1);
+      Value patchCount = getOrCreateIndexConstant(rewriter, anchor, outputPatchCount);
+      Value pixelsPerBatch = getOrCreateIndexConstant(
+        rewriter, anchor, outputHeight * outputWidth);
+      Value outputWidthValue = getOrCreateIndexConstant(rewriter, anchor, outputWidth);
+      Value strideHeight = getOrCreateIndexConstant(rewriter, anchor, strides[0]);
+      Value strideWidth = getOrCreateIndexConstant(rewriter, anchor, strides[1]);
+
+      auto loop = buildNormalizedScfFor(
+        rewriter, planOp.getLoc(), zero, patchCount, one, ValueRange {outputInit},
+        [&](OpBuilder&, Location loc, Value patch, ValueRange iterArgs,
+            SmallVectorImpl<Value>& yielded) {
+          Value batch = arith::DivUIOp::create(rewriter, loc, patch, pixelsPerBatch);
+          Value batchPatch = arith::RemUIOp::create(rewriter, loc, patch, pixelsPerBatch);
+          Value outputRow = arith::DivUIOp::create(rewriter, loc, batchPatch, outputWidthValue);
+          Value outputColumn = arith::RemUIOp::create(rewriter, loc, batchPatch, outputWidthValue);
+          Value windowRow = arith::MulIOp::create(rewriter, loc, outputRow, strideHeight);
+          Value windowColumn = arith::MulIOp::create(rewriter, loc, outputColumn, strideWidth);
+          Value updated = iterArgs.front();
+
+          for (int64_t tile = 0; tile < channelTileCount; ++tile) {
+            const int64_t tileChannels = std::min<int64_t>(tileWidth, channels - tile * tileWidth);
+            auto tileType = RankedTensorType::get(
+              {1, tileChannels, 1, 1}, outputType.getElementType());
+            Value reduced = createPoolFillTensor(
+              rewriter, loc, tileType, /*useMinimumValue=*/true);
+            for (int64_t kernelRow = 0; kernelRow < kernel[0]; ++kernelRow) {
+              Value sourceRow = windowRow;
+              if (kernelRow * dilations[0] != 0)
+                sourceRow = arith::AddIOp::create(
+                  rewriter, loc, sourceRow,
+                  getOrCreateIndexConstant(rewriter, anchor, kernelRow * dilations[0]));
+              for (int64_t kernelColumn = 0; kernelColumn < kernel[1]; ++kernelColumn) {
+                Value sourceColumn = windowColumn;
+                if (kernelColumn * dilations[1] != 0)
+                  sourceColumn = arith::AddIOp::create(
+                    rewriter, loc, sourceColumn,
+                    getOrCreateIndexConstant(rewriter, anchor, kernelColumn * dilations[1]));
+                Value point = tensor::ExtractSliceOp::create(
+                  rewriter, loc, tileType, paddedInput,
+                  SmallVector<OpFoldResult> {
+                    batch, rewriter.getIndexAttr(tile * tileWidth), sourceRow, sourceColumn},
+                  SmallVector<OpFoldResult> {
+                    rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels),
+                    rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)},
+                  getUnitStrides(rewriter, 4));
+                point = materializeTileTensor(rewriter, loc, point);
+                reduced = spatial::SpatVMaxOp::create(
+                  rewriter, loc, tileType, reduced, point);
+              }
+            }
+            updated = tensor::InsertSliceOp::create(
+              rewriter, loc, reduced, updated,
+              SmallVector<OpFoldResult> {
+                batch, rewriter.getIndexAttr(tile * tileWidth), outputRow, outputColumn},
+              SmallVector<OpFoldResult> {
+                rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels),
+                rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)},
+              getUnitStrides(rewriter, 4));
+          }
+          yielded.push_back(updated);
+          return success();
+        });
+      if (failed(loop))
+        return failure();
+      spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), loop->results.front());
+      return success();
+    });
+  if (failed(compute))
+    return failure();
+  return compute->getResult(0);
+}
+
 static Value createClampedPoolIndexTable(PatternRewriter& rewriter,
                                          Operation* anchorOp,
                                          int64_t outputSize,
@@ -497,8 +616,9 @@ static Value extractPoolIndex(PatternRewriter& rewriter,
 
 FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
                                             std::optional<Value> rowStripInput,
+                                            const spatial::SpatialTargetInfo& target,
                                             PatternRewriter& rewriter) {
-  if (failed(canLowerMaxPoolPlanToRowStrip(planOp)))
+  if (failed(canLowerMaxPoolPlanToRowStrip(planOp, target)))
     return failure();
 
   Location loc = planOp.getLoc();
@@ -590,8 +710,8 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
                                                                                 rewriter.getIndexAttr(1),
                                                                                 rewriter.getIndexAttr(inputWidth)},
                                                      getUnitStrides(rewriter, 4));
-          inputRows.push_back(ONNXTransposeOp::create(
+          inputRows.push_back(createLinalgTranspose(
-            rewriter, loc, inputFragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1})));
+            nchw, inputFragmentType, {0, 2, 3, 1}, rewriter, loc));
         }
       }
 
@@ -685,7 +805,8 @@ FailureOr<Value> lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
   return batch->getResult(0);
 }
 
-LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp) {
+LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(
+  spatial::SpatGlobalAveragePoolPlanOp planOp, const spatial::SpatialTargetInfo&) {
   auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape())
@@ -697,10 +818,87 @@ LogicalResult canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAverage
   return success();
 }
 
+FailureOr<Value> lowerDenseGlobalAveragePoolPlan(
+    spatial::SpatGlobalAveragePoolPlanOp planOp,
+    const spatial::SpatialTargetInfo& target,
+    PatternRewriter& rewriter) {
+  auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
+  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+  if (!inputType || !outputType || !inputType.hasStaticShape()
+      || !outputType.hasStaticShape() || inputType.getRank() != 4
+      || outputType.getRank() != 4 || inputType.getDimSize(0) != 1
+      || outputType.getDimSize(0) != 1 || inputType.getDimSize(1) != outputType.getDimSize(1)
+      || outputType.getDimSize(2) != 1 || outputType.getDimSize(3) != 1)
+    return planOp.emitOpError("dense global AveragePool lowering requires static rank-4 floating-point tensors"),
+           failure();
+  auto elementType = dyn_cast<FloatType>(inputType.getElementType());
+  if (!elementType)
+    return planOp.emitOpError("dense global AveragePool lowering requires floating-point tensors"),
+           failure();
+
+  const int64_t channels = inputType.getDimSize(1);
+  const int64_t height = inputType.getDimSize(2);
+  const int64_t width = inputType.getDimSize(3);
+  const int64_t tileWidth = std::max<int64_t>(1, target.matrixShape.rows);
+  const int64_t channelTileCount = (channels + tileWidth - 1) / tileWidth;
+  const double scaleValue = 1.0 / static_cast<double>(height * width);
+
+  auto compute = createSpatCompute<1>(
+    rewriter, planOp.getLoc(), outputType, {}, planOp.getInput(),
+    [&](Value input) -> LogicalResult {
+      Value output = tensor::EmptyOp::create(
+        rewriter, planOp.getLoc(), outputType.getShape(), outputType.getElementType());
+      Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
+      for (int64_t tile = 0; tile < channelTileCount; ++tile) {
+        const int64_t tileChannels = std::min<int64_t>(tileWidth, channels - tile * tileWidth);
+        auto tileType = RankedTensorType::get(
+          {1, tileChannels, 1, 1}, outputType.getElementType());
+        Value reduced = createPoolFillTensor(
+          rewriter, planOp.getLoc(), tileType, /*useMinimumValue=*/false);
+        for (int64_t row = 0; row < height; ++row) {
+          for (int64_t column = 0; column < width; ++column) {
+            Value point = tensor::ExtractSliceOp::create(
+              rewriter, planOp.getLoc(), tileType, input,
+              SmallVector<OpFoldResult> {
+                rewriter.getIndexAttr(0), rewriter.getIndexAttr(tile * tileWidth),
+                rewriter.getIndexAttr(row), rewriter.getIndexAttr(column)},
+              SmallVector<OpFoldResult> {
+                rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels),
+                rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)},
+              getUnitStrides(rewriter, 4));
+            point = materializeTileTensor(rewriter, planOp.getLoc(), point);
+            reduced = spatial::SpatVAddOp::create(
+              rewriter, planOp.getLoc(), tileType, reduced, point);
+          }
+        }
+        auto scaleAttr = DenseElementsAttr::get(
+          tileType, rewriter.getFloatAttr(elementType, scaleValue));
+        Value scale = getOrCreateConstant(rewriter, anchor, scaleAttr, tileType);
+        reduced = spatial::SpatVMulOp::create(
+          rewriter, planOp.getLoc(), tileType, reduced, scale);
+        output = tensor::InsertSliceOp::create(
+          rewriter, planOp.getLoc(), reduced, output,
+          SmallVector<OpFoldResult> {
+            rewriter.getIndexAttr(0), rewriter.getIndexAttr(tile * tileWidth),
+            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)},
+          SmallVector<OpFoldResult> {
+            rewriter.getIndexAttr(1), rewriter.getIndexAttr(tileChannels),
+            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)},
+          getUnitStrides(rewriter, 4));
+      }
+      spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), output);
+      return success();
+    });
+  if (failed(compute))
+    return failure();
+  return compute->getResult(0);
+}
+
 FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
                                                     std::optional<Value> rowStripInput,
+                                                    const spatial::SpatialTargetInfo& target,
                                                     PatternRewriter& rewriter) {
-  if (failed(canLowerGlobalAveragePoolPlanToRowStrip(planOp)))
+  if (failed(canLowerGlobalAveragePoolPlanToRowStrip(planOp, target)))
     return failure();
 
   Location loc = planOp.getLoc();
@@ -777,8 +975,8 @@ FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo
                                       rewriter.getIndexAttr(1),
                                       rewriter.getIndexAttr(width)},
             getUnitStrides(rewriter, 4));
-          fragment = ONNXTransposeOp::create(
+          fragment = createLinalgTranspose(
-            rewriter, loc, inputFragmentType, nchw, rewriter.getI64ArrayAttr({0, 2, 3, 1}));
+            nchw, inputFragmentType, {0, 2, 3, 1}, rewriter, loc);
         }
         for (int64_t column = 0; column < width; ++column) {
           Value point = tensor::ExtractSliceOp::create(
@@ -811,9 +1009,11 @@ FailureOr<Value> lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePo
   return batch->getResult(0);
 }
 
-void populatePoolPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+void populatePoolPatterns(RewritePatternSet& patterns,
-  patterns.insert<PoolToSpatialCompute<ONNXMaxPoolSingleOutOp>>(ctx);
+                          MLIRContext* ctx,
-  patterns.insert<PoolToSpatialCompute<ONNXAveragePoolOp>>(ctx);
+                          const spatial::SpatialTargetInfo& target) {
+  patterns.insert<PoolToSpatialCompute<ONNXMaxPoolSingleOutOp>>(ctx, target);
+  patterns.insert<PoolToSpatialCompute<ONNXAveragePoolOp>>(ctx, target);
 }
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Relu.cpp

@@ -17,7 +17,7 @@ struct ReluToSpatialCompute : OpConversionPattern<ONNXReluOp> {
     Location loc = reluOp.getLoc();
     Type resultType = reluOp.getResult().getType();
     auto reluPlan = spatial::SpatReluPlanOp::create(
-      rewriter, loc, resultType, adaptor.getX(), rewriter.getStringAttr("nchw"));
+      rewriter, loc, resultType, adaptor.getX(), spatial::getNCHWLayout(rewriter.getContext()));
     rewriter.replaceOp(reluOp, reluPlan.getResult());
     return success();
   }

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Concat.cpp

@@ -32,7 +32,8 @@ struct Concat : public OpConversionPattern<ONNXConcatOp> {
              return type && type.hasStaticShape() && type.getRank() == 4;
            })) {
       rewriter.replaceOpWithNewOp<spatial::SpatConcatPlanOp>(
-        maxpoolOp, resultType, inputs, rewriter.getI64IntegerAttr(axis), rewriter.getStringAttr("nchw"));
+        maxpoolOp, resultType, inputs, rewriter.getI64IntegerAttr(axis),
+        spatial::getNCHWLayout(rewriter.getContext()));
       return success();
     }
 

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Flatten.cpp

@@ -4,7 +4,6 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
@@ -119,7 +118,8 @@ struct RowStripFlattenAnalysis {
   DenseElementsAttr weight;
 };
 
-static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(spatial::SpatGraphCompute flattenOp) {
+static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(
+  spatial::SpatGraphCompute flattenOp, const spatial::SpatialTargetInfo& target) {
   if (flattenOp.getWeights().size() != 0 || flattenOp.getInputs().size() != 1
       || flattenOp.getOutputs().size() != 1)
     return failure();
@@ -130,7 +130,7 @@ static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(spatial::SpatGr
       || resultType.getDimSize(0) != 1 || resultType.getDimSize(1) != sourceType.getNumElements())
     return failure();
   const int64_t channels = sourceType.getDimSize(1);
-  const int64_t xbarDim = static_cast<int64_t>(crossbarSize.getValue());
+  const int64_t xbarDim = static_cast<int64_t>(target.matrixShape.rows);
   if (channels > xbarDim && channels % xbarDim != 0)
     return failure();
 
@@ -162,14 +162,16 @@ static FailureOr<RowStripFlattenAnalysis> analyzeRowStripFlatten(spatial::SpatGr
 
 void populateFlattenPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Flatten>(ctx); }
 
-LogicalResult canLowerFlattenFromRowStrip(spatial::SpatGraphCompute flattenOp) {
+LogicalResult canLowerFlattenFromRowStrip(spatial::SpatGraphCompute flattenOp,
-  return succeeded(analyzeRowStripFlatten(flattenOp)) ? success() : failure();
+                                          const spatial::SpatialTargetInfo& target) {
+  return succeeded(analyzeRowStripFlatten(flattenOp, target)) ? success() : failure();
 }
 
 LogicalResult lowerFlattenFromRowStrip(const RowStripPhysicalValue& input,
                                        spatial::SpatGraphCompute flattenOp,
+                                       const spatial::SpatialTargetInfo& target,
                                        PatternRewriter& rewriter) {
-  FailureOr<RowStripFlattenAnalysis> analysis = analyzeRowStripFlatten(flattenOp);
+  FailureOr<RowStripFlattenAnalysis> analysis = analyzeRowStripFlatten(flattenOp, target);
   if (failed(analysis))
     return failure();
   auto storageType = dyn_cast<RankedTensorType>(input.storage.getType());

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Resize.cpp

@@ -183,7 +183,7 @@ struct Resize : OpConversionPattern<ONNXResizeOp> {
       return rewriter.notifyMatchFailure(resizeOp, "resize lowering requires positive static dimensions.");
 
     auto plan = spatial::SpatResizeNearestPlanOp::create(
-      rewriter, resizeOp.getLoc(), resultType, adaptor.getX(), rewriter.getStringAttr("nchw"));
+      rewriter, resizeOp.getLoc(), resultType, adaptor.getX(), spatial::getNCHWLayout(rewriter.getContext()));
     rewriter.replaceOp(resizeOp, plan.getResult());
     return success();
   }
@@ -192,7 +192,8 @@ struct Resize : OpConversionPattern<ONNXResizeOp> {
 } // namespace
 
 LogicalResult canLowerResizeNearestPlanToRowStrip(
-    spatial::SpatResizeNearestPlanOp planOp) {
+    spatial::SpatResizeNearestPlanOp planOp,
+    const spatial::SpatialTargetInfo&) {
   auto inputType = dyn_cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
   return success(inputType && outputType && inputType.hasStaticShape()
@@ -204,6 +205,7 @@ LogicalResult canLowerResizeNearestPlanToRowStrip(
 
 FailureOr<Value> lowerSelectedResizeNearestPlan(
     spatial::SpatResizeNearestPlanOp planOp, std::optional<Value> rowStripInput,
+    const spatial::SpatialTargetInfo&,
     PatternRewriter& rewriter) {
   auto inputType = cast<RankedTensorType>(planOp.getInput().getType());
   auto outputType = cast<RankedTensorType>(planOp.getOutput().getType());

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Transpose.cpp

@@ -70,7 +70,7 @@ static FailureOr<Value> transposeFragmentAssemblyBlueprint(spatial::SpatBlueprin
   auto sourceOffsets = blueprint.getFragmentSourceOffsets();
   auto fragmentStrides = blueprint.getFragmentStrides();
   if (!storageType || !storageType.hasStaticShape() || !resultType.hasStaticShape()
-      || !blueprint.getFragments().empty() || blueprint.getMode() != "fragment_assembly"
+      || !blueprint.getFragments().empty() || !spatial::isFragmentAssembly(blueprint.getMode())
       || !blueprint.getFragmentOperandIndices() || !sourceOffsets || !fragmentStrides
       || llvm::any_of(*sourceOffsets, [](int64_t offset) { return offset != 0; })
       || storageType.getRank() != resultType.getRank() + 1)
@@ -113,7 +113,7 @@ static FailureOr<Value> transposeFragmentAssemblyBlueprint(spatial::SpatBlueprin
                                                                *mapped,
                                                                ValueRange {},
                                                                blueprint.getLogicalLayoutAttr(),
-                                                               rewriter.getStringAttr("fragmented"),
+                                                               spatial::getFragmentedLayout(rewriter.getContext()),
                                                                rewriter.getDenseI64ArrayAttr(offsets),
                                                                rewriter.getDenseI64ArrayAttr(sizes),
                                                                rewriter.getStringAttr("permuted_fragments"),

src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp

@@ -15,38 +15,50 @@ mlir::FailureOr<mlir::Value>
 lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
                         std::optional<mlir::Value> rowStripInput,
                         bool emitRowStripLayout,
+                        const spatial::SpatialTargetInfo& target,
                         mlir::PatternRewriter& rewriter);
 
-mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp);
+mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp,
-mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp);
+                                               const spatial::SpatialTargetInfo& target);
+mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp,
+                                                 const spatial::SpatialTargetInfo& target);
 
 mlir::LogicalResult canLowerResizeNearestPlanToRowStrip(
-  spatial::SpatResizeNearestPlanOp planOp);
+  spatial::SpatResizeNearestPlanOp planOp, const spatial::SpatialTargetInfo& target);
 
 mlir::FailureOr<mlir::Value> lowerSelectedResizeNearestPlan(
   spatial::SpatResizeNearestPlanOp planOp,
   std::optional<mlir::Value> rowStripInput,
+  const spatial::SpatialTargetInfo& target,
   mlir::PatternRewriter& rewriter);
 
-mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp);
+mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
+                                                 const spatial::SpatialTargetInfo& target);
+
+mlir::FailureOr<mlir::Value>
+lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
+                        const spatial::SpatialTargetInfo& target,
+                       mlir::PatternRewriter& rewriter);
 
 mlir::FailureOr<mlir::Value>
 lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
                            std::optional<mlir::Value> rowStripInput,
+                           const spatial::SpatialTargetInfo& target,
                            mlir::PatternRewriter& rewriter);
 
 mlir::LogicalResult
-canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp);
+canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                        const spatial::SpatialTargetInfo& target);
+
+mlir::FailureOr<mlir::Value>
+lowerDenseGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
+                                const spatial::SpatialTargetInfo& target,
+                                mlir::PatternRewriter& rewriter);
 
 mlir::FailureOr<mlir::Value>
 lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
                                    std::optional<mlir::Value> rowStripInput,
+                                   const spatial::SpatialTargetInfo& target,
                                    mlir::PatternRewriter& rewriter);
 
-mlir::LogicalResult canLowerFlattenFromRowStrip(spatial::SpatGraphCompute flattenOp);
-
-mlir::LogicalResult lowerFlattenFromRowStrip(const RowStripPhysicalValue& input,
-                                             spatial::SpatGraphCompute flattenOp,
-                                             mlir::PatternRewriter& rewriter);
-
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutCapabilities.cpp

@@ -0,0 +1,133 @@
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir::spatial {
+
+static LayoutAlternative denseAlternative(Operation *op) {
+  LayoutAlternative alternative;
+  alternative.operandLayouts.assign(op->getNumOperands(), PhysicalLayout::DenseNCHW);
+  alternative.resultLayout = PhysicalLayout::DenseNCHW;
+  return alternative;
+}
+
+static LayoutAlternative rowStripAlternative(Operation *op,
+                                             ArrayRef<PhysicalLayout> operandLayouts) {
+  LayoutAlternative alternative;
+  alternative.operandLayouts.assign(operandLayouts.begin(), operandLayouts.end());
+  alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
+  alternative.intrinsicCost = -2;
+  return alternative;
+}
+
+static bool hasRowStripInput(ArrayRef<PhysicalLayout> operandLayouts, unsigned index) {
+  return index < operandLayouts.size()
+      && operandLayouts[index] == PhysicalLayout::NHWCRowStrip;
+}
+
+SmallVector<LayoutAlternative> SpatConv2DPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (hasRowStripInput(operandLayouts, 0)) {
+    if (succeeded(canConsumeAndProduceRowStrip(*this, target)))
+      alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
+  }
+  else if (succeeded(canLowerConvPlanToRowStrip(*this, target))) {
+    LayoutAlternative alternative = denseAlternative(getOperation());
+    alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
+    alternative.intrinsicCost = -2;
+    alternatives.push_back(std::move(alternative));
+  }
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatReluPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (hasRowStripInput(operandLayouts, 0))
+    alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatSiluPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (hasRowStripInput(operandLayouts, 0)) {
+    LayoutAlternative alternative = rowStripAlternative(getOperation(), operandLayouts);
+    alternative.intrinsicCost = -3;
+    alternatives.push_back(std::move(alternative));
+  }
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatResizeNearestPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (hasRowStripInput(operandLayouts, 0)
+      && succeeded(canLowerResizeNearestPlanToRowStrip(*this, target)))
+    alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatMaxPool2DPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (succeeded(canLowerMaxPoolPlanToRowStrip(*this, target))) {
+    LayoutAlternative alternative = denseAlternative(getOperation());
+    if (hasRowStripInput(operandLayouts, 0))
+      alternative = rowStripAlternative(getOperation(), operandLayouts);
+    alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
+    alternative.intrinsicCost = -2;
+    alternatives.push_back(std::move(alternative));
+  }
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatGlobalAveragePoolPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo& target, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (succeeded(canLowerGlobalAveragePoolPlanToRowStrip(*this, target))) {
+    LayoutAlternative alternative = denseAlternative(getOperation());
+    if (hasRowStripInput(operandLayouts, 0))
+      alternative = rowStripAlternative(getOperation(), operandLayouts);
+    alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
+    alternative.intrinsicCost = -2;
+    alternatives.push_back(std::move(alternative));
+  }
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatBiasAddPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  auto resultType = dyn_cast<RankedTensorType>(getOutput().getType());
+  if (resultType && hasRowStripInput(operandLayouts, 0)
+      && isSupportedBiasAddValue(getBias(), resultType))
+    alternatives.push_back(rowStripAlternative(getOperation(),
+                                                {PhysicalLayout::NHWCRowStrip,
+                                                 PhysicalLayout::DenseNCHW}));
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatAddPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (operandLayouts.size() >= 2 && hasRowStripInput(operandLayouts, 0)
+      && hasRowStripInput(operandLayouts, 1))
+    alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
+  return alternatives;
+}
+
+SmallVector<LayoutAlternative> SpatConcatPlanOp::getLayoutAlternatives(
+    const SpatialTargetInfo&, ArrayRef<PhysicalLayout> operandLayouts) {
+  SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
+  if (!operandLayouts.empty() && llvm::all_of(operandLayouts, [](PhysicalLayout layout) {
+        return layout == PhysicalLayout::NHWCRowStrip;
+      }))
+    alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
+  return alternatives;
+}
+
+} // namespace onnx_mlir::spatial

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -6,358 +6,260 @@
 
 #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 <algorithm>
+
 using namespace mlir;
 
 namespace onnx_mlir {
 namespace {
 
-static constexpr StringLiteral kLogicalLayout = "nchw";
+using LayoutMap = llvm::DenseMap<Value, spatial::PhysicalLayout>;
-static constexpr StringLiteral kDenseLayout = "dense_nchw";
-static constexpr StringLiteral kRowStripLayout = "nhwc_row_strip";
 
-enum class SelectedLayout {
+static spatial::PhysicalLayout getSelectedLayout(const LayoutMap& layouts, Value value) {
-  DenseNchw,
+  if (auto it = layouts.find(value); it != layouts.end())
-  PixelMajorRowStrip,
+    return it->second;
-};
+  if (auto materialize = value.getDefiningOp<spatial::SpatMaterializeLayoutOp>())
-
+    return materialize.getTargetPhysicalLayout();
-static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& layouts, Value value) {
+  if (auto blueprint = value.getDefiningOp<spatial::SpatBlueprintOp>())
-  auto it = layouts.find(value);
+    return blueprint.getPhysicalLayout();
-  return it == layouts.end() ? SelectedLayout::DenseNchw : it->second;
+  return spatial::PhysicalLayout::DenseNCHW;
 }
 
-static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, SelectedLayout>& layouts) {
+static SmallVector<spatial::PhysicalLayout> getOperandLayouts(
-  if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(user))
+    Operation* op, const LayoutMap& layouts) {
-    return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
+  SmallVector<spatial::PhysicalLayout> operandLayouts;
-  if (auto siluPlan = dyn_cast<spatial::SpatSiluPlanOp>(user))
+  operandLayouts.reserve(op->getNumOperands());
-    return getSelectedLayout(layouts, siluPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
+  for (Value operand : op->getOperands())
-  if (auto resizePlan = dyn_cast<spatial::SpatResizeNearestPlanOp>(user))
+    operandLayouts.push_back(getSelectedLayout(layouts, operand));
-    return getSelectedLayout(layouts, resizePlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
+  return operandLayouts;
-  if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user))
-    return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(user))
-    return getSelectedLayout(layouts, addPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto concatPlan = dyn_cast<spatial::SpatConcatPlanOp>(user))
-    return getSelectedLayout(layouts, concatPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
-    return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
-    return getSelectedLayout(layouts, maxPoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(user))
-    return getSelectedLayout(layouts, averagePoolPlan.getResult()) == SelectedLayout::PixelMajorRowStrip;
-  if (auto flattenCompute = dyn_cast<spatial::SpatGraphCompute>(user))
-    return succeeded(canLowerFlattenFromRowStrip(flattenCompute));
-  return false;
 }
 
-static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap<Value, SelectedLayout>& layouts) {
+static FailureOr<SmallVector<spatial::LayoutAlternative>> getAlternatives(
-  for (Operation* user : value.getUsers()) {
+    Operation* op, const LayoutMap& layouts, const spatial::SpatialTargetInfo& target) {
-    if (usesSelectedRowStrip(user, layouts))
+  auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(op);
+  if (!capability)
+    return failure();
+  SmallVector<spatial::LayoutAlternative> alternatives =
+      capability.getLayoutAlternatives(target, getOperandLayouts(op, layouts));
+  if (alternatives.empty())
+    return op->emitOpError("does not advertise a legal Spatial layout alternative"), failure();
+  for (const spatial::LayoutAlternative& alternative : alternatives)
+    if (alternative.operandLayouts.size() != op->getNumOperands())
+      return op->emitOpError("advertises a layout alternative with the wrong operand count"), failure();
+  return alternatives;
+}
+
+static unsigned findCurrentAlternative(
+    Operation* op, ArrayRef<spatial::LayoutAlternative> alternatives,
+    spatial::PhysicalLayout selectedResult) {
+  for (auto [index, alternative] : llvm::enumerate(alternatives))
+    if (alternative.resultLayout == selectedResult)
+      return index;
+  return 0;
+}
+
+static int64_t alternativeCost(Operation* op,
+                               const spatial::LayoutAlternative& alternative,
+                               const LayoutMap& layouts,
+                               const LayoutMap& selectedResults,
+                               const spatial::SpatialTargetInfo& target) {
+  int64_t cost = alternative.intrinsicCost;
+  SmallVector<spatial::PhysicalLayout> operandLayouts = getOperandLayouts(op, layouts);
+  for (auto [actual, required] : llvm::zip(operandLayouts, alternative.operandLayouts))
+    cost += actual != required;
+
+  Value result = op->getResult(0);
+  for (OpOperand& use : result.getUses()) {
+    auto user = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(use.getOwner());
+    if (!user) {
+      if (alternative.resultLayout != spatial::PhysicalLayout::DenseNCHW) {
+        auto flatten = dyn_cast<spatial::SpatGraphCompute>(use.getOwner());
+        if (!flatten || failed(canLowerFlattenFromRowStrip(flatten, target)))
+          ++cost;
+      }
       continue;
-    // Dense-only users must be materialized explicitly.
+    }
-    continue;
+    auto userAlternatives = getAlternatives(use.getOwner(), selectedResults, target);
-  }
+    if (failed(userAlternatives))
-  return true;
-}
-
-static bool canConsumeRowStripAsUser(Operation* user) {
-  if (isa<spatial::SpatReluPlanOp, spatial::SpatSiluPlanOp>(user))
-    return true;
-  if (auto resizePlan = dyn_cast<spatial::SpatResizeNearestPlanOp>(user))
-    return succeeded(canLowerResizeNearestPlanToRowStrip(resizePlan));
-  if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user)) {
-    auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
-    return resultType && isSupportedBiasAddValue(biasAddPlan.getBias(), resultType);
-  }
-  if (isa<spatial::SpatAddPlanOp>(user))
-    return true;
-  if (isa<spatial::SpatConcatPlanOp>(user))
-    return true;
-  if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
-    return succeeded(canConsumeAndProduceRowStrip(convPlan));
-  if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
-    return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan));
-  if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(user))
-    return succeeded(canLowerGlobalAveragePoolPlanToRowStrip(averagePoolPlan));
-  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());
-  if (inputLayout == SelectedLayout::PixelMajorRowStrip)
-    return succeeded(canConsumeAndProduceRowStrip(convPlan));
-  return succeeded(canLowerConvPlanToRowStrip(convPlan));
-}
-
-static SelectedLayout chooseConvLayout(spatial::SpatConv2DPlanOp convPlan,
-                                       llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (!canSelectConvRowStrip(convPlan, layouts))
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(convPlan.getResult(), layouts))
-    return SelectedLayout::DenseNchw;
-  return SelectedLayout::PixelMajorRowStrip;
-}
-
-static SelectedLayout chooseActivationLayout(Value input,
-                                             Value result,
-                                             llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (getSelectedLayout(layouts, input) != SelectedLayout::PixelMajorRowStrip)
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(result, layouts))
-    return SelectedLayout::DenseNchw;
-  return SelectedLayout::PixelMajorRowStrip;
-}
-
-static SelectedLayout chooseResizeLayout(
-    spatial::SpatResizeNearestPlanOp resizePlan,
-    llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  return getSelectedLayout(layouts, resizePlan.getInput()) == SelectedLayout::PixelMajorRowStrip
-      && succeeded(canLowerResizeNearestPlanToRowStrip(resizePlan))
-    ? SelectedLayout::PixelMajorRowStrip : SelectedLayout::DenseNchw;
-}
-
-static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan,
-                                          llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::PixelMajorRowStrip)
-    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::PixelMajorRowStrip;
-}
-
-static SelectedLayout chooseAddLayout(spatial::SpatAddPlanOp addPlan, llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (getSelectedLayout(layouts, addPlan.getLhs()) != SelectedLayout::PixelMajorRowStrip
-      || getSelectedLayout(layouts, addPlan.getRhs()) != SelectedLayout::PixelMajorRowStrip)
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(addPlan.getResult(), layouts))
-    return SelectedLayout::DenseNchw;
-  return SelectedLayout::PixelMajorRowStrip;
-}
-
-static SelectedLayout chooseConcatLayout(spatial::SpatConcatPlanOp concatPlan,
-                                         llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  if (llvm::any_of(concatPlan.getInputs(), [&](Value input) {
-        return getSelectedLayout(layouts, input) != SelectedLayout::PixelMajorRowStrip;
-      }))
-    return SelectedLayout::DenseNchw;
-  if (!allUsersCanHandleRowStrip(concatPlan.getResult(), layouts))
-    return SelectedLayout::DenseNchw;
-  return SelectedLayout::PixelMajorRowStrip;
-}
-
-static SelectedLayout chooseMaxPoolLayout(spatial::SpatMaxPool2DPlanOp maxPoolPlan) {
-  return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan)) ? SelectedLayout::PixelMajorRowStrip
-                                                               : SelectedLayout::DenseNchw;
-}
-
-static SelectedLayout chooseGlobalAveragePoolLayout(
-  spatial::SpatGlobalAveragePoolPlanOp averagePoolPlan) {
-  return succeeded(canLowerGlobalAveragePoolPlanToRowStrip(averagePoolPlan))
-           ? SelectedLayout::PixelMajorRowStrip
-           : SelectedLayout::DenseNchw;
-}
-
-static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Value value) {
-  auto outputType = cast<RankedTensorType>(value.getType());
-  auto [offsets, sizes] = buildRowStripMetadata(outputType);
-  return spatial::SpatBlueprintOp::create(rewriter,
-                                              value.getLoc(),
-                                              outputType,
-                                              value,
-                                              ValueRange {},
-                                              rewriter.getStringAttr(kLogicalLayout),
-                                              rewriter.getStringAttr(kRowStripLayout),
-                                              rewriter.getDenseI64ArrayAttr(offsets),
-                                              rewriter.getDenseI64ArrayAttr(sizes),
-                                              rewriter.getStringAttr(kRowStripIndexMap),
-                                              nullptr,
-                                              nullptr,
-                                              nullptr,
-                                              nullptr,
-                                              nullptr,
-                                              nullptr,
-                                              nullptr);
-}
-
-static void materializeDenseUses(IRRewriter& rewriter,
-                                 Value layoutValue,
-                                 llvm::DenseMap<Value, SelectedLayout>& layouts) {
-  SmallVector<OpOperand*> denseUses;
-  for (OpOperand& use : layoutValue.getUses()) {
-    if (usesSelectedRowStrip(use.getOwner(), layouts))
       continue;
-    denseUses.push_back(&use);
+    spatial::PhysicalLayout userResult =
+        selectedResults.lookup(use.getOwner()->getResult(0));
+    unsigned userIndex = findCurrentAlternative(use.getOwner(), *userAlternatives, userResult);
+    if (use.getOperandNumber() < (*userAlternatives)[userIndex].operandLayouts.size()
+        && (*userAlternatives)[userIndex].operandLayouts[use.getOperandNumber()]
+             != alternative.resultLayout)
+      ++cost;
+  }
+  return cost;
+}
+
+static LogicalResult materializeMismatchedUses(
+    IRRewriter& rewriter, Value value, const LayoutMap& layouts,
+    const spatial::SpatialTargetInfo& target) {
+  spatial::PhysicalLayout sourceLayout = getSelectedLayout(layouts, value);
+  SmallVector<std::pair<OpOperand*, spatial::PhysicalLayout>> mismatches;
+  for (OpOperand& use : value.getUses()) {
+    Operation* userOp = use.getOwner();
+    spatial::PhysicalLayout required = spatial::PhysicalLayout::DenseNCHW;
+    if (auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(userOp)) {
+      auto alternatives = getAlternatives(userOp, layouts, target);
+      if (failed(alternatives))
+        return failure();
+      spatial::PhysicalLayout selected =
+          getSelectedLayout(layouts, userOp->getResult(0));
+      unsigned selectedIndex = findCurrentAlternative(userOp, *alternatives, selected);
+      required = (*alternatives)[selectedIndex].operandLayouts[use.getOperandNumber()];
+    }
+    else if (auto flatten = dyn_cast<spatial::SpatGraphCompute>(userOp);
+             flatten && sourceLayout == spatial::PhysicalLayout::NHWCRowStrip
+                  && succeeded(canLowerFlattenFromRowStrip(flatten, target))) {
+      continue;
+    }
+    if (required != sourceLayout)
+      mismatches.push_back({&use, required});
   }
 
-  for (OpOperand* use : denseUses) {
+  for (auto [use, required] : mismatches) {
-    Operation* owner = use->getOwner();
+    Operation* userOp = use->getOwner();
-    rewriter.setInsertionPoint(owner);
+    rewriter.setInsertionPoint(userOp);
-    auto materialized = spatial::SpatMaterializeLayoutOp::create(rewriter,
+    auto materialized = spatial::SpatMaterializeLayoutOp::create(
-                                                                 owner->getLoc(),
+        rewriter, userOp->getLoc(), use->get().getType(), use->get(),
-                                                                 use->get().getType(),
+        spatial::LogicalLayoutAttr::get(
-                                                                 use->get(),
+          rewriter.getContext(), spatial::LogicalLayout::NCHW),
-                                                                 rewriter.getStringAttr(kLogicalLayout),
+        spatial::PhysicalLayoutAttr::get(rewriter.getContext(), sourceLayout),
-                                                                 rewriter.getStringAttr(kRowStripLayout),
+        spatial::PhysicalLayoutAttr::get(rewriter.getContext(),
-                                                                 rewriter.getStringAttr(kDenseLayout));
+                                          required));
     use->set(materialized.getResult());
   }
+  return success();
 }
 
-struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass, OperationPass<ModuleOp>> {
+static LogicalResult verifySelectedLayouts(
+    ArrayRef<Operation*> planOps, const LayoutMap& layouts,
+    const spatial::SpatialTargetInfo& target) {
+  for (Operation* op : planOps) {
+    auto selected = spatial::getSelectedPhysicalLayout(op);
+    if (!selected)
+      return op->emitOpError("requires a selected physical layout"), failure();
+    auto alternatives = getAlternatives(op, layouts, target);
+    if (failed(alternatives))
+      return failure();
+    if (llvm::none_of(*alternatives, [&](const spatial::LayoutAlternative& alternative) {
+          return alternative.resultLayout == *selected;
+        }))
+      return op->emitOpError("selected physical layout is not advertised by its layout contract"), failure();
+  }
+  return success();
+}
+
+struct SpatialLayoutPlanningPass final
+    : PassWrapper<SpatialLayoutPlanningPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialLayoutPlanningPass)
 
   StringRef getArgument() const override { return "spatial-layout-planning"; }
-  StringRef getDescription() const override { return "Select conservative Spatial layouts and insert reconciliation barriers."; }
+  StringRef getDescription() const override {
+    return "Select Spatial layout alternatives and insert explicit reconciliation barriers.";
+  }
+
+  SpatialLayoutPlanningPass() = default;
+  explicit SpatialLayoutPlanningPass(const spatial::SpatialTargetInfo& target)
+      : target(target), hasTarget(true) {}
 
   void runOnOperation() override {
-    auto entryFunc = getPimEntryFunc(getOperation());
+    ModuleOp moduleOp = getOperation();
+    if (!hasTarget) {
+      moduleOp.emitError("Spatial layout planning requires an injected SpatialTargetInfo");
+      signalPassFailure();
+      return;
+    }
+    auto entryFunc = getPimEntryFunc(moduleOp);
     if (failed(entryFunc)) {
-      getOperation().emitError("failed to locate the PIM entry function during Spatial layout planning");
+      moduleOp.emitError("failed to locate the PIM entry function during Spatial layout planning");
       signalPassFailure();
       return;
     }
 
     func::FuncOp funcOp = *entryFunc;
-    IRRewriter rewriter(&getContext());
+    SmallVector<Operation*> planOps;
-    llvm::DenseMap<Value, SelectedLayout> layouts;
+    for (Operation& op : funcOp.getBody().front())
+      if (isa<spatial::SpatialLayoutCapabilityInterface>(&op))
+        planOps.push_back(&op);
 
-    bool changed = true;
+    LayoutMap layouts;
-    while (changed) {
+    for (Operation* op : planOps)
-      changed = false;
+      layouts[op->getResult(0)] = spatial::PhysicalLayout::DenseNCHW;
-      for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
+
-        if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(&op)) {
+    const size_t maxRounds = 2 * planOps.size() + 1;
-          SelectedLayout selected = chooseConvLayout(convPlan, layouts);
+    bool converged = false;
-          if (layouts[convPlan.getResult()] != selected) {
+    for (size_t round = 0; round < maxRounds && !converged; ++round) {
-            layouts[convPlan.getResult()] = selected;
+      converged = true;
-            changed = true;
+      SmallVector<Operation*> order(planOps);
-          }
+      if (round % 2)
-          continue;
+        std::reverse(order.begin(), order.end());
+      for (Operation* op : order) {
+        auto alternatives = getAlternatives(op, layouts, target);
+        if (failed(alternatives)) {
+          signalPassFailure();
+          return;
         }
-        if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op)) {
+        spatial::PhysicalLayout current = layouts.lookup(op->getResult(0));
-          SelectedLayout selected = chooseActivationLayout(reluPlan.getInput(), reluPlan.getResult(), layouts);
+        unsigned currentIndex = findCurrentAlternative(op, *alternatives, current);
-          if (layouts[reluPlan.getResult()] != selected) {
+        int64_t bestCost = alternativeCost(
-            layouts[reluPlan.getResult()] = selected;
+            op, (*alternatives)[currentIndex], layouts, layouts, target);
-            changed = true;
+        unsigned bestIndex = currentIndex;
+        for (auto [index, alternative] : llvm::enumerate(*alternatives)) {
+          int64_t cost = alternativeCost(op, alternative, layouts, layouts, target);
+          if (cost < bestCost) {
+            bestCost = cost;
+            bestIndex = index;
           }
-          continue;
         }
-        if (auto siluPlan = dyn_cast<spatial::SpatSiluPlanOp>(&op)) {
+        spatial::PhysicalLayout selected = (*alternatives)[bestIndex].resultLayout;
-          SelectedLayout selected = chooseActivationLayout(siluPlan.getInput(), siluPlan.getResult(), layouts);
+        if (selected != current) {
-          if (layouts[siluPlan.getResult()] != selected) {
+          layouts[op->getResult(0)] = selected;
-            layouts[siluPlan.getResult()] = selected;
+          converged = false;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto resizePlan = dyn_cast<spatial::SpatResizeNearestPlanOp>(&op)) {
-          SelectedLayout selected = chooseResizeLayout(resizePlan, layouts);
-          if (layouts[resizePlan.getResult()] != selected) {
-            layouts[resizePlan.getResult()] = selected;
-            changed = true;
-          }
-          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;
-        }
-        if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(&op)) {
-          SelectedLayout selected = chooseAddLayout(addPlan, layouts);
-          if (layouts[addPlan.getResult()] != selected) {
-            layouts[addPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto concatPlan = dyn_cast<spatial::SpatConcatPlanOp>(&op)) {
-          SelectedLayout selected = chooseConcatLayout(concatPlan, layouts);
-          if (layouts[concatPlan.getResult()] != selected) {
-            layouts[concatPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op)) {
-          SelectedLayout selected = chooseMaxPoolLayout(maxPoolPlan);
-          if (layouts[maxPoolPlan.getResult()] != selected) {
-            layouts[maxPoolPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
-        }
-        if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op)) {
-          SelectedLayout selected = chooseGlobalAveragePoolLayout(averagePoolPlan);
-          if (layouts[averagePoolPlan.getResult()] != selected) {
-            layouts[averagePoolPlan.getResult()] = selected;
-            changed = true;
-          }
-          continue;
         }
       }
     }
-
+    if (!converged) {
-    for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
+      moduleOp.emitError("Spatial layout selection did not converge within its bounded iteration budget");
-      Value producedValue;
+      signalPassFailure();
-      if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(&op))
+      return;
-        producedValue = convPlan.getResult();
-      else if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op))
-        producedValue = biasAddPlan.getResult();
-      else if (auto addPlan = dyn_cast<spatial::SpatAddPlanOp>(&op))
-        producedValue = addPlan.getResult();
-      else if (auto concatPlan = dyn_cast<spatial::SpatConcatPlanOp>(&op))
-        producedValue = concatPlan.getResult();
-      else if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op))
-        producedValue = reluPlan.getResult();
-      else if (auto siluPlan = dyn_cast<spatial::SpatSiluPlanOp>(&op))
-        producedValue = siluPlan.getResult();
-      else if (auto resizePlan = dyn_cast<spatial::SpatResizeNearestPlanOp>(&op))
-        producedValue = resizePlan.getResult();
-      else if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op))
-        producedValue = maxPoolPlan.getResult();
-      else if (auto averagePoolPlan = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op))
-        producedValue = averagePoolPlan.getResult();
-      else
-        continue;
-
-      if (getSelectedLayout(layouts, producedValue) != SelectedLayout::PixelMajorRowStrip)
-        continue;
-
-      rewriter.setInsertionPointAfter(&op);
-      auto blueprint = insertRowStripBlueprint(rewriter, producedValue);
-      rewriter.replaceAllUsesExcept(producedValue, blueprint.getResult(), blueprint);
-      materializeDenseUses(rewriter, blueprint.getResult(), layouts);
     }
-    if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+    IRRewriter rewriter(&getContext());
-      getOperation().emitError("logical Spatial graph verification failed after SpatialLayoutPlanning");
+    for (Operation* op : planOps) {
+      op->setAttr(spatial::kSelectedLayoutAttrName,
+                  spatial::PhysicalLayoutAttr::get(
+                      rewriter.getContext(), layouts.lookup(op->getResult(0))));
+      if (failed(materializeMismatchedUses(rewriter, op->getResult(0), layouts, target))) {
+        signalPassFailure();
+        return;
+      }
+    }
+    if (failed(verifySelectedLayouts(planOps, layouts, target))
+        || failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+      moduleOp.emitError("Spatial layout planning verification failed");
       signalPassFailure();
     }
   }
+
+  spatial::SpatialTargetInfo target;
+  bool hasTarget = false;
 };
 
 } // namespace
 
-std::unique_ptr<Pass> createSpatialLayoutPlanningPass() { return std::make_unique<SpatialLayoutPlanningPass>(); }
+std::unique_ptr<Pass> createSpatialLayoutPlanningPass() {
+  return std::make_unique<SpatialLayoutPlanningPass>();
+}
+
+std::unique_ptr<Pass> createSpatialLayoutPlanningPass(
+    const spatial::SpatialTargetInfo& target) {
+  return std::make_unique<SpatialLayoutPlanningPass>(target);
+}
 
 } // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

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

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -42,12 +42,11 @@ static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
   if (!resultType || !resultType.hasStaticShape())
     return blueprint.emitOpError("fragment assembly lowering requires a static ranked tensor result");
 
-  std::optional<StringRef> modeAttr = blueprint.getMode();
   std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
   std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
   std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
   std::optional<ArrayRef<int64_t>> fragmentStridesAttr = blueprint.getFragmentStrides();
-  if (!modeAttr || *modeAttr != "fragment_assembly" || !operandIndicesAttr || !sourceSlotsAttr
+  if (!spatial::isFragmentAssembly(blueprint.getMode()) || !operandIndicesAttr || !sourceSlotsAttr
       || !sourceOffsetsAttr || !fragmentStridesAttr)
     return blueprint.emitOpError("fragment assembly lowering requires explicit fragment metadata");
 
@@ -203,8 +202,7 @@ static bool isHostMaterializableHelperOp(Operation* op) {
   if (isa<arith::ConstantOp>(op) || op->hasTrait<OpTrait::ConstantLike>())
     return true;
   if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-    std::optional<StringRef> mode = blueprint.getMode();
+    return spatial::isFragmentAssembly(blueprint.getMode());
-    return mode && *mode == "fragment_assembly";
   }
   return isShapingOnlyOp(op) || isPureIndexComputationOp(op);
 }
@@ -291,8 +289,7 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatSchedule
   }
   for (Operation& op : block.without_terminator()) {
     if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
-      std::optional<StringRef> modeAttr = blueprint.getMode();
+      if (spatial::isFragmentAssembly(blueprint.getMode())) {
-      if (modeAttr && *modeAttr == "fragment_assembly") {
         auto lowered = lowerFragmentAssemblyBlueprint(rewriter, blueprint, mapping);
         if (failed(lowered))
           return false;

src/PIM/Conversion/SpatialToPim/Patterns.cpp

@@ -22,8 +22,7 @@ struct LowerFragmentAssemblyBlueprintPattern
   LogicalResult matchAndRewrite(spatial::SpatBlueprintOp op,
                                 OpAdaptor adaptor,
                                 ConversionPatternRewriter& rewriter) const override {
-    std::optional<StringRef> modeAttr = op.getMode();
+    if (!spatial::isFragmentAssembly(op.getMode()))
-    if (!modeAttr || *modeAttr != "fragment_assembly")
       return failure();
 
     auto resultType = dyn_cast<ShapedType>(op.getOutput().getType());

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp

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

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

@@ -33,6 +33,9 @@ using namespace pim;
 
 namespace onnx_mlir {
 
+static void annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp);
+static FailureOr<func::FuncOp> requirePimEntryFunc(ModuleOp moduleOp, StringRef phase);
+
 namespace {
 
 struct MemRefCopyWorkItem {
@@ -333,22 +336,6 @@ static LogicalResult verifyPimCopyEndpoints(Operation* copy,
   return success(valid);
 }
 
-struct PimBufferizationPass : PassWrapper<PimBufferizationPass, OperationPass<ModuleOp>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationPass)
-  StringRef getArgument() const override { return "bufferize-pim"; }
-  StringRef getDescription() const override { return "Bufferize PIM and Spatial ops."; }
-
-  PimBufferizationPass() = default;
-  PimBufferizationPass(const PimBufferizationPass& pass) {}
-
-  void runOnOperation() final;
-
-private:
-  void annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) const;
-  LogicalResult verifyContiguousRuntimeOperands(ModuleOp moduleOp) const;
-  LogicalResult verifyPimCopyAddressSpaces(ModuleOp moduleOp) const;
-};
-
 static void materializeWritableConstantDestinations(func::FuncOp funcOp) {
   SmallVector<OpOperand*> constantBackedRoots;
   llvm::SmallPtrSet<OpOperand*, 8> seenRoots;
@@ -387,65 +374,23 @@ static void materializeWritableConstantDestinations(func::FuncOp funcOp) {
   }
 }
 
-static LogicalResult verifyPimCoresNeedNoTensorCopies(
+static bufferization::OneShotBufferizationOptions makePimBufferizationOptions() {
-  ModuleOp module, const bufferization::OneShotBufferizationOptions& baseOptions) {
-  static constexpr StringLiteral kExistingAlloc = "raptor.existing_core_alloc";
-  OwningOpRef<ModuleOp> clone = module.clone();
-  clone->walk([&](bufferization::AllocTensorOp alloc) {
-    if (alloc->getParentOfType<pim::PimCoreOp>()
-        || alloc->getParentOfType<pim::PimCoreBatchOp>())
-      alloc->setAttr(kExistingAlloc, UnitAttr::get(module.getContext()));
-  });
-
-  auto options = baseOptions;
-  options.bufferizeFunctionBoundaries = false;
-  options.opFilter.allowOperation([](Operation* op) {
-    return isa<pim::PimCoreOp, pim::PimCoreBatchOp>(op)
-        || op->getParentOfType<pim::PimCoreOp>()
-        || op->getParentOfType<pim::PimCoreBatchOp>();
-  });
-
-  bufferization::BufferizationState state;
-  if (failed(bufferization::insertTensorCopies(*clone, options, state))) {
-    module.emitError("official one-shot analysis failed while verifying PIM core copy freedom");
-    return failure();
-  }
-
-  CappedDiagnosticReporter diagnostics;
-  clone->walk([&](bufferization::AllocTensorOp alloc) {
-    if (alloc->hasAttr(kExistingAlloc)
-        || (!alloc->getParentOfType<pim::PimCoreOp>()
-            && !alloc->getParentOfType<pim::PimCoreBatchOp>()))
-      return;
-    Operation* requiredBy = alloc->getUsers().empty()
-      ? alloc.getOperation() : *alloc->getUsers().begin();
-    diagnostics.report(requiredBy, [](Operation* op) {
-      op->emitOpError("official one-shot bufferization requires a tensor copy inside a PIM core");
-    });
-  });
-  diagnostics.emitSuppressedSummary(module, "required PIM core tensor copies");
-  return success(!diagnostics.hasFailure());
-}
-
-} // namespace
-
-void PimBufferizationPass::runOnOperation() {
-  auto moduleOp = getOperation();
-  auto funcOp = *getPimEntryFunc(moduleOp);
-
   bufferization::OneShotBufferizationOptions options;
   options.allowUnknownOps = true;
   options.bufferizeFunctionBoundaries = true;
   options.setFunctionBoundaryTypeConversion(bufferization::LayoutMapOption::IdentityLayoutMap);
+  return options;
+}
 
+static LogicalResult preparePimBufferization(func::FuncOp funcOp) {
   materializeWritableConstantDestinations(funcOp);
-  if (failed(verifyPimCoresNeedNoTensorCopies(moduleOp, options))) {
+  return success();
-    signalPassFailure();
+}
-    return;
-  }
 
+static LogicalResult runOneShotPimBufferization(
+    ModuleOp moduleOp, const bufferization::OneShotBufferizationOptions& options) {
   auto hostOptions = options;
-  hostOptions.opFilter.denyOperation([](Operation *op) {
+  hostOptions.opFilter.denyOperation([](Operation* op) {
     return op->getParentOfType<pim::PimCoreOp>()
         || op->getParentOfType<pim::PimCoreBatchOp>();
   });
@@ -453,84 +398,14 @@ void PimBufferizationPass::runOnOperation() {
   if (failed(bufferization::insertTensorCopies(moduleOp, hostOptions, state))
       || failed(bufferization::bufferizeModuleOp(moduleOp, options, state))) {
     moduleOp.emitError("Failed to bufferize PIM and Spatial ops");
-    signalPassFailure();
+    return failure();
-    return;
   }
-
+  return success();
-  forwardSingleConsumerReceiveCopies(funcOp);
-  forwardSingleConsumerContiguousInputCopies(funcOp);
-  forwardSingleConsumerPimOutputCopies(funcOp);
-
-  MLIRContext* ctx = moduleOp.getContext();
-  PatternRewriter rewriter(ctx);
-
-  SmallVector<MemRefCopyWorkItem> copyWorklist;
-  llvm::SmallPtrSet<Operation*, 16> seenCopyOps;
-  auto addCopyOp = [&](memref::CopyOp copyOp, const StaticValueKnowledge& knowledge) {
-    if (seenCopyOps.insert(copyOp.getOperation()).second)
-      copyWorklist.push_back({copyOp, knowledge});
-  };
-
-  moduleOp.walk([&](pim::PimCoreOp coreOp) {
-    StaticValueKnowledge knowledge = seedCoreKnowledge(coreOp);
-    (void) walkPimCoreBlockStructurally(
-      coreOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) {
-        if (auto copyOp = dyn_cast<memref::CopyOp>(&op))
-          addCopyOp(copyOp, opKnowledge);
-        return success();
-      });
-  });
-  moduleOp.walk([&](pim::PimCoreBatchOp coreBatchOp) {
-    for (unsigned lane = 0; lane < coreBatchOp.getLaneCount(); ++lane) {
-      StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, lane);
-      (void) walkPimCoreBlockStructurally(
-        coreBatchOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) {
-          if (auto copyOp = dyn_cast<memref::CopyOp>(&op))
-            addCopyOp(copyOp, opKnowledge);
-          return success();
-      });
-    }
-  });
-
-  bool hasFailed = false;
-  Value zeroOffset = getOrCreateIndexConstant(rewriter, funcOp, 0);
-  for (const MemRefCopyWorkItem& workItem : copyWorklist) {
-    memref::CopyOp copyOp = workItem.copyOp;
-    rewriter.setInsertionPoint(copyOp);
-    if (failed(lowerMemRefCopyToPimCopy(copyOp, zeroOffset, rewriter, workItem.knowledge)))
-      hasFailed = true;
-  }
-  if (hasFailed) {
-    signalPassFailure();
-    return;
-  }
-
-  RewritePatternSet contiguityPatterns(ctx);
-  populatePimContiguityNormalizationPatterns(contiguityPatterns);
-
-  GreedyRewriteConfig contiguityConfig;
-  contiguityConfig.enableFolding(false);
-  if (failed(applyPatternsGreedily(moduleOp, std::move(contiguityPatterns), contiguityConfig))) {
-    moduleOp.emitError("failed to normalize PIM copy contiguity during bufferization");
-    signalPassFailure();
-    return;
-  }
-  if (failed(verifyContiguousRuntimeOperands(moduleOp))) {
-    signalPassFailure();
-    return;
-  }
-  if (failed(verifyPimCopyAddressSpaces(moduleOp))) {
-    signalPassFailure();
-    return;
-  }
-
-  annotateWeightsMemrefs(moduleOp, funcOp);
-
-  // Dump to file for debug
-  dumpModule(moduleOp, "pim1_buff");
 }
 
-void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) const {
+} // namespace
+
+static void annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) {
   auto markWeights = [&](Operation* op) {
     walkPimMvmVmmWeightUses(op, [&](OpOperand& weightUse) {
       Value weight = weightUse.get();
@@ -548,7 +423,7 @@ void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncO
   funcOp.walk([&](PimCoreBatchOp coreBatchOp) { markWeights(coreBatchOp); });
 }
 
-LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp moduleOp) const {
+static LogicalResult verifyContiguousRuntimeOperands(ModuleOp moduleOp) {
   bool hasFailure = false;
 
   auto verifyWithKnowledge = [&](auto coreLikeOp, const StaticValueKnowledge& initialKnowledge) {
@@ -640,7 +515,7 @@ LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp mod
   return success();
 }
 
-LogicalResult PimBufferizationPass::verifyPimCopyAddressSpaces(ModuleOp moduleOp) const {
+static LogicalResult verifyPimCopyAddressSpaces(ModuleOp moduleOp) {
   size_t failureCount = 0;
   auto verifyWithKnowledge = [&](auto coreLikeOp, const StaticValueKnowledge& initialKnowledge) {
     (void) walkPimCoreBlockStructurally(
@@ -675,6 +550,201 @@ LogicalResult PimBufferizationPass::verifyPimCopyAddressSpaces(ModuleOp moduleOp
   return success(failureCount == 0);
 }
 
-std::unique_ptr<Pass> createPimBufferizationPass() { return std::make_unique<PimBufferizationPass>(); }
+static LogicalResult normalizePimMemory(ModuleOp moduleOp, func::FuncOp funcOp) {
+  forwardSingleConsumerReceiveCopies(funcOp);
+  forwardSingleConsumerContiguousInputCopies(funcOp);
+  forwardSingleConsumerPimOutputCopies(funcOp);
+
+  MLIRContext* ctx = moduleOp.getContext();
+  PatternRewriter rewriter(ctx);
+
+  SmallVector<MemRefCopyWorkItem> copyWorklist;
+  llvm::SmallPtrSet<Operation*, 16> seenCopyOps;
+  auto addCopyOp = [&](memref::CopyOp copyOp, const StaticValueKnowledge& knowledge) {
+    if (seenCopyOps.insert(copyOp.getOperation()).second)
+      copyWorklist.push_back({copyOp, knowledge});
+  };
+
+  moduleOp.walk([&](pim::PimCoreOp coreOp) {
+    StaticValueKnowledge knowledge = seedCoreKnowledge(coreOp);
+    (void) walkPimCoreBlockStructurally(
+      coreOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) {
+        if (auto copyOp = dyn_cast<memref::CopyOp>(&op))
+          addCopyOp(copyOp, opKnowledge);
+        return success();
+      });
+  });
+  moduleOp.walk([&](pim::PimCoreBatchOp coreBatchOp) {
+    for (unsigned lane = 0; lane < coreBatchOp.getLaneCount(); ++lane) {
+      StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, lane);
+      (void) walkPimCoreBlockStructurally(
+        coreBatchOp.getBody().front(), knowledge, [&](Operation& op, const StaticValueKnowledge& opKnowledge) {
+          if (auto copyOp = dyn_cast<memref::CopyOp>(&op))
+            addCopyOp(copyOp, opKnowledge);
+          return success();
+        });
+    }
+  });
+
+  bool hasFailed = false;
+  Value zeroOffset = getOrCreateIndexConstant(rewriter, funcOp, 0);
+  for (const MemRefCopyWorkItem& workItem : copyWorklist) {
+    memref::CopyOp copyOp = workItem.copyOp;
+    rewriter.setInsertionPoint(copyOp);
+    if (failed(lowerMemRefCopyToPimCopy(copyOp, zeroOffset, rewriter, workItem.knowledge)))
+      hasFailed = true;
+  }
+  if (hasFailed)
+    return failure();
+
+  RewritePatternSet contiguityPatterns(ctx);
+  populatePimContiguityNormalizationPatterns(contiguityPatterns);
+
+  GreedyRewriteConfig contiguityConfig;
+  contiguityConfig.enableFolding(false);
+  if (failed(applyPatternsGreedily(moduleOp, std::move(contiguityPatterns), contiguityConfig))) {
+    moduleOp.emitError("failed to normalize PIM copy contiguity during bufferization");
+    return failure();
+  }
+  annotateWeightsMemrefs(moduleOp, funcOp);
+  dumpModule(moduleOp, "pim1_buff");
+  return success();
+}
+
+static FailureOr<func::FuncOp> requirePimEntryFunc(ModuleOp moduleOp, StringRef phase) {
+  auto entryFunc = getPimEntryFunc(moduleOp);
+  if (failed(entryFunc)) {
+    moduleOp.emitError("failed to locate the PIM entry function during ") << phase;
+    return failure();
+  }
+  return *entryFunc;
+}
+
+namespace {
+
+struct PimBufferizationPreparationPass
+    : PassWrapper<PimBufferizationPreparationPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationPreparationPass)
+
+  StringRef getArgument() const override { return "pim-bufferization-preparation"; }
+  StringRef getDescription() const override {
+    return "Prepare writable tensor destinations for PIM one-shot bufferization.";
+  }
+
+  void runOnOperation() final {
+    ModuleOp moduleOp = getOperation();
+    auto funcOp = requirePimEntryFunc(moduleOp, "PIM bufferization preparation");
+    if (failed(funcOp)) {
+      signalPassFailure();
+      return;
+    }
+    if (failed(preparePimBufferization(*funcOp)))
+      signalPassFailure();
+  }
+};
+
+struct PimOneShotBufferizationPass
+    : PassWrapper<PimOneShotBufferizationPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimOneShotBufferizationPass)
+
+  StringRef getArgument() const override { return "pim-one-shot-bufferization"; }
+  StringRef getDescription() const override {
+    return "Run one-shot bufferization for PIM and Spatial tensors.";
+  }
+
+  void runOnOperation() final {
+    if (failed(runOneShotPimBufferization(getOperation(), makePimBufferizationOptions())))
+      signalPassFailure();
+  }
+};
+
+struct PimMemoryNormalizationPass
+    : PassWrapper<PimMemoryNormalizationPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimMemoryNormalizationPass)
+
+  StringRef getArgument() const override { return "pim-memory-normalization"; }
+  StringRef getDescription() const override {
+    return "Normalize PIM memory copies and verify addressable operands.";
+  }
+
+  void runOnOperation() final {
+    ModuleOp moduleOp = getOperation();
+    auto funcOp = requirePimEntryFunc(moduleOp, "PIM memory normalization");
+    if (failed(funcOp)) {
+      signalPassFailure();
+      return;
+    }
+    if (failed(normalizePimMemory(moduleOp, *funcOp)))
+      signalPassFailure();
+  }
+};
+
+static LogicalResult verifyNoTensorValues(ModuleOp moduleOp) {
+  size_t failureCount = 0;
+  moduleOp.walk([&](Operation* op) {
+    if (failureCount >= 8)
+      return;
+    if (op->getDialect()->getNamespace() == "tensor") {
+      op->emitOpError("tensor operation remains after PIM bufferization");
+      ++failureCount;
+      return;
+    }
+    for (Value value : op->getOperands()) {
+      if (isa<TensorType>(value.getType())) {
+        op->emitOpError("tensor operand remains after PIM bufferization");
+        ++failureCount;
+        return;
+      }
+    }
+    for (Value value : op->getResults()) {
+      if (isa<TensorType>(value.getType())) {
+        op->emitOpError("tensor result remains after PIM bufferization");
+        ++failureCount;
+        return;
+      }
+    }
+  });
+  if (failureCount != 0)
+    moduleOp.emitError() << "found " << failureCount
+                         << " tensor value(s) after PIM bufferization"
+                         << (failureCount == 8 ? " (first 8 reported)" : "");
+  return success(failureCount == 0);
+}
+
+struct PimBufferizationVerificationPass
+    : PassWrapper<PimBufferizationVerificationPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationVerificationPass)
+
+  StringRef getArgument() const override { return "pim-bufferization-verification"; }
+  StringRef getDescription() const override {
+    return "Verify tensor elimination, contiguity, and PIM copy address spaces.";
+  }
+
+  void runOnOperation() final {
+    ModuleOp moduleOp = getOperation();
+    if (failed(verifyNoTensorValues(moduleOp))
+        || failed(verifyContiguousRuntimeOperands(moduleOp))
+        || failed(verifyPimCopyAddressSpaces(moduleOp)))
+      signalPassFailure();
+  }
+};
+
+} // namespace
+
+std::unique_ptr<Pass> createPimBufferizationPreparationPass() {
+  return std::make_unique<PimBufferizationPreparationPass>();
+}
+
+std::unique_ptr<Pass> createPimOneShotBufferizationPass() {
+  return std::make_unique<PimOneShotBufferizationPass>();
+}
+
+std::unique_ptr<Pass> createPimMemoryNormalizationPass() {
+  return std::make_unique<PimMemoryNormalizationPass>();
+}
+
+std::unique_ptr<Pass> createPimBufferizationVerificationPass() {
+  return std::make_unique<PimBufferizationVerificationPass>();
+}
 
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -1,6 +1,16 @@
 add_onnx_mlir_dialect(Spatial spat)
 add_onnx_mlir_dialect_doc(spat Spatial.td)
 
+set(LLVM_TARGET_DEFINITIONS Spatial.td)
+mlir_tablegen(SpatialEnums.hpp.inc -gen-enum-decls "-I${ONNX_MLIR_SRC_ROOT}")
+mlir_tablegen(SpatialEnums.cpp.inc -gen-enum-defs "-I${ONNX_MLIR_SRC_ROOT}")
+add_public_tablegen_target(OMSpatialEnumsIncGen)
+
+set(LLVM_TARGET_DEFINITIONS SpatialLayoutInterface.td)
+mlir_tablegen(SpatialLayoutInterface.hpp.inc -gen-op-interface-decls "-I${ONNX_MLIR_SRC_ROOT}")
+mlir_tablegen(SpatialLayoutInterface.cpp.inc -gen-op-interface-defs "-I${ONNX_MLIR_SRC_ROOT}")
+add_public_tablegen_target(OMSpatialLayoutInterfaceIncGen)
+
 add_pim_library(SpatialOps
   SpatialOps.cpp
   SpatialOpsAsm.cpp
@@ -18,7 +28,7 @@ add_pim_library(SpatialOps
   Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp
   Transforms/MergeComputeNodes/DeferredResultRealization.cpp
   Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp
-  Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
+  Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp
   Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp
   Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp
   Transforms/MergeComputeNodes/ScheduledComputeReport.cpp
@@ -33,6 +43,8 @@ add_pim_library(SpatialOps
   DEPENDS
   OMONNXIncGen
   OMSpatialIncGen
+  OMSpatialEnumsIncGen
+  OMSpatialLayoutInterfaceIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR

src/PIM/Dialect/Spatial/Spatial.td

@@ -5,20 +5,77 @@ include "mlir/IR/OpBase.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/BuiltinTypes.td"
 include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/EnumAttr.td"
 include "mlir/IR/RegionKindInterface.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/ParallelCombiningOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "src/Accelerators/PIM/Dialect/Spatial/SpatialLayoutInterface.td"
 
 def SpatialDialect : Dialect {
   let name = "spat";
   let summary = "Dialect designed for deep learning computation in a spatial architecture";
   let cppNamespace = "::onnx_mlir::spatial";
+  let useDefaultAttributePrinterParser = 0;
+  let extraClassDeclaration = [{
+    ::mlir::Attribute parseAttribute(::mlir::DialectAsmParser &parser,
+                                     ::mlir::Type type) const override;
+    void printAttribute(::mlir::Attribute attr,
+                        ::mlir::DialectAsmPrinter &printer) const override;
+  }];
+}
+
+def SpatLogicalLayoutNCHW : I32EnumAttrCase<"NCHW", 0, "nchw">;
+def SpatLogicalLayout : I32EnumAttr<"LogicalLayout", "Logical tensor layout", [
+  SpatLogicalLayoutNCHW
+]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::onnx_mlir::spatial";
+}
+
+def SpatLogicalLayoutAttr : EnumAttr<SpatialDialect, SpatLogicalLayout, "logical_layout"> {
+  let assemblyFormat = "$value";
+}
+
+def SpatPhysicalLayoutDenseNCHW : I32EnumAttrCase<"DenseNCHW", 0, "dense_nchw">;
+def SpatPhysicalLayoutNCHWRowStrip : I32EnumAttrCase<"NCHWRowStrip", 1, "nchw_row_strip">;
+def SpatPhysicalLayoutNHWCRowStrip : I32EnumAttrCase<"NHWCRowStrip", 2, "nhwc_row_strip">;
+def SpatPhysicalLayoutFragmented : I32EnumAttrCase<"Fragmented", 3, "fragmented">;
+def SpatPhysicalLayout : I32EnumAttr<"PhysicalLayout", "Physical tensor layout", [
+  SpatPhysicalLayoutDenseNCHW,
+  SpatPhysicalLayoutNCHWRowStrip,
+  SpatPhysicalLayoutNHWCRowStrip,
+  SpatPhysicalLayoutFragmented
+]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::onnx_mlir::spatial";
+}
+
+def SpatPhysicalLayoutAttr : EnumAttr<SpatialDialect, SpatPhysicalLayout, "physical_layout"> {
+  let assemblyFormat = "$value";
+}
+
+def SpatBlueprintModePhysicalView : I32EnumAttrCase<"PhysicalView", 0, "physical_view">;
+def SpatBlueprintModeFragmentAssembly : I32EnumAttrCase<"FragmentAssembly", 1, "fragment_assembly">;
+def SpatBlueprintMode : I32EnumAttr<"BlueprintMode", "Blueprint reconstruction mode", [
+  SpatBlueprintModePhysicalView,
+  SpatBlueprintModeFragmentAssembly
+]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::onnx_mlir::spatial";
+}
+
+def SpatBlueprintModeAttr : EnumAttr<SpatialDialect, SpatBlueprintMode, "blueprint_mode"> {
+  let assemblyFormat = "$value";
 }
 
 class SpatOp<string mnemonic, list<Trait> traits = []> :
     Op<SpatialDialect, mnemonic, traits>;
 
+class SpatLayoutPlanOp<string mnemonic> : SpatOp<mnemonic,
+    [SpatialLayoutCapabilityInterface,
+     DeclareOpInterfaceMethods<SpatialLayoutCapabilityInterface>]>;
+
 // TODO maybe remove and use AnyRankedTensor directly
 def SpatTensor :
   AnyTypeOf<[AnyMemRef, AnyRankedTensor], "", "::mlir::ShapedType">;
@@ -252,7 +309,7 @@ def SpatConcatOp : SpatOp<"concat", []> {
 // Planning
 //===----------------------------------------------------------------------===//
 
-def SpatConv2DPlanOp : SpatOp<"conv2d_plan", []> {
+def SpatConv2DPlanOp : SpatLayoutPlanOp<"conv2d_plan"> {
   let summary = "Structured Conv2D planning op that preserves logical ONNX geometry";
 
   let arguments = (ins
@@ -263,7 +320,7 @@ def SpatConv2DPlanOp : SpatOp<"conv2d_plan", []> {
     DenseI64ArrayAttr:$strides,
     DenseI64ArrayAttr:$dilations,
     I64Attr:$group,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -273,12 +330,12 @@ def SpatConv2DPlanOp : SpatOp<"conv2d_plan", []> {
   let hasVerifier = 1;
 }
 
-def SpatReluPlanOp : SpatOp<"relu_plan", []> {
+def SpatReluPlanOp : SpatLayoutPlanOp<"relu_plan"> {
   let summary = "Layout-aware ReLU planning op";
 
   let arguments = (ins
     SpatTensor:$input,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -288,12 +345,12 @@ def SpatReluPlanOp : SpatOp<"relu_plan", []> {
   let hasVerifier = 1;
 }
 
-def SpatSiluPlanOp : SpatOp<"silu_plan", []> {
+def SpatSiluPlanOp : SpatLayoutPlanOp<"silu_plan"> {
   let summary = "Layout-aware SiLU planning op";
 
   let arguments = (ins
     SpatTensor:$input,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -303,12 +360,12 @@ def SpatSiluPlanOp : SpatOp<"silu_plan", []> {
   let hasVerifier = 1;
 }
 
-def SpatResizeNearestPlanOp : SpatOp<"resize_nearest_plan", []> {
+def SpatResizeNearestPlanOp : SpatLayoutPlanOp<"resize_nearest_plan"> {
   let summary = "Layout-aware nearest asymmetric Resize planning op";
 
   let arguments = (ins
     SpatTensor:$input,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -318,7 +375,7 @@ def SpatResizeNearestPlanOp : SpatOp<"resize_nearest_plan", []> {
   let hasVerifier = 1;
 }
 
-def SpatMaxPool2DPlanOp : SpatOp<"max_pool2d_plan", []> {
+def SpatMaxPool2DPlanOp : SpatLayoutPlanOp<"max_pool2d_plan"> {
   let summary = "Layout-aware 2D NCHW MaxPool planning op";
 
   let arguments = (ins
@@ -327,7 +384,7 @@ def SpatMaxPool2DPlanOp : SpatOp<"max_pool2d_plan", []> {
     DenseI64ArrayAttr:$pads,
     DenseI64ArrayAttr:$strides,
     DenseI64ArrayAttr:$dilations,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -337,12 +394,12 @@ def SpatMaxPool2DPlanOp : SpatOp<"max_pool2d_plan", []> {
   let hasVerifier = 1;
 }
 
-def SpatGlobalAveragePoolPlanOp : SpatOp<"global_average_pool_plan", []> {
+def SpatGlobalAveragePoolPlanOp : SpatLayoutPlanOp<"global_average_pool_plan"> {
   let summary = "Layout-aware NCHW global average-pool planning op";
 
   let arguments = (ins
     SpatTensor:$input,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -352,13 +409,13 @@ def SpatGlobalAveragePoolPlanOp : SpatOp<"global_average_pool_plan", []> {
   let hasVerifier = 1;
 }
 
-def SpatBiasAddPlanOp : SpatOp<"bias_add_plan", []> {
+def SpatBiasAddPlanOp : SpatLayoutPlanOp<"bias_add_plan"> {
   let summary = "Layout-aware Conv-style bias add planning op";
 
   let arguments = (ins
     SpatTensor:$input,
     SpatTensor:$bias,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -368,13 +425,13 @@ def SpatBiasAddPlanOp : SpatOp<"bias_add_plan", []> {
   let hasVerifier = 1;
 }
 
-def SpatAddPlanOp : SpatOp<"add_plan", []> {
+def SpatAddPlanOp : SpatLayoutPlanOp<"add_plan"> {
   let summary = "Layout-aware elementwise add planning op";
 
   let arguments = (ins
     SpatTensor:$lhs,
     SpatTensor:$rhs,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -384,13 +441,13 @@ def SpatAddPlanOp : SpatOp<"add_plan", []> {
   let hasVerifier = 1;
 }
 
-def SpatConcatPlanOp : SpatOp<"concat_plan", []> {
+def SpatConcatPlanOp : SpatLayoutPlanOp<"concat_plan"> {
   let summary = "Layout-aware tensor concatenation planning op";
 
   let arguments = (ins
     Variadic<SpatTensor>:$inputs,
     I64Attr:$axis,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -406,12 +463,12 @@ def SpatBlueprintOp : SpatOp<"blueprint", []> {
   let arguments = (ins
     SpatTensor:$input,
     Variadic<SpatTensor>:$fragments,
-    StrAttr:$logicalLayout,
+    SpatLogicalLayoutAttr:$logicalLayout,
-    StrAttr:$physicalLayout,
+    SpatPhysicalLayoutAttr:$physicalLayout,
     DenseI64ArrayAttr:$fragmentOffsets,
     DenseI64ArrayAttr:$fragmentSizes,
     StrAttr:$indexMap,
-    OptionalAttr<StrAttr>:$mode,
+    OptionalAttr<SpatBlueprintModeAttr>:$mode,
     OptionalAttr<DenseI64ArrayAttr>:$fragmentOperandIndices,
     OptionalAttr<DenseI64ArrayAttr>:$fragmentSourceSlots,
     OptionalAttr<DenseI64ArrayAttr>:$fragmentSourceOffsets,
@@ -433,9 +490,9 @@ def SpatMaterializeLayoutOp : SpatOp<"materialize_layout", []> {
 
   let arguments = (ins
     SpatTensor:$input,
-    StrAttr:$logicalLayout,
+    SpatLogicalLayoutAttr:$logicalLayout,
-    StrAttr:$sourcePhysicalLayout,
+    SpatPhysicalLayoutAttr:$sourcePhysicalLayout,
-    StrAttr:$targetPhysicalLayout
+    SpatPhysicalLayoutAttr:$targetPhysicalLayout
   );
 
   let results = (outs

src/PIM/Dialect/Spatial/SpatialLayoutInterface.td

@@ -0,0 +1,24 @@
+#ifndef SPATIAL_LAYOUT_INTERFACE_TD
+#define SPATIAL_LAYOUT_INTERFACE_TD
+
+include "mlir/IR/OpBase.td"
+
+def SpatialLayoutCapabilityInterface : OpInterface<"SpatialLayoutCapabilityInterface"> {
+  let description = [{
+    Contract implemented by logical Spatial planning operations that expose
+    their legal physical layout alternatives to the Spatial planner.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      "Return legal physical layout alternatives for this operation and its current operand layouts.",
+      "::llvm::SmallVector<::onnx_mlir::spatial::LayoutAlternative>",
+      "getLayoutAlternatives",
+      (ins "const ::onnx_mlir::spatial::SpatialTargetInfo &":$target,
+           "::llvm::ArrayRef<::onnx_mlir::spatial::PhysicalLayout>":$operandLayouts)>
+  ];
+
+  let cppNamespace = "::onnx_mlir::spatial";
+}
+
+#endif

src/PIM/Dialect/Spatial/SpatialOps.cpp

@@ -4,6 +4,8 @@
 
 #include <string>
 
+#include "mlir/IR/DialectImplementation.h"
+
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -47,7 +49,7 @@ bool isCanonicalContiguousRowMajorFragmentAssembly(SpatBlueprintOp blueprint) {
   auto fragmentStrides = blueprint.getFragmentStrides();
   if (!logicalType || !physicalType || !logicalType.hasStaticShape() || !physicalType.hasStaticShape()
       || logicalType.getRank() < 2 || !blueprint.getFragments().empty()
-      || blueprint.getMode() != "fragment_assembly" || !operandIndices || !sourceSlots || !sourceOffsets
+      || !isFragmentAssembly(blueprint.getMode()) || !operandIndices || !sourceSlots || !sourceOffsets
       || !fragmentStrides)
     return false;
 
@@ -438,6 +440,11 @@ OpResult SpatInParallelOp::getParentResult(int64_t idx) {
 llvm::iterator_range<Block::iterator> SpatInParallelOp::getYieldingOps() { return getRegion().front().getOperations(); }
 
 void SpatialDialect::initialize() {
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialAttributes.cpp.inc"
+
+      >();
   addTypes<
 #define GET_TYPEDEF_LIST
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialTypes.cpp.inc"
@@ -459,6 +466,33 @@ void SpatialDialect::initialize() {
 #define GET_OP_CLASSES
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.cpp.inc"
 
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialEnums.cpp.inc"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialLayoutInterface.cpp.inc"
+
+#define GET_ATTRDEF_CLASSES
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialAttributes.cpp.inc"
+
+namespace onnx_mlir {
+namespace spatial {
+
+Attribute SpatialDialect::parseAttribute(DialectAsmParser& parser, Type type) const {
+  StringRef attrTag;
+  if (Attribute attr; generatedAttributeParser(parser, &attrTag, type, attr).has_value())
+    return attr;
+  parser.emitError(parser.getCurrentLocation()) << "unknown attribute `" << attrTag
+                                                << "` in dialect `spat`";
+  return {};
+}
+
+void SpatialDialect::printAttribute(Attribute attr, DialectAsmPrinter& printer) const {
+  if (succeeded(generatedAttributePrinter(attr, printer)))
+    return;
+  llvm_unreachable("unknown attribute in Spatial dialect");
+}
+
+} // namespace spatial
+} // namespace onnx_mlir
+
 #define GET_TYPEDEF_CLASSES
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialDialect.cpp.inc"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialTypes.cpp.inc"

src/PIM/Dialect/Spatial/SpatialOps.hpp

@@ -11,15 +11,37 @@
 #include "mlir/Interfaces/ParallelCombiningOpInterface.h"
 
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 #include <map>
 #include <optional>
 #include <string>
 #include <tuple>
 
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetInfo.hpp"
+
 /// Include the auto-generated header files containing the declarations
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialDialect.hpp.inc"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialEnums.hpp.inc"
+
+namespace onnx_mlir {
+namespace spatial {
+
+struct LayoutAlternative {
+  llvm::SmallVector<PhysicalLayout> operandLayouts;
+  PhysicalLayout resultLayout = PhysicalLayout::DenseNCHW;
+  int64_t intrinsicCost = 0;
+};
+
+} // namespace spatial
+} // namespace onnx_mlir
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialLayoutInterface.hpp.inc"
+
+#define GET_ATTRDEF_CLASSES
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialAttributes.hpp.inc"
 
 #define GET_TYPEDEF_CLASSES
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialTypes.hpp.inc"
@@ -31,6 +53,48 @@ namespace onnx_mlir {
 namespace spatial {
 
 inline constexpr llvm::StringLiteral kContiguousRowMajorFragments = "contiguous_row_major_fragments";
+inline constexpr llvm::StringLiteral kSelectedLayoutAttrName = "spat.selected_layout";
+
+inline LogicalLayoutAttr getNCHWLayout(mlir::MLIRContext* context) {
+  return LogicalLayoutAttr::get(context, LogicalLayout::NCHW);
+}
+
+inline PhysicalLayoutAttr getDenseNCHWLayout(mlir::MLIRContext* context) {
+  return PhysicalLayoutAttr::get(context, PhysicalLayout::DenseNCHW);
+}
+
+inline PhysicalLayoutAttr getNCHWRowStripLayout(mlir::MLIRContext* context) {
+  return PhysicalLayoutAttr::get(context, PhysicalLayout::NCHWRowStrip);
+}
+
+inline PhysicalLayoutAttr getNHWCRowStripLayout(mlir::MLIRContext* context) {
+  return PhysicalLayoutAttr::get(context, PhysicalLayout::NHWCRowStrip);
+}
+
+inline PhysicalLayoutAttr getFragmentedLayout(mlir::MLIRContext* context) {
+  return PhysicalLayoutAttr::get(context, PhysicalLayout::Fragmented);
+}
+
+inline BlueprintModeAttr getFragmentAssemblyMode(mlir::MLIRContext* context) {
+  return BlueprintModeAttr::get(context, BlueprintMode::FragmentAssembly);
+}
+
+inline BlueprintModeAttr getPhysicalViewMode(mlir::MLIRContext* context) {
+  return BlueprintModeAttr::get(context, BlueprintMode::PhysicalView);
+}
+
+inline bool isPhysicalView(std::optional<BlueprintMode> mode) {
+  return mode && *mode == BlueprintMode::PhysicalView;
+}
+
+inline bool isFragmentAssembly(std::optional<BlueprintMode> mode) {
+  return mode && *mode == BlueprintMode::FragmentAssembly;
+}
+
+inline std::optional<PhysicalLayout> getSelectedPhysicalLayout(mlir::Operation* op) {
+  auto attr = op->getAttrOfType<PhysicalLayoutAttr>(kSelectedLayoutAttrName);
+  return attr ? std::optional<PhysicalLayout>(attr.getValue()) : std::nullopt;
+}
 
 bool hasCanonicalContiguousRowMajorFragments(mlir::RankedTensorType logicalType,
                                              llvm::ArrayRef<int64_t> offsets,

src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp

@@ -616,7 +616,7 @@ void SpatBlueprintOp::print(OpAsmPrinter& printer) {
   printer << " sizes ";
   printCompressedIntegerList(printer, getFragmentSizes());
   printer << " map " << getIndexMap();
-  if (std::optional<StringRef> mode = getMode())
+  if (auto mode = getMode())
     printer << " mode " << *mode;
   if (std::optional<ArrayRef<int64_t>> operandIndices = getFragmentOperandIndices()) {
     printer << " operandIndices ";
@@ -712,14 +712,25 @@ ParseResult SpatBlueprintOp::parse(OpAsmParser& parser, OperationState& result)
   if (operands.size() != operandTypes.size())
     return parser.emitError(parser.getCurrentLocation(), "number of fragment operands and types must match");
 
+  auto logicalLayoutValue = symbolizeLogicalLayout(logicalLayout.getValue());
+  auto physicalLayoutValue = symbolizePhysicalLayout(physicalLayout.getValue());
+  if (!logicalLayoutValue || !physicalLayoutValue)
+    return parser.emitError(parser.getCurrentLocation(), "unknown Blueprint layout");
+  std::optional<BlueprintMode> modeValue;
+  if (mode) {
+    modeValue = symbolizeBlueprintMode(mode.getValue());
+    if (!modeValue)
+      return parser.emitError(parser.getCurrentLocation(), "unknown Blueprint mode");
+  }
+
   auto& builder = parser.getBuilder();
-  result.addAttribute("logicalLayout", logicalLayout);
+  result.addAttribute("logicalLayout", LogicalLayoutAttr::get(builder.getContext(), *logicalLayoutValue));
-  result.addAttribute("physicalLayout", physicalLayout);
+  result.addAttribute("physicalLayout", PhysicalLayoutAttr::get(builder.getContext(), *physicalLayoutValue));
   result.addAttribute("fragmentOffsets", builder.getDenseI64ArrayAttr(fragmentOffsets));
   result.addAttribute("fragmentSizes", builder.getDenseI64ArrayAttr(fragmentSizes));
   result.addAttribute("indexMap", indexMap);
-  if (mode)
+  if (modeValue)
-    result.addAttribute("mode", mode);
+    result.addAttribute("mode", BlueprintModeAttr::get(builder.getContext(), *modeValue));
   if (!fragmentOperandIndices.empty())
     result.addAttribute("fragmentOperandIndices", builder.getDenseI64ArrayAttr(fragmentOperandIndices));
   if (!fragmentSourceSlots.empty())

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -17,7 +17,6 @@
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -405,7 +404,7 @@ static LogicalResult verifyConcatTypes(Operation* op, ValueRange inputs, Value o
 LogicalResult SpatConcatOp::verify() { return verifyConcatTypes(getOperation(), getInputs(), getOutput(), getAxis()); }
 
 LogicalResult SpatConcatPlanOp::verify() {
-  if (getLogicalLayout() != "nchw")
+  if (getLogicalLayout() != LogicalLayout::NCHW)
     return emitError("requires logicalLayout = \"nchw\"");
   auto outputType = dyn_cast<RankedTensorType>(getOutput().getType());
   if (!outputType || !outputType.hasStaticShape() || outputType.getRank() != 4)
@@ -415,11 +414,11 @@ LogicalResult SpatConcatPlanOp::verify() {
   return verifyConcatTypes(getOperation(), getInputs(), getOutput(), getAxis());
 }
 
-static bool isKnownLogicalLayout(StringRef layout) { return layout == "nchw"; }
+static bool isKnownLogicalLayout(LogicalLayout layout) { return layout == LogicalLayout::NCHW; }
 
-static bool isKnownPhysicalLayout(StringRef layout) {
+static bool isKnownPhysicalLayout(PhysicalLayout layout) {
-  return layout == "dense_nchw" || layout == "nchw_row_strip" || layout == "nhwc_row_strip"
+  return layout == PhysicalLayout::DenseNCHW || layout == PhysicalLayout::NCHWRowStrip
-         || layout == "fragmented";
+         || layout == PhysicalLayout::NHWCRowStrip || layout == PhysicalLayout::Fragmented;
 }
 
 static LogicalResult verifyPlanTensorTypes(Operation* op, Value input, Value output, StringRef kind) {
@@ -489,7 +488,7 @@ LogicalResult SpatResizeNearestPlanOp::verify() {
   if (!inputType.hasStaticShape() || !outputType.hasStaticShape()
       || inputType.getRank() != 4 || outputType.getRank() != 4)
     return emitError("requires static rank-4 input and output tensors");
-  if (getLogicalLayout() != "nchw")
+  if (getLogicalLayout() != LogicalLayout::NCHW)
     return emitError("requires logical layout \"nchw\"");
   if (llvm::any_of(inputType.getShape(), [](int64_t dim) { return dim <= 0; })
       || llvm::any_of(outputType.getShape(), [](int64_t dim) { return dim <= 0; }))
@@ -505,7 +504,7 @@ LogicalResult SpatMaxPool2DPlanOp::verify() {
   if (!inputType.hasStaticShape() || !outputType.hasStaticShape() || inputType.getRank() != 4
       || outputType.getRank() != 4)
     return emitError("requires static rank-4 input and output tensors");
-  if (getLogicalLayout() != "nchw")
+  if (getLogicalLayout() != LogicalLayout::NCHW)
     return emitError("requires logical layout \"nchw\"");
   if (getKernelShape().size() != 2 || getStrides().size() != 2 || getDilations().size() != 2)
     return emitError("requires two kernel, stride, and dilation values");
@@ -526,7 +525,7 @@ LogicalResult SpatGlobalAveragePoolPlanOp::verify() {
   if (!inputType.hasStaticShape() || !outputType.hasStaticShape() || inputType.getRank() != 4
       || outputType.getRank() != 4)
     return emitError("requires static rank-4 input and output tensors");
-  if (getLogicalLayout() != "nchw")
+  if (getLogicalLayout() != LogicalLayout::NCHW)
     return emitError("requires logical layout \"nchw\"");
   if (inputType.getDimSize(0) != 1 || outputType.getDimSize(0) != 1
       || inputType.getDimSize(1) != outputType.getDimSize(1)
@@ -552,7 +551,7 @@ LogicalResult SpatBiasAddPlanOp::verify() {
     return emitError("requires matching input and output tensor types");
   if (outputType.getRank() != 4)
     return emitError("requires rank-4 input/output tensors");
-  if (getLogicalLayout() != "nchw")
+  if (getLogicalLayout() != LogicalLayout::NCHW)
     return emitError("requires logical layout \"nchw\"");
   if (biasType.getElementType() != outputType.getElementType())
     return emitError("requires bias element type to match the output element type");
@@ -580,14 +579,31 @@ LogicalResult SpatAddPlanOp::verify() {
     return emitError("requires matching operand and output tensor types");
   if (outputType.getRank() != 4)
     return emitError("requires rank-4 operands and output");
-  if (getLogicalLayout() != "nchw")
+  if (getLogicalLayout() != LogicalLayout::NCHW)
     return emitError("requires logical layout \"nchw\"");
   return success();
 }
 
 LogicalResult SpatBlueprintOp::verify() {
   auto modeAttr = getModeAttr();
-  bool isFragmentAssembly = modeAttr && modeAttr.getValue() == "fragment_assembly";
+  bool isPhysicalView = modeAttr && modeAttr.getValue() == BlueprintMode::PhysicalView;
+  bool isFragmentAssembly = modeAttr && modeAttr.getValue() == BlueprintMode::FragmentAssembly;
+  if (isPhysicalView) {
+    auto inputType = dyn_cast<RankedTensorType>(getInput().getType());
+    auto outputType = dyn_cast<RankedTensorType>(getOutput().getType());
+    if (!inputType || !outputType || !inputType.hasStaticShape() || !outputType.hasStaticShape())
+      return emitError("physical view requires static ranked tensor input and output");
+    if (inputType.getRank() != outputType.getRank() + 1)
+      return emitError("physical view requires one leading physical slot dimension");
+    if (!getFragments().empty() || !getFragmentOffsets().empty() || !getFragmentSizes().empty()
+        || getFragmentOperandIndicesAttr() || getFragmentSourceSlotsAttr()
+        || getFragmentSourceOffsetsAttr() || getFragmentStridesAttr()
+        || getConflictPolicyAttr() || getCoveragePolicyAttr())
+      return emitError("physical view does not accept fragment assembly metadata");
+    if (!isKnownLogicalLayout(getLogicalLayout()) || !isKnownPhysicalLayout(getPhysicalLayout()))
+      return emitError("physical view requires known logical and physical layouts");
+    return success();
+  }
   if (!isFragmentAssembly && failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.blueprint")))
     return failure();
   if (!isKnownLogicalLayout(getLogicalLayout()))

src/PIM/Dialect/Spatial/SpatialTargetInfo.hpp

@@ -0,0 +1,37 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+namespace onnx_mlir::spatial {
+
+struct MatrixUnitShape {
+  size_t rows = 128;
+  size_t columns = 128;
+};
+
+enum class ConvLoweringStrategy : uint8_t {
+  Auto,
+  Legacy,
+  Depthwise,
+  PackedIm2Col,
+  StreamedPatch,
+  StreamedPacked,
+  OutputChannelTiled,
+  InputKTiled,
+  Tiled2D,
+};
+
+struct SpatialTargetInfo {
+  MatrixUnitShape matrixShape;
+  size_t matrixUnitsPerProcessor = 64;
+  size_t processorCount = 1;
+  size_t vectorWidth = 16;
+
+  uint64_t convIm2colMaxElements = 1ull << 20;
+  uint64_t convStreamChunkPositions = 1024;
+  ConvLoweringStrategy convLoweringStrategy = ConvLoweringStrategy::Auto;
+  bool useExperimentalConvImplementation = false;
+};
+
+} // namespace onnx_mlir::spatial

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

@@ -387,7 +387,7 @@ static void collectClosure(Value value, Block &body, const DeferredInputPlan &pl
 
 bool isDeferredFragmentAssemblyInput(Value input, size_t processorCount) {
   auto blueprint = input.getDefiningOp<SpatBlueprintOp>();
-  if (!blueprint || blueprint.getMode() != "fragment_assembly")
+  if (!blueprint || !isFragmentAssembly(blueprint.getMode()))
     return false;
   return llvm::all_of(getBlueprintFragments(blueprint), [&](Value fragment) {
     return getProducerValueRef(fragment, nullptr, processorCount).has_value();

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

@@ -395,7 +395,7 @@ static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& p
 
 static LogicalResult
 retargetBlueprint(DeferredTransferPlan& plan, SpatBlueprintOp blueprint, GraphBatchPublicationCache& publicationCache) {
-  if (blueprint.getMode() != "fragment_assembly")
+  if (!isFragmentAssembly(blueprint.getMode()))
     return success();
   bool escapesScheduledGraph = llvm::any_of(
     blueprint.getOutput().getUses(), [](OpOperand &use) {

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

@@ -1,128 +0,0 @@
-#include "mlir/Pass/Pass.h"
-
-#include "DeferredCommunicationRealization.hpp"
-#include "ScheduledComputeMaterialization.hpp"
-#include "ScheduledComputeReport.hpp"
-#include "ScheduledComputeVerification.hpp"
-#include "Scheduling/MergeSchedulingAnalysis.hpp"
-#include "SpatialDataflowCsvExporter.hpp"
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
-
-using namespace mlir;
-
-namespace onnx_mlir {
-namespace spatial {
-namespace {
-
-struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, OperationPass<ModuleOp>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeComputeNodesPass)
-
-  MergeComputeNodesPass() = default;
-  explicit MergeComputeNodesPass(const SchedulingTarget& schedulingTarget)
-  : target(schedulingTarget), hasTarget(true) {}
-
-  StringRef getArgument() const override { return "pim-merge-compute-nodes"; }
-  StringRef getDescription() const override {
-    return "Materialize scheduled Spatial compute with deferred communication placeholders.";
-  }
-
-  void runOnOperation() override {
-    ModuleOp moduleOp = getOperation();
-    if (!hasTarget || target.processorCount == 0 || target.residentWeightCapacity == 0 || target.transferWidthBytes == 0
-        || target.interProcessorLatencyNs.size() != target.processorCount * target.processorCount
-        || (target.processorCount > 1 && target.averageInterProcessorLatencyNs == 0)) {
-      moduleOp.emitError("MergeComputeNodes requires an explicit valid Spatial scheduling target");
-      signalPassFailure();
-      return;
-    }
-    auto entryFunc = getPimEntryFunc(moduleOp);
-    if (failed(entryFunc)) {
-      moduleOp.emitError("failed to locate the PIM entry function during MergeComputeNodes");
-      signalPassFailure();
-      return;
-    }
-
-    func::FuncOp funcOp = *entryFunc;
-    MergeScheduleResult logicalSchedule = MergeSchedulingAnalysis(funcOp, target).getResult();
-    PatternRewriter rewriter(moduleOp.getContext());
-    FailureOr<ScheduledComputeMaterializationResult> materialization =
-      materializeScheduledCompute(funcOp, logicalSchedule, rewriter);
-    if (failed(materialization)) {
-      signalPassFailure();
-      return;
-    }
-    // Phase 1 is intentionally dumped before its verifier: malformed deferred
-    // payloads must be diagnosed from the producer-owned body.
-    dumpModule(moduleOp, "spatial3_scheduled_no_comm", /*assumeVerified=*/true);
-    if (failed(verifyMaterializedScheduleMapping(funcOp,
-                                                 logicalSchedule,
-                                                 materialization->peftClassPlans,
-                                                 materialization->graphComputeToBlockMap,
-                                                 materialization->materializedSchedules))) {
-      moduleOp.emitError("scheduled Spatial materialization mapping verification failed");
-      signalPassFailure();
-      return;
-    }
-    if (failed(verifyDeferredTransferPhase1Invariants(funcOp))) {
-      moduleOp.emitError("scheduled Spatial deferred communication verification failed");
-      signalPassFailure();
-      return;
-    }
-    if (failed(verifyScheduledMaterializationRecords(materialization->materializedSchedules))) {
-      moduleOp.emitError("scheduled Spatial materialization record verification failed");
-      signalPassFailure();
-      return;
-    }
-    if (failed(verifyScheduledSpatialInvariants(funcOp))) {
-      moduleOp.emitError("scheduled Spatial phase 1 verification failed");
-      signalPassFailure();
-      return;
-    }
-
-    SpatialDataflowExportStage exportMode = getSpatialDataflowExportStage();
-    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial3)
-        && failed(exportSpatialDataflowCsvScheduled(
-          funcOp, materialization->materializedSchedules, "spatial3_scheduled_no_comm", "spatial3"))) {
-      signalPassFailure();
-      return;
-    }
-
-    dumpScheduledComputeReport(
-      moduleOp, funcOp, logicalSchedule, materialization->peftClassPlans, materialization->materializedSchedules);
-    if (failed(realizeDeferredCommunication(funcOp, *materialization, target))) {
-      moduleOp.emitError("MergeComputeNodes phase 2 communication realization failed");
-      signalPassFailure();
-      return;
-    }
-    dumpModule(moduleOp, "spatial4_scheduled", /*assumeVerified=*/true);
-    if (failed(verifyScheduledResultsLive(materialization->materializedSchedules))
-        || failed(verifyScheduledSpatialInvariants(funcOp))) {
-      moduleOp.emitError("scheduled Spatial phase 2 verification failed");
-      signalPassFailure();
-      return;
-    }
-    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial4)
-        && failed(exportSpatialDataflowCsvScheduled(
-          funcOp, materialization->materializedSchedules, "spatial4_scheduled", "spatial4"))) {
-      signalPassFailure();
-    }
-  }
-
-private:
-  SchedulingTarget target;
-  bool hasTarget = false;
-};
-
-} // namespace
-} // namespace spatial
-
-std::unique_ptr<Pass> createMergeComputeNodesPass() { return std::make_unique<spatial::MergeComputeNodesPass>(); }
-
-std::unique_ptr<Pass> createMergeComputeNodesPass(const spatial::SchedulingTarget& target) {
-  return std::make_unique<spatial::MergeComputeNodesPass>(target);
-}
-
-} // namespace onnx_mlir

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

@@ -24,7 +24,7 @@ bool requiresScheduledPublication(Value value, DenseSet<Value> &visited) {
             SpatDeferredCommunicationOp>(user))
       return false;
     auto blueprint = dyn_cast<SpatBlueprintOp>(user);
-    return !blueprint || blueprint.getMode() != "fragment_assembly"
+    return !blueprint || !isFragmentAssembly(blueprint.getMode())
       || requiresScheduledPublication(blueprint.getOutput(), visited);
   });
 }

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

@@ -0,0 +1,279 @@
+#include "ScheduledSpatialPasses.hpp"
+
+#include "mlir/Pass/Pass.h"
+
+#include "DeferredCommunicationRealization.hpp"
+#include "ScheduledComputeReport.hpp"
+#include "ScheduledComputeVerification.hpp"
+#include "SpatialDataflowCsvExporter.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace spatial {
+namespace {
+
+static bool hasValidTarget(const SchedulingTarget& target) {
+  return target.processorCount != 0 && target.residentWeightCapacity != 0
+      && target.transferWidthBytes != 0
+      && target.interProcessorLatencyNs.size() == target.processorCount * target.processorCount
+      && (target.processorCount == 1 || target.averageInterProcessorLatencyNs != 0);
+}
+
+static FailureOr<func::FuncOp> requireEntry(ModuleOp moduleOp, StringRef passName) {
+  auto entry = getPimEntryFunc(moduleOp);
+  if (failed(entry)) {
+    moduleOp.emitError("failed to locate the PIM entry function during ") << passName;
+    return failure();
+  }
+  return *entry;
+}
+
+static LogicalResult requireState(ModuleOp moduleOp,
+                                   const std::shared_ptr<ScheduledSpatialState>& state,
+                                   StringRef passName) {
+  if (state && state->logicalSchedule && state->materialization)
+    return success();
+  moduleOp.emitError() << passName << " requires scheduling state from ScheduleSpatialGraph";
+  return failure();
+}
+
+struct ScheduleSpatialGraphPass final
+    : PassWrapper<ScheduleSpatialGraphPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ScheduleSpatialGraphPass)
+
+  ScheduleSpatialGraphPass() = default;
+  ScheduleSpatialGraphPass(const SchedulingTarget& target,
+                           std::shared_ptr<ScheduledSpatialState> state)
+      : target(target), state(std::move(state)), hasTarget(true) {}
+
+  StringRef getArgument() const override { return "schedule-spatial-graph"; }
+  StringRef getDescription() const override {
+    return "Schedule Spatial graph computes and materialize deferred communication boundaries.";
+  }
+
+  void runOnOperation() override {
+    ModuleOp moduleOp = getOperation();
+    if (!hasTarget || !hasValidTarget(target) || !state) {
+      moduleOp.emitError("ScheduleSpatialGraph requires an explicit target and shared pass state");
+      signalPassFailure();
+      return;
+    }
+    auto entry = requireEntry(moduleOp, "ScheduleSpatialGraph");
+    if (failed(entry)) {
+      signalPassFailure();
+      return;
+    }
+
+    MergeSchedulingAnalysis analysis(*entry, target);
+    MergeScheduleResult schedule = std::move(analysis.getResult());
+    PatternRewriter rewriter(moduleOp.getContext());
+    FailureOr<ScheduledComputeMaterializationResult> materialization =
+        materializeScheduledCompute(*entry, schedule, rewriter);
+    if (failed(materialization)) {
+      signalPassFailure();
+      return;
+    }
+
+    state->logicalSchedule = std::move(schedule);
+    state->materialization = std::move(*materialization);
+    dumpModule(moduleOp, "spatial3_scheduled_no_comm", /*assumeVerified=*/true);
+  }
+
+private:
+  SchedulingTarget target;
+  std::shared_ptr<ScheduledSpatialState> state;
+  bool hasTarget = false;
+};
+
+struct VerifyScheduledSpatialPass final
+    : PassWrapper<VerifyScheduledSpatialPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VerifyScheduledSpatialPass)
+
+  explicit VerifyScheduledSpatialPass(std::shared_ptr<ScheduledSpatialState> state = {})
+      : state(std::move(state)) {}
+
+  StringRef getArgument() const override { return "verify-scheduled-spatial"; }
+  StringRef getDescription() const override {
+    return "Verify scheduled Spatial compute, deferred communication, and materialization records.";
+  }
+
+  void runOnOperation() override {
+    ModuleOp moduleOp = getOperation();
+    if (failed(requireState(moduleOp, state, "VerifyScheduledSpatial"))) {
+      signalPassFailure();
+      return;
+    }
+    auto entry = requireEntry(moduleOp, "VerifyScheduledSpatial");
+    if (failed(entry)) {
+      signalPassFailure();
+      return;
+    }
+    const auto& schedule = *state->logicalSchedule;
+    const auto& materialization = *state->materialization;
+    if (failed(verifyMaterializedScheduleMapping(
+            *entry, schedule, materialization.peftClassPlans,
+            materialization.graphComputeToBlockMap,
+            materialization.materializedSchedules))
+        || failed(verifyDeferredTransferPhase1Invariants(*entry))
+        || failed(verifyScheduledMaterializationRecords(materialization.materializedSchedules))
+        || failed(verifyScheduledSpatialInvariants(*entry))) {
+      moduleOp.emitError("scheduled Spatial phase verification failed");
+      signalPassFailure();
+      return;
+    }
+    SpatialDataflowExportStage exportMode = getSpatialDataflowExportStage();
+    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial3)
+        && failed(exportSpatialDataflowCsvScheduled(
+            *entry, materialization.materializedSchedules,
+            "spatial3_scheduled_no_comm", "spatial3"))) {
+      signalPassFailure();
+      return;
+    }
+    dumpScheduledComputeReport(
+        moduleOp, *entry, schedule, materialization.peftClassPlans,
+        materialization.materializedSchedules);
+  }
+
+private:
+  std::shared_ptr<ScheduledSpatialState> state;
+};
+
+struct RealizeSpatialCommunicationPass final
+    : PassWrapper<RealizeSpatialCommunicationPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RealizeSpatialCommunicationPass)
+
+  RealizeSpatialCommunicationPass() = default;
+  RealizeSpatialCommunicationPass(const SchedulingTarget& target,
+                                  std::shared_ptr<ScheduledSpatialState> state)
+      : target(target), state(std::move(state)), hasTarget(true) {}
+
+  StringRef getArgument() const override { return "realize-spatial-communication"; }
+  StringRef getDescription() const override {
+    return "Realize deferred Spatial communication after scheduled graph verification.";
+  }
+
+  void runOnOperation() override {
+    ModuleOp moduleOp = getOperation();
+    if (!hasTarget || !hasValidTarget(target)
+        || failed(requireState(moduleOp, state, "RealizeSpatialCommunication"))) {
+      signalPassFailure();
+      return;
+    }
+    auto entry = requireEntry(moduleOp, "RealizeSpatialCommunication");
+    if (failed(entry)) {
+      signalPassFailure();
+      return;
+    }
+    if (failed(realizeDeferredCommunication(*entry, *state->materialization, target))) {
+      moduleOp.emitError("Spatial communication realization failed");
+      signalPassFailure();
+      return;
+    }
+    dumpModule(moduleOp, "spatial4_scheduled", /*assumeVerified=*/true);
+  }
+
+private:
+  SchedulingTarget target;
+  std::shared_ptr<ScheduledSpatialState> state;
+  bool hasTarget = false;
+};
+
+struct VerifyRealizedSpatialPass final
+    : PassWrapper<VerifyRealizedSpatialPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VerifyRealizedSpatialPass)
+
+  explicit VerifyRealizedSpatialPass(std::shared_ptr<ScheduledSpatialState> state = {})
+      : state(std::move(state)) {}
+
+  StringRef getArgument() const override { return "verify-realized-spatial"; }
+  StringRef getDescription() const override {
+    return "Verify realized Spatial communication and scheduled result liveness.";
+  }
+
+  void runOnOperation() override {
+    ModuleOp moduleOp = getOperation();
+    if (failed(requireState(moduleOp, state, "VerifyRealizedSpatial"))) {
+      signalPassFailure();
+      return;
+    }
+    auto entry = requireEntry(moduleOp, "VerifyRealizedSpatial");
+    if (failed(entry)) {
+      signalPassFailure();
+      return;
+    }
+    const auto& records = state->materialization->materializedSchedules;
+    bool deferredRemains = false;
+    (*entry).walk([&](SpatDeferredCommunicationOp deferred) {
+      if (deferredRemains)
+        return;
+      deferred.emitOpError("realized Spatial graph still contains deferred communication");
+      deferredRemains = true;
+    });
+    if (deferredRemains
+        || failed(verifyScheduledResultsLive(records))
+        || failed(verifyScheduledSpatialInvariants(*entry))) {
+      moduleOp.emitError("realized Spatial communication verification failed");
+      signalPassFailure();
+      return;
+    }
+    SpatialDataflowExportStage exportMode = getSpatialDataflowExportStage();
+    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial4)
+        && failed(exportSpatialDataflowCsvScheduled(
+            *entry, records, "spatial4_scheduled", "spatial4")))
+      signalPassFailure();
+  }
+
+private:
+  std::shared_ptr<ScheduledSpatialState> state;
+};
+
+} // namespace
+
+std::unique_ptr<Pass> createScheduleSpatialGraphPass() {
+  return std::make_unique<ScheduleSpatialGraphPass>();
+}
+
+std::unique_ptr<Pass> createScheduleSpatialGraphPass(const SchedulingTarget& target) {
+  return std::make_unique<ScheduleSpatialGraphPass>(
+      target, std::make_shared<ScheduledSpatialState>());
+}
+
+std::unique_ptr<Pass> createScheduleSpatialGraphPass(
+    const SchedulingTarget& target, std::shared_ptr<ScheduledSpatialState> state) {
+  return std::make_unique<ScheduleSpatialGraphPass>(target, std::move(state));
+}
+
+std::unique_ptr<Pass> createVerifyScheduledSpatialPass() {
+  return std::make_unique<VerifyScheduledSpatialPass>();
+}
+
+std::unique_ptr<Pass> createVerifyScheduledSpatialPass(
+    std::shared_ptr<ScheduledSpatialState> state) {
+  return std::make_unique<VerifyScheduledSpatialPass>(std::move(state));
+}
+
+std::unique_ptr<Pass> createRealizeSpatialCommunicationPass() {
+  return std::make_unique<RealizeSpatialCommunicationPass>();
+}
+
+std::unique_ptr<Pass> createRealizeSpatialCommunicationPass(
+    const SchedulingTarget& target, std::shared_ptr<ScheduledSpatialState> state) {
+  return std::make_unique<RealizeSpatialCommunicationPass>(target, std::move(state));
+}
+
+std::unique_ptr<Pass> createVerifyRealizedSpatialPass() {
+  return std::make_unique<VerifyRealizedSpatialPass>();
+}
+
+std::unique_ptr<Pass> createVerifyRealizedSpatialPass(
+    std::shared_ptr<ScheduledSpatialState> state) {
+  return std::make_unique<VerifyRealizedSpatialPass>(std::move(state));
+}
+
+} // namespace spatial
+} // namespace onnx_mlir

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

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "ScheduledComputeMaterialization.hpp"
+#include "Scheduling/MergeSchedulingAnalysis.hpp"
+
+#include <memory>
+#include <optional>
+
+namespace onnx_mlir::spatial {
+
+struct ScheduledSpatialState {
+  std::optional<MergeScheduleResult> logicalSchedule;
+  std::optional<ScheduledComputeMaterializationResult> materialization;
+};
+
+} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp

@@ -181,7 +181,7 @@ FailureOr<LanePublicationSignatures> buildLanePublicationSignatures(SpatComputeB
 
     for (auto [useIndex, use] : llvm::enumerate(result.getUses())) {
       auto blueprint = dyn_cast<SpatBlueprintOp>(use.getOwner());
-      if (!blueprint || blueprint.getMode() != "fragment_assembly")
+      if (!blueprint || !isFragmentAssembly(blueprint.getMode()))
         continue;
       auto operandIndices = blueprint.getFragmentOperandIndices();
       auto sourceSlots = blueprint.getFragmentSourceSlots();

src/PIM/Dialect/Spatial/Transforms/TrivialGraphComputeMergePass.cpp

@@ -14,7 +14,6 @@
 #include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp"

src/PIM/Pass/PIMPasses.h

@@ -9,18 +9,40 @@
 namespace onnx_mlir {
 namespace spatial {
 struct SchedulingTarget;
+struct ScheduledSpatialState;
+struct SpatialTargetInfo;
+
+std::unique_ptr<mlir::Pass> createScheduleSpatialGraphPass();
+std::unique_ptr<mlir::Pass> createScheduleSpatialGraphPass(const SchedulingTarget& target);
+std::unique_ptr<mlir::Pass> createScheduleSpatialGraphPass(
+  const SchedulingTarget& target,
+  std::shared_ptr<ScheduledSpatialState> state);
+std::unique_ptr<mlir::Pass> createVerifyScheduledSpatialPass();
+std::unique_ptr<mlir::Pass> createVerifyScheduledSpatialPass(
+  std::shared_ptr<ScheduledSpatialState> state);
+std::unique_ptr<mlir::Pass> createRealizeSpatialCommunicationPass();
+std::unique_ptr<mlir::Pass> createRealizeSpatialCommunicationPass(
+  const SchedulingTarget& target,
+  std::shared_ptr<ScheduledSpatialState> state);
+std::unique_ptr<mlir::Pass> createVerifyRealizedSpatialPass();
+std::unique_ptr<mlir::Pass> createVerifyRealizedSpatialPass(
+  std::shared_ptr<ScheduledSpatialState> state);
 }
 
 std::unique_ptr<mlir::Pass> createONNXToSpatialPass();
+std::unique_ptr<mlir::Pass> createONNXToSpatialPass(const spatial::SpatialTargetInfo& target);
 std::unique_ptr<mlir::Pass> createSpatialLayoutPlanningPass();
+std::unique_ptr<mlir::Pass> createSpatialLayoutPlanningPass(const spatial::SpatialTargetInfo& target);
 std::unique_ptr<mlir::Pass> createLowerSpatialPlansPass();
+std::unique_ptr<mlir::Pass> createLowerSpatialPlansPass(const spatial::SpatialTargetInfo& target);
 
 std::unique_ptr<mlir::Pass> createSpatialToPimPass();
 
-std::unique_ptr<mlir::Pass> createPimBufferizationPass();
+std::unique_ptr<mlir::Pass> createPimBufferizationPreparationPass();
+std::unique_ptr<mlir::Pass> createPimOneShotBufferizationPass();
+std::unique_ptr<mlir::Pass> createPimMemoryNormalizationPass();
+std::unique_ptr<mlir::Pass> createPimBufferizationVerificationPass();
 
-std::unique_ptr<mlir::Pass> createMergeComputeNodesPass();
-std::unique_ptr<mlir::Pass> createMergeComputeNodesPass(const spatial::SchedulingTarget& target);
 
 std::unique_ptr<mlir::Pass> createTrivialGraphComputeMergePass();
 std::unique_ptr<mlir::Pass> createTrivialGraphComputeMergePass(

src/PIM/PimAccelerator.cpp

@@ -71,13 +71,19 @@ void PimAccelerator::registerDialects(mlir::DialectRegistry& registry) const {
 
 void PimAccelerator::registerPasses(int optLevel) const {
   LLVM_DEBUG(llvm::dbgs() << "Registering passes for PIM accelerator\n");
-  registerPass(createONNXToSpatialPass);
+  mlir::registerPass([] { return createONNXToSpatialPass(); });
-  registerPass(createSpatialLayoutPlanningPass);
+  mlir::registerPass([] { return createSpatialLayoutPlanningPass(); });
-  registerPass(createLowerSpatialPlansPass);
+  mlir::registerPass([] { return createLowerSpatialPlansPass(); });
   registerPass(createSpatialToPimPass);
-  registerPass(createPimBufferizationPass);
+  registerPass(createPimBufferizationPreparationPass);
+  registerPass(createPimOneShotBufferizationPass);
+  registerPass(createPimMemoryNormalizationPass);
+  registerPass(createPimBufferizationVerificationPass);
   mlir::registerPass([] { return createTrivialGraphComputeMergePass(); });
-  mlir::registerPass([] { return createMergeComputeNodesPass(); });
+  mlir::registerPass([] { return spatial::createScheduleSpatialGraphPass(); });
+  mlir::registerPass([] { return spatial::createVerifyScheduledSpatialPass(); });
+  mlir::registerPass([] { return spatial::createRealizeSpatialCommunicationPass(); });
+  mlir::registerPass([] { return spatial::createVerifyRealizedSpatialPass(); });
   registerPass(createPimHostConstantFoldingPass);
   registerPass(createPimInstructionSelectionPass);
   registerPass(createPimLocalMemoryPlanningPass);

validation/operations/validation_results.csv

@@ -1,169 +1,169 @@
 Operation,Result,Compile,Host mem,Cores mem,Cores,Xbars,Latency,Power,Energy
-add/after_gemm,PASS,0.072 s,0.01 MiB,0.01 MiB,5,4,0.007784 ms,104.703618 mW,815012.960000 pJ
+add/after_gemm,PASS,0.062 s,0.01 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-add/basic,PASS,0.049 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ
+add/basic,PASS,0.053 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-add/broadcast_row,PASS,0.048 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ
+add/broadcast_row,PASS,0.052 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-add/channel_broadcast_1024,PASS,0.051 s,0.02 MiB,0.01 MiB,1,0,0.006913 ms,78.118038 mW,540030.000000 pJ
+add/channel_broadcast_1024,PASS,0.061 s,0.02 MiB,0.01 MiB,1,0,SKIP,SKIP,SKIP
-add/leading_dimension_broadcast,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ
+add/leading_dimension_broadcast,PASS,0.057 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-concat/channel_axis,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,0.000457 ms,78.157549 mW,35718.000000 pJ
+concat/channel_axis,PASS,0.052 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-concat/negative_axis,PASS,0.061 s,0.00 MiB,0.00 MiB,1,0,0.001043 ms,78.092042 mW,81450.000000 pJ
+concat/negative_axis,PASS,0.061 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-concat/three_inputs_channel_axis,PASS,0.054 s,0.00 MiB,0.00 MiB,1,0,0.000644 ms,78.149068 mW,50328.000000 pJ
+concat/three_inputs_channel_axis,PASS,0.053 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-conv/batch_2,PASS,0.059 s,0.00 MiB,0.00 MiB,2,2,0.013694 ms,82.623885 mW,1131451.480000 pJ
+conv/batch_2,PASS,0.056 s,0.00 MiB,0.00 MiB,2,2,SKIP,SKIP,SKIP
-conv/batch_4_pointwise,PASS,0.056 s,0.00 MiB,0.01 MiB,5,4,0.003932 ms,116.078576 mW,456420.960000 pJ
+conv/batch_4_pointwise,PASS,0.058 s,0.00 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-conv/depthwise_1024_channels,PASS,0.081 s,0.19 MiB,0.38 MiB,129,128,0.220751 ms,178.454307 mW,39393966.720000 pJ
+conv/depthwise_1024_channels,PASS,0.085 s,0.19 MiB,0.38 MiB,129,128,SKIP,SKIP,SKIP
-conv/depthwise_grouped,PASS,0.059 s,0.01 MiB,0.00 MiB,5,4,0.006024 ms,108.326521 mW,652558.960000 pJ
+conv/depthwise_grouped,PASS,0.074 s,0.01 MiB,0.00 MiB,5,4,SKIP,SKIP,SKIP
-conv/dilated_3x3,PASS,0.059 s,0.00 MiB,0.00 MiB,3,3,0.004045 ms,110.234541 mW,445898.720000 pJ
+conv/dilated_3x3,PASS,0.067 s,0.00 MiB,0.00 MiB,3,3,SKIP,SKIP,SKIP
-conv/dynamic,PASS,0.059 s,0.00 MiB,0.00 MiB,5,0,0.001835 ms,92.281199 mW,169336.000000 pJ
+conv/dynamic,PASS,0.059 s,0.00 MiB,0.00 MiB,5,0,SKIP,SKIP,SKIP
-conv/explicit_padding,PASS,0.058 s,0.00 MiB,0.00 MiB,4,4,0.004327 ms,115.794768 mW,501043.960000 pJ
+conv/explicit_padding,PASS,0.065 s,0.00 MiB,0.00 MiB,4,4,SKIP,SKIP,SKIP
-conv/grouped_many_groups,PASS,0.505 s,0.05 MiB,0.09 MiB,65,64,0.181845 ms,142.210104 mW,25860196.360000 pJ
+conv/grouped_many_groups,PASS,0.500 s,0.05 MiB,0.09 MiB,65,64,SKIP,SKIP,SKIP
-conv/grouped_two_groups,PASS,0.061 s,0.00 MiB,0.00 MiB,3,2,0.005360 ms,101.459418 mW,543822.480000 pJ
+conv/grouped_two_groups,PASS,0.070 s,0.00 MiB,0.00 MiB,3,2,SKIP,SKIP,SKIP
-conv/huge_pointwise_1024,PASS,0.643 s,0.01 MiB,0.01 MiB,1,64,0.028261 ms,133.488743 mW,3772525.360000 pJ
+conv/huge_pointwise_1024,PASS,0.652 s,0.01 MiB,0.01 MiB,1,64,SKIP,SKIP,SKIP
-conv/huge_pointwise_1024_dynamic,PASS,0.081 s,8.04 MiB,12.61 MiB,168,0,2.627964 ms,169.518697 mW,445489032.000000 pJ
+conv/huge_pointwise_1024_dynamic,PASS,0.083 s,8.04 MiB,12.61 MiB,168,0,SKIP,SKIP,SKIP
-conv/kernel_3x3,PASS,0.056 s,0.00 MiB,0.00 MiB,3,3,0.003091 ms,115.862414 mW,358130.720000 pJ
+conv/kernel_3x3,PASS,0.067 s,0.00 MiB,0.00 MiB,3,3,SKIP,SKIP,SKIP
-conv/kernel_equals_input_spatial,PASS,0.060 s,0.00 MiB,0.00 MiB,1,2,0.008443 ms,83.863849 mW,708062.480000 pJ
+conv/kernel_equals_input_spatial,PASS,0.072 s,0.00 MiB,0.00 MiB,1,2,SKIP,SKIP,SKIP
-conv/large_input_channels_1x1,PASS,0.088 s,0.01 MiB,0.01 MiB,1,8,0.017167 ms,89.416900 mW,1535019.920000 pJ
+conv/large_input_channels_1x1,PASS,0.098 s,0.01 MiB,0.01 MiB,1,8,SKIP,SKIP,SKIP
-conv/large_output_channels_1x1,PASS,0.087 s,0.00 MiB,0.01 MiB,1,8,0.004964 ms,117.628106 mW,583905.920000 pJ
+conv/large_output_channels_1x1,PASS,0.089 s,0.00 MiB,0.01 MiB,1,8,SKIP,SKIP,SKIP
-conv/large_spatial,PASS,0.055 s,0.00 MiB,0.01 MiB,6,6,0.004096 ms,129.015000 mW,528445.440000 pJ
+conv/large_spatial,PASS,0.063 s,0.00 MiB,0.01 MiB,6,6,SKIP,SKIP,SKIP
-conv/multi_channel,PASS,0.056 s,0.00 MiB,0.00 MiB,3,3,0.005148 ms,106.453520 mW,548022.720000 pJ
+conv/multi_channel,PASS,0.068 s,0.00 MiB,0.00 MiB,3,3,SKIP,SKIP,SKIP
-conv/non_square_kernel_1x3,PASS,0.056 s,0.00 MiB,0.00 MiB,5,5,0.004029 ms,123.600943 mW,497988.200000 pJ
+conv/non_square_kernel_1x3,PASS,0.080 s,0.00 MiB,0.00 MiB,5,5,SKIP,SKIP,SKIP
-conv/non_square_kernel_3x1,PASS,0.058 s,0.00 MiB,0.00 MiB,3,3,0.005526 ms,105.464843 mW,582798.720000 pJ
+conv/non_square_kernel_3x1,PASS,0.068 s,0.00 MiB,0.00 MiB,3,3,SKIP,SKIP,SKIP
-conv/non_uniform_stride,PASS,0.062 s,0.00 MiB,0.00 MiB,4,4,0.005808 ms,110.081433 mW,639352.960000 pJ
+conv/non_uniform_stride,PASS,0.068 s,0.00 MiB,0.00 MiB,4,4,SKIP,SKIP,SKIP
-conv/pointwise_1x1,PASS,0.057 s,0.00 MiB,0.00 MiB,4,4,0.004539 ms,114.835858 mW,521239.960000 pJ
+conv/pointwise_1x1,PASS,0.071 s,0.00 MiB,0.00 MiB,4,4,SKIP,SKIP,SKIP
-conv/pointwise_tiled_chain,PASS,0.771 s,0.01 MiB,0.02 MiB,2,80,0.084437 ms,102.289307 mW,8637002.200000 pJ
+conv/pointwise_tiled_chain,PASS,0.911 s,0.01 MiB,0.02 MiB,2,80,SKIP,SKIP,SKIP
-conv/real_asymmetric_padding,PASS,0.055 s,0.00 MiB,0.00 MiB,4,4,0.005232 ms,111.870214 mW,585304.960000 pJ
+conv/real_asymmetric_padding,PASS,0.057 s,0.00 MiB,0.00 MiB,4,4,SKIP,SKIP,SKIP
-conv/relu_conv_store,PASS,0.076 s,0.02 MiB,0.10 MiB,32,32,0.064390 ms,243.916397 mW,15705776.800000 pJ
+conv/relu_conv_store,PASS,0.070 s,0.02 MiB,0.10 MiB,32,32,SKIP,SKIP,SKIP
-conv/same_lower_3x3,PASS,0.057 s,0.00 MiB,0.00 MiB,5,5,0.004700 ms,119.232170 mW,560391.200000 pJ
+conv/same_lower_3x3,PASS,0.061 s,0.00 MiB,0.00 MiB,5,5,SKIP,SKIP,SKIP
-conv/same_padding_3x3,PASS,0.060 s,0.00 MiB,0.00 MiB,5,5,0.004700 ms,119.232170 mW,560391.200000 pJ
+conv/same_padding_3x3,PASS,0.075 s,0.00 MiB,0.00 MiB,5,5,SKIP,SKIP,SKIP
-conv/simple,PASS,0.055 s,0.00 MiB,0.00 MiB,2,2,0.003148 ms,94.665972 mW,298008.480000 pJ
+conv/simple,PASS,0.055 s,0.00 MiB,0.00 MiB,2,2,SKIP,SKIP,SKIP
-conv/stride_2,PASS,0.053 s,0.00 MiB,0.00 MiB,2,2,0.002827 ms,96.393873 mW,272505.480000 pJ
+conv/stride_2,PASS,0.057 s,0.00 MiB,0.00 MiB,2,2,SKIP,SKIP,SKIP
-conv/with_bias_3x3,PASS,0.055 s,0.00 MiB,0.00 MiB,3,3,0.004898 ms,107.176546 mW,524950.720000 pJ
+conv/with_bias_3x3,PASS,0.067 s,0.00 MiB,0.00 MiB,3,3,SKIP,SKIP,SKIP
-conv/with_constant,PASS,0.073 s,0.00 MiB,0.00 MiB,3,3,0.004273 ms,109.362677 mW,467306.720000 pJ
+conv/with_constant,PASS,0.063 s,0.00 MiB,0.00 MiB,3,3,SKIP,SKIP,SKIP
-conv/without_kernel_shape_attr,PASS,0.057 s,0.00 MiB,0.00 MiB,3,3,0.003091 ms,115.862414 mW,358130.720000 pJ
+conv/without_kernel_shape_attr,PASS,0.059 s,0.00 MiB,0.00 MiB,3,3,SKIP,SKIP,SKIP
-conv/yolo11n_depthwise_head,PASS,0.509 s,4.82 MiB,15.92 MiB,160,720,4.105780 ms,492.826315 mW,2023436428.000010 pJ
+conv/yolo11n_depthwise_head,PASS,0.624 s,4.82 MiB,15.92 MiB,160,720,SKIP,SKIP,SKIP
-conv/yolo11n_heavy,PASS,0.488 s,4.82 MiB,20.66 MiB,160,800,6.275385 ms,418.028199 mW,2623287892.000010 pJ
+conv/yolo11n_heavy,PASS,0.496 s,4.82 MiB,20.66 MiB,160,800,SKIP,SKIP,SKIP
-conv/yolo11n_stem,PASS,1.730 s,12.86 MiB,31.38 MiB,168,488,9.799204 ms,361.075380 mW,3538251304.000010 pJ
+conv/yolo11n_stem,PASS,0.935 s,12.86 MiB,31.38 MiB,168,488,SKIP,SKIP,SKIP
-div/after_gemm,PASS,0.059 s,0.01 MiB,0.01 MiB,5,4,0.007784 ms,104.703618 mW,815012.960000 pJ
+div/after_gemm,PASS,0.060 s,0.01 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-div/basic,PASS,0.048 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ
+div/basic,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-div/channel_broadcast_1024,PASS,0.051 s,0.02 MiB,0.01 MiB,1,0,0.006913 ms,78.118038 mW,540030.000000 pJ
+div/channel_broadcast_1024,PASS,0.049 s,0.02 MiB,0.01 MiB,1,0,SKIP,SKIP,SKIP
-div/leading_dimension_broadcast,PASS,0.063 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ
+div/leading_dimension_broadcast,PASS,0.052 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-div/runtime_scalar_rhs,PASS,0.057 s,0.02 MiB,0.01 MiB,1,0,0.006913 ms,78.118038 mW,540030.000000 pJ
+div/runtime_scalar_rhs,PASS,0.053 s,0.02 MiB,0.01 MiB,1,0,SKIP,SKIP,SKIP
-div/scalar_constant,PASS,0.049 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ
+div/scalar_constant,PASS,0.078 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-gather/3d_input_axis1,PASS,0.050 s,0.00 MiB,0.00 MiB,1,0,0.000589 ms,78.081494 mW,45990.000000 pJ
+gather/3d_input_axis1,PASS,0.059 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-gather/axis0_matrix_indices,PASS,0.050 s,0.00 MiB,0.00 MiB,1,0,0.000697 ms,78.068867 mW,54414.000000 pJ
+gather/axis0_matrix_indices,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-gather/axis1,PASS,0.053 s,0.00 MiB,0.00 MiB,1,0,0.000801 ms,78.059925 mW,62526.000000 pJ
+gather/axis1,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-gather/negative_axis,PASS,0.053 s,0.00 MiB,0.00 MiB,1,0,0.001437 ms,78.033403 mW,112134.000000 pJ
+gather/negative_axis,PASS,0.056 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-gather/negative_indices,PASS,0.052 s,0.00 MiB,0.00 MiB,1,0,0.000376 ms,78.127660 mW,29376.000000 pJ
+gather/negative_indices,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-gemm/alpha_beta,PASS,0.056 s,0.01 MiB,0.01 MiB,5,4,0.007456 ms,105.272125 mW,784908.960000 pJ
+gemm/alpha_beta,PASS,0.078 s,0.01 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-gemm/bias_rank2_broadcast,PASS,0.056 s,0.00 MiB,0.01 MiB,5,4,0.007072 ms,105.979208 mW,749484.960000 pJ
+gemm/bias_rank2_broadcast,PASS,0.055 s,0.00 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-gemm/dynamic,PASS,0.054 s,0.00 MiB,0.00 MiB,5,0,0.002421 ms,91.480793 mW,221475.000000 pJ
+gemm/dynamic,PASS,0.058 s,0.00 MiB,0.00 MiB,5,0,SKIP,SKIP,SKIP
-gemm/dynamic_alpha,PASS,0.056 s,0.00 MiB,0.00 MiB,5,0,0.003262 ms,91.415696 mW,298198.000000 pJ
+gemm/dynamic_alpha,PASS,0.056 s,0.00 MiB,0.00 MiB,5,0,SKIP,SKIP,SKIP
-gemm/dynamic_beta,PASS,0.059 s,0.00 MiB,0.00 MiB,5,0,0.004365 ms,91.316151 mW,398595.000000 pJ
+gemm/dynamic_beta,PASS,0.053 s,0.00 MiB,0.00 MiB,5,0,SKIP,SKIP,SKIP
-gemm/dynamic_bias,PASS,0.056 s,0.00 MiB,0.00 MiB,5,0,0.002665 ms,91.445779 mW,243703.000000 pJ
+gemm/dynamic_bias,PASS,0.062 s,0.00 MiB,0.00 MiB,5,0,SKIP,SKIP,SKIP
-gemm/dynamic_bias_alpha_beta,PASS,0.053 s,0.00 MiB,0.00 MiB,5,0,0.005629 ms,91.279268 mW,513811.000000 pJ
+gemm/dynamic_bias_alpha_beta,PASS,0.114 s,0.00 MiB,0.00 MiB,5,0,SKIP,SKIP,SKIP
-gemm/dynamic_transB,PASS,0.054 s,0.00 MiB,0.00 MiB,5,0,0.001301 ms,91.378171 mW,118883.000000 pJ
+gemm/dynamic_transB,PASS,0.055 s,0.00 MiB,0.00 MiB,5,0,SKIP,SKIP,SKIP
-gemm/huge_1024,PASS,0.157 s,0.01 MiB,0.10 MiB,73,64,0.017522 ms,215.037402 mW,3767885.360000 pJ
+gemm/huge_1024,PASS,0.150 s,0.01 MiB,0.10 MiB,73,64,SKIP,SKIP,SKIP
-gemm/large,PASS,0.067 s,0.02 MiB,0.03 MiB,17,16,0.011229 ms,140.152181 mW,1573768.840000 pJ
+gemm/large,PASS,0.059 s,0.02 MiB,0.03 MiB,17,16,SKIP,SKIP,SKIP
-gemm/large_k_small_n,PASS,0.097 s,0.01 MiB,0.01 MiB,9,8,0.004748 ms,133.481449 mW,633769.920000 pJ
+gemm/large_k_small_n,PASS,0.087 s,0.01 MiB,0.01 MiB,9,8,SKIP,SKIP,SKIP
-gemm/non_square,PASS,0.060 s,0.00 MiB,0.01 MiB,5,4,0.003527 ms,118.958310 mW,419565.960000 pJ
+gemm/non_square,PASS,0.058 s,0.00 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-gemm/scalar_bias,PASS,0.054 s,0.00 MiB,0.01 MiB,5,4,0.007072 ms,105.979208 mW,749484.960000 pJ
+gemm/scalar_bias,PASS,0.058 s,0.00 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-gemm/simple,PASS,0.071 s,0.03 MiB,0.08 MiB,42,40,0.021640 ms,151.774196 mW,3284393.600000 pJ
+gemm/simple,PASS,0.075 s,0.03 MiB,0.08 MiB,42,40,SKIP,SKIP,SKIP
-gemm/small,PASS,0.052 s,0.00 MiB,0.00 MiB,2,2,0.004420 ms,90.144000 mW,398436.480000 pJ
+gemm/small,PASS,0.056 s,0.00 MiB,0.00 MiB,2,2,SKIP,SKIP,SKIP
-gemm/small_k_large_n,PASS,0.089 s,0.01 MiB,0.02 MiB,17,8,0.007962 ms,131.005014 mW,1043061.920000 pJ
+gemm/small_k_large_n,PASS,0.091 s,0.01 MiB,0.02 MiB,17,8,SKIP,SKIP,SKIP
-gemm/transA,PASS,0.057 s,0.00 MiB,0.01 MiB,5,4,0.005762 ms,109.140743 mW,628868.960000 pJ
+gemm/transA,PASS,0.057 s,0.00 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-gemm/transA_transB,PASS,0.053 s,0.00 MiB,0.01 MiB,5,4,0.005762 ms,109.140743 mW,628868.960000 pJ
+gemm/transA_transB,PASS,0.079 s,0.00 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-gemm/transB,PASS,0.062 s,0.00 MiB,0.01 MiB,5,4,0.003527 ms,118.958310 mW,419565.960000 pJ
+gemm/transB,PASS,0.065 s,0.00 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-gemm/transB_with_bias,PASS,0.057 s,0.01 MiB,0.01 MiB,5,4,0.005046 ms,110.546762 mW,557818.960000 pJ
+gemm/transB_with_bias,PASS,0.058 s,0.01 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-gemm/with_bias,PASS,0.056 s,0.01 MiB,0.01 MiB,5,4,0.005562 ms,108.767882 mW,604966.960000 pJ
+gemm/with_bias,PASS,0.054 s,0.01 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-gemv/constant,PASS,0.051 s,0.00 MiB,0.00 MiB,0,0,0.000000 ms,2.000000 mW,0.000000 pJ
+gemv/constant,PASS,0.049 s,0.00 MiB,0.00 MiB,0,0,SKIP,SKIP,SKIP
-gemv/simple,PASS,0.066 s,0.00 MiB,0.01 MiB,6,4,0.005160 ms,111.150380 mW,573535.960000 pJ
+gemv/simple,PASS,0.066 s,0.00 MiB,0.01 MiB,6,4,SKIP,SKIP,SKIP
-gemv/with_heterogeneous_constant,PASS,0.066 s,0.00 MiB,0.01 MiB,6,4,0.005549 ms,109.816536 mW,609371.960000 pJ
+gemv/with_heterogeneous_constant,PASS,0.069 s,0.00 MiB,0.01 MiB,6,4,SKIP,SKIP,SKIP
-gemv/with_homogeneous_constant,PASS,0.063 s,0.00 MiB,0.01 MiB,6,4,0.005549 ms,109.816536 mW,609371.960000 pJ
+gemv/with_homogeneous_constant,PASS,0.065 s,0.00 MiB,0.01 MiB,6,4,SKIP,SKIP,SKIP
-gemv/with_scalar_constant,PASS,0.064 s,0.00 MiB,0.01 MiB,6,4,0.005549 ms,109.816536 mW,609371.960000 pJ
+gemv/with_scalar_constant,PASS,0.100 s,0.00 MiB,0.01 MiB,6,4,SKIP,SKIP,SKIP
-matmul/basic,PASS,0.064 s,0.00 MiB,0.00 MiB,2,2,0.004420 ms,90.144000 mW,398436.480000 pJ
+matmul/basic,PASS,0.056 s,0.00 MiB,0.00 MiB,2,2,SKIP,SKIP,SKIP
-matmul/batched_3d,PASS,0.058 s,0.00 MiB,0.01 MiB,5,4,0.005958 ms,108.588949 mW,646972.960000 pJ
+matmul/batched_3d,PASS,0.057 s,0.00 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-matmul/batched_3d_dynamic,PASS,0.053 s,0.00 MiB,0.00 MiB,4,0,0.001822 ms,92.192645 mW,167975.000000 pJ
+matmul/batched_3d_dynamic,PASS,0.053 s,0.00 MiB,0.00 MiB,4,0,SKIP,SKIP,SKIP
-matmul/batched_left_constant,PASS,0.058 s,0.00 MiB,0.02 MiB,9,8,0.008822 ms,114.385164 mW,1009105.920000 pJ
+matmul/batched_left_constant,PASS,0.062 s,0.00 MiB,0.02 MiB,9,8,SKIP,SKIP,SKIP
-matmul/batched_lhs_broadcast,PASS,0.056 s,0.00 MiB,0.01 MiB,5,4,0.005681 ms,109.389361 mW,621440.960000 pJ
+matmul/batched_lhs_broadcast,PASS,0.059 s,0.00 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-matmul/batched_rhs_broadcast,PASS,0.059 s,0.00 MiB,0.01 MiB,5,4,0.005958 ms,108.588949 mW,646972.960000 pJ
+matmul/batched_rhs_broadcast,PASS,0.087 s,0.00 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-matmul/dynamic,PASS,0.054 s,0.00 MiB,0.00 MiB,5,0,0.001621 ms,91.421962 mW,148195.000000 pJ
+matmul/dynamic,PASS,0.058 s,0.00 MiB,0.00 MiB,5,0,SKIP,SKIP,SKIP
-matmul/huge_1024,PASS,0.142 s,0.01 MiB,0.10 MiB,73,64,0.017522 ms,215.037402 mW,3767885.360000 pJ
+matmul/huge_1024,PASS,0.164 s,0.01 MiB,0.10 MiB,73,64,SKIP,SKIP,SKIP
-matmul/left_constant,PASS,0.064 s,0.00 MiB,0.01 MiB,5,4,0.005853 ms,108.861944 mW,637168.960000 pJ
+matmul/left_constant,PASS,0.058 s,0.00 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-matmul/matrix_vector,PASS,0.097 s,0.52 MiB,0.78 MiB,168,173,0.384660 ms,202.131271 mW,77751814.880000 pJ
+matmul/matrix_vector,PASS,0.100 s,0.52 MiB,0.78 MiB,168,173,SKIP,SKIP,SKIP
-matmul/vector_matrix,PASS,0.086 s,0.01 MiB,0.01 MiB,9,8,0.007409 ms,118.680243 mW,879301.920000 pJ
+matmul/vector_matrix,PASS,0.148 s,0.01 MiB,0.01 MiB,9,8,SKIP,SKIP,SKIP
-matmul/yolo_attention,PASS,0.417 s,1.02 MiB,43.44 MiB,168,0,8.151445 ms,170.003707 mW,1385775865.000000 pJ
+matmul/yolo_attention,PASS,0.466 s,1.02 MiB,43.44 MiB,168,0,SKIP,SKIP,SKIP
-mul/after_conv,PASS,0.060 s,0.00 MiB,0.00 MiB,4,3,0.005453 ms,107.639046 mW,586955.720000 pJ
+mul/after_conv,PASS,0.059 s,0.00 MiB,0.00 MiB,4,3,SKIP,SKIP,SKIP
-mul/basic,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ
+mul/basic,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-mul/channel_broadcast_1024,PASS,0.060 s,0.02 MiB,0.01 MiB,1,0,0.006913 ms,78.118038 mW,540030.000000 pJ
+mul/channel_broadcast_1024,PASS,0.051 s,0.02 MiB,0.01 MiB,1,0,SKIP,SKIP,SKIP
-mul/leading_dimension_broadcast,PASS,0.058 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ
+mul/leading_dimension_broadcast,PASS,0.071 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-mul/scalar_constant,PASS,0.053 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ
+mul/scalar_constant,PASS,0.053 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-pool/avg_basic,PASS,0.057 s,0.00 MiB,0.00 MiB,1,0,0.011939 ms,78.022112 mW,931506.000000 pJ
+pool/avg_basic,PASS,0.055 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-pool/avg_ceil_mode,PASS,0.054 s,0.00 MiB,0.00 MiB,1,0,0.004359 ms,78.033035 mW,340146.000000 pJ
+pool/avg_ceil_mode,PASS,0.059 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-pool/avg_explicit_padding,PASS,0.062 s,0.00 MiB,0.00 MiB,1,0,0.008822 ms,78.027205 mW,688356.000000 pJ
+pool/avg_explicit_padding,PASS,0.059 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-pool/avg_include_pad,PASS,0.055 s,0.00 MiB,0.00 MiB,1,0,0.008506 ms,78.016929 mW,663612.000000 pJ
+pool/avg_include_pad,PASS,0.062 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-pool/avg_large_channels,PASS,0.067 s,0.04 MiB,0.02 MiB,1,0,0.178249 ms,78.280327 mW,13953390.000000 pJ
+pool/avg_large_channels,PASS,0.078 s,0.04 MiB,0.02 MiB,1,0,SKIP,SKIP,SKIP
-pool/avg_non_uniform_stride,PASS,0.060 s,0.00 MiB,0.00 MiB,1,0,0.014513 ms,78.016537 mW,1132254.000000 pJ
+pool/avg_non_uniform_stride,PASS,0.057 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-pool/avg_real_asymmetric_padding,PASS,0.071 s,0.00 MiB,0.00 MiB,1,0,0.025206 ms,78.024756 mW,1966692.000000 pJ
+pool/avg_real_asymmetric_padding,PASS,0.061 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-pool/max_after_conv,PASS,0.066 s,0.00 MiB,0.00 MiB,6,4,0.006452 ms,96.374606 mW,621808.960000 pJ
+pool/max_after_conv,PASS,0.061 s,0.00 MiB,0.00 MiB,6,4,SKIP,SKIP,SKIP
-pool/max_basic,PASS,0.063 s,0.00 MiB,0.00 MiB,3,0,0.001634 ms,92.132191 mW,150544.000000 pJ
+pool/max_basic,PASS,0.060 s,0.00 MiB,0.00 MiB,3,0,SKIP,SKIP,SKIP
-pool/max_ceil_mode,PASS,0.078 s,0.00 MiB,0.00 MiB,2,0,0.001297 ms,79.111025 mW,102607.000000 pJ
+pool/max_ceil_mode,PASS,0.055 s,0.00 MiB,0.00 MiB,2,0,SKIP,SKIP,SKIP
-pool/max_global_style_kernel_equals_input,PASS,0.067 s,0.00 MiB,0.00 MiB,1,0,0.004366 ms,78.010994 mW,340596.000000 pJ
+pool/max_global_style_kernel_equals_input,PASS,0.100 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-pool/max_non_square_kernel,PASS,0.073 s,0.00 MiB,0.00 MiB,4,0,0.003409 ms,93.253447 mW,317901.000000 pJ
+pool/max_non_square_kernel,PASS,0.062 s,0.00 MiB,0.00 MiB,4,0,SKIP,SKIP,SKIP
-pool/max_real_asymmetric_padding,PASS,0.064 s,0.00 MiB,0.00 MiB,4,0,0.003078 ms,93.124756 mW,286638.000000 pJ
+pool/max_real_asymmetric_padding,PASS,0.062 s,0.00 MiB,0.00 MiB,4,0,SKIP,SKIP,SKIP
-pool/max_same_upper,PASS,0.063 s,0.00 MiB,0.00 MiB,3,0,0.003024 ms,92.095238 mW,278496.000000 pJ
+pool/max_same_upper,PASS,0.059 s,0.00 MiB,0.00 MiB,3,0,SKIP,SKIP,SKIP
-pool/max_stride2_multichannel,PASS,0.066 s,0.00 MiB,0.00 MiB,3,0,0.004012 ms,92.269192 mW,370184.000000 pJ
+pool/max_stride2_multichannel,PASS,0.056 s,0.00 MiB,0.00 MiB,3,0,SKIP,SKIP,SKIP
-reduce_mean/4d_spatial,PASS,0.078 s,0.00 MiB,0.00 MiB,3,0,0.000321 ms,92.448598 mW,29676.000000 pJ
+reduce_mean/4d_spatial,PASS,0.053 s,0.00 MiB,0.00 MiB,3,0,SKIP,SKIP,SKIP
-reduce_mean/4d_spatial_keepdims_0,PASS,0.066 s,0.00 MiB,0.00 MiB,4,0,0.000655 ms,94.352672 mW,61801.000000 pJ
+reduce_mean/4d_spatial_keepdims_0,PASS,0.070 s,0.00 MiB,0.00 MiB,4,0,SKIP,SKIP,SKIP
-reduce_mean/after_conv,PASS,0.068 s,0.00 MiB,0.00 MiB,5,3,0.005342 ms,106.951089 mW,571332.720000 pJ
+reduce_mean/after_conv,PASS,0.063 s,0.00 MiB,0.00 MiB,5,3,SKIP,SKIP,SKIP
-reduce_mean/all_axes_keepdims_0,PASS,0.067 s,0.00 MiB,0.00 MiB,2,0,0.000391 ms,79.237852 mW,30982.000000 pJ
+reduce_mean/all_axes_keepdims_0,PASS,0.053 s,0.00 MiB,0.00 MiB,2,0,SKIP,SKIP,SKIP
-reduce_mean/all_axes_keepdims_1,PASS,0.056 s,0.00 MiB,0.00 MiB,1,0,0.000221 ms,78.217195 mW,17286.000000 pJ
+reduce_mean/all_axes_keepdims_1,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-reduce_mean/basic,PASS,0.052 s,0.00 MiB,0.00 MiB,4,0,0.000373 ms,93.514745 mW,34881.000000 pJ
+reduce_mean/basic,PASS,0.052 s,0.00 MiB,0.00 MiB,4,0,SKIP,SKIP,SKIP
-reduce_mean/channel_axis_nchw,PASS,0.054 s,0.03 MiB,0.02 MiB,4,0,0.164926 ms,93.596631 mW,15436518.000000 pJ
+reduce_mean/channel_axis_nchw,PASS,0.053 s,0.03 MiB,0.02 MiB,4,0,SKIP,SKIP,SKIP
-reduce_mean/keepdims_0,PASS,0.054 s,0.00 MiB,0.00 MiB,5,0,0.000748 ms,91.401070 mW,68368.000000 pJ
+reduce_mean/keepdims_0,PASS,0.053 s,0.00 MiB,0.00 MiB,5,0,SKIP,SKIP,SKIP
-reduce_mean/large_dimension_1024,PASS,0.056 s,0.01 MiB,0.00 MiB,1,0,0.002785 ms,78.017235 mW,217278.000000 pJ
+reduce_mean/large_dimension_1024,PASS,0.061 s,0.01 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-reduce_mean/legacy_axes_1_2_keepdims_1,PASS,0.064 s,0.00 MiB,0.00 MiB,2,0,0.000271 ms,79.354244 mW,21505.000000 pJ
+reduce_mean/legacy_axes_1_2_keepdims_1,PASS,0.052 s,0.00 MiB,0.00 MiB,2,0,SKIP,SKIP,SKIP
-reduce_mean/legacy_axis1_keepdims_0,PASS,0.071 s,0.00 MiB,0.00 MiB,9,0,0.001986 ms,92.501511 mW,183708.000000 pJ
+reduce_mean/legacy_axis1_keepdims_0,PASS,0.052 s,0.00 MiB,0.00 MiB,9,0,SKIP,SKIP,SKIP
-reduce_mean/legacy_axis1_keepdims_1,PASS,0.065 s,0.00 MiB,0.00 MiB,8,0,0.001373 ms,94.559359 mW,129830.000000 pJ
+reduce_mean/legacy_axis1_keepdims_1,PASS,0.058 s,0.00 MiB,0.00 MiB,8,0,SKIP,SKIP,SKIP
-reduce_mean/legacy_empty_axes_noop,PASS,0.053 s,0.00 MiB,0.00 MiB,1,0,0.000221 ms,78.217195 mW,17286.000000 pJ
+reduce_mean/legacy_empty_axes_noop,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-reduce_mean/legacy_nchw_spatial,PASS,0.050 s,0.00 MiB,0.00 MiB,3,0,0.000321 ms,92.448598 mW,29676.000000 pJ
+reduce_mean/legacy_nchw_spatial,PASS,0.058 s,0.00 MiB,0.00 MiB,3,0,SKIP,SKIP,SKIP
-reduce_mean/legacy_negative_axis,PASS,0.064 s,0.00 MiB,0.00 MiB,6,0,0.000553 ms,93.520796 mW,51717.000000 pJ
+reduce_mean/legacy_negative_axis,PASS,0.059 s,0.00 MiB,0.00 MiB,6,0,SKIP,SKIP,SKIP
-reduce_mean/legacy_reduce_all_keepdims_1,PASS,0.054 s,0.00 MiB,0.00 MiB,1,0,0.000221 ms,78.217195 mW,17286.000000 pJ
+reduce_mean/legacy_reduce_all_keepdims_1,PASS,0.068 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-reduce_mean/negative_axis,PASS,0.059 s,0.00 MiB,0.00 MiB,6,0,0.000553 ms,93.520796 mW,51717.000000 pJ
+reduce_mean/negative_axis,PASS,0.065 s,0.00 MiB,0.00 MiB,6,0,SKIP,SKIP,SKIP
-relu/4d,PASS,0.068 s,0.00 MiB,0.00 MiB,1,0,0.000521 ms,78.184261 mW,40734.000000 pJ
+relu/4d,PASS,0.066 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-relu/after_conv,PASS,0.057 s,0.00 MiB,0.00 MiB,3,3,0.004956 ms,106.998935 mW,530286.720000 pJ
+relu/after_conv,PASS,0.075 s,0.00 MiB,0.00 MiB,3,3,SKIP,SKIP,SKIP
-relu/after_gemm,PASS,0.064 s,0.01 MiB,0.01 MiB,5,4,0.007513 ms,105.158653 mW,790056.960000 pJ
+relu/after_gemm,PASS,0.079 s,0.01 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-relu/basic,PASS,0.056 s,0.00 MiB,0.00 MiB,1,0,0.000221 ms,78.217195 mW,17286.000000 pJ
+relu/basic,PASS,0.049 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-reshape/4d_to_2d_flatten,PASS,0.068 s,0.00 MiB,0.00 MiB,1,0,0.000258 ms,78.279070 mW,20196.000000 pJ
+reshape/4d_to_2d_flatten,PASS,0.049 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-reshape/infer_dim_minus_one,PASS,0.059 s,0.00 MiB,0.00 MiB,1,0,0.000162 ms,78.296296 mW,12684.000000 pJ
+reshape/infer_dim_minus_one,PASS,0.061 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-reshape/same_rank,PASS,0.065 s,0.00 MiB,0.00 MiB,1,0,0.000162 ms,78.296296 mW,12684.000000 pJ
+reshape/same_rank,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-reshape/zero_copies_input_dim,PASS,0.057 s,0.00 MiB,0.00 MiB,1,0,0.000162 ms,78.296296 mW,12684.000000 pJ
+reshape/zero_copies_input_dim,PASS,0.052 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-resize/height_only,PASS,0.063 s,0.00 MiB,0.00 MiB,4,0,0.000693 ms,93.554113 mW,64833.000000 pJ
+resize/height_only,PASS,0.079 s,0.00 MiB,0.00 MiB,4,0,SKIP,SKIP,SKIP
-resize/nearest_2x,PASS,0.064 s,0.00 MiB,0.00 MiB,4,0,0.001173 ms,93.572890 mW,109761.000000 pJ
+resize/nearest_2x,PASS,0.053 s,0.00 MiB,0.00 MiB,4,0,SKIP,SKIP,SKIP
-resize/nearest_downsample,PASS,0.062 s,0.00 MiB,0.00 MiB,2,0,0.000427 ms,79.449649 mW,33925.000000 pJ
+resize/nearest_downsample,PASS,0.055 s,0.00 MiB,0.00 MiB,2,0,SKIP,SKIP,SKIP
-resize/non_uniform,PASS,0.068 s,0.00 MiB,0.00 MiB,6,0,0.001753 ms,93.575014 mW,164037.000000 pJ
+resize/non_uniform,PASS,0.053 s,0.00 MiB,0.00 MiB,6,0,SKIP,SKIP,SKIP
-resize/width_only,PASS,0.070 s,0.00 MiB,0.00 MiB,2,0,0.000667 ms,79.503748 mW,53029.000000 pJ
+resize/width_only,PASS,0.052 s,0.00 MiB,0.00 MiB,2,0,SKIP,SKIP,SKIP
-resize/with_sizes,PASS,0.052 s,0.00 MiB,0.00 MiB,3,0,0.000797 ms,92.542033 mW,73756.000000 pJ
+resize/with_sizes,PASS,0.053 s,0.00 MiB,0.00 MiB,3,0,SKIP,SKIP,SKIP
-sigmoid/4d,PASS,0.067 s,0.00 MiB,0.00 MiB,1,0,0.000521 ms,78.184261 mW,40734.000000 pJ
+sigmoid/4d,PASS,0.077 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-sigmoid/after_gemm,PASS,0.070 s,0.01 MiB,0.01 MiB,5,4,0.007513 ms,105.158653 mW,790056.960000 pJ
+sigmoid/after_gemm,PASS,0.058 s,0.01 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-sigmoid/basic,PASS,0.062 s,0.00 MiB,0.00 MiB,1,0,0.000221 ms,78.217195 mW,17286.000000 pJ
+sigmoid/basic,PASS,0.052 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-slice/2d_basic,PASS,0.049 s,0.00 MiB,0.00 MiB,1,0,0.000242 ms,78.297521 mW,18948.000000 pJ
+slice/2d_basic,PASS,0.055 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-slice/after_conv,PASS,0.062 s,0.00 MiB,0.01 MiB,7,6,0.011296 ms,118.190765 mW,1335082.880000 pJ
+slice/after_conv,PASS,0.061 s,0.00 MiB,0.01 MiB,7,6,SKIP,SKIP,SKIP
-slice/default_axes,PASS,0.056 s,0.00 MiB,0.00 MiB,1,0,0.000242 ms,78.297521 mW,18948.000000 pJ
+slice/default_axes,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-slice/large_channel_1024,PASS,0.049 s,0.01 MiB,0.00 MiB,1,0,0.002832 ms,78.144068 mW,221304.000000 pJ
+slice/large_channel_1024,PASS,0.068 s,0.01 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-slice/nchw_spatial_crop,PASS,0.063 s,0.00 MiB,0.00 MiB,1,0,0.001302 ms,78.239631 mW,101868.000000 pJ
+slice/nchw_spatial_crop,PASS,0.052 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-slice/negative_axis,PASS,0.058 s,0.00 MiB,0.00 MiB,1,0,0.000562 ms,78.298932 mW,44004.000000 pJ
+slice/negative_axis,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-slice/negative_indices,PASS,0.059 s,0.00 MiB,0.00 MiB,1,0,0.000322 ms,78.298137 mW,25212.000000 pJ
+slice/negative_indices,PASS,0.048 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-slice/step2,PASS,0.054 s,0.00 MiB,0.00 MiB,1,0,0.002042 ms,78.293830 mW,159876.000000 pJ
+slice/step2,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-softmax/3d_last_axis,PASS,0.059 s,0.00 MiB,0.00 MiB,1,0,UNSUPPORTED,UNSUPPORTED,UNSUPPORTED
+softmax/3d_last_axis,PASS,0.056 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-softmax/basic,PASS,0.069 s,0.00 MiB,0.00 MiB,1,0,UNSUPPORTED,UNSUPPORTED,UNSUPPORTED
+softmax/basic,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-softmax/channel_axis,PASS,0.061 s,0.00 MiB,0.00 MiB,1,0,UNSUPPORTED,UNSUPPORTED,UNSUPPORTED
+softmax/channel_axis,PASS,0.056 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-softmax/large_dimension_1024,PASS,0.045 s,0.01 MiB,0.01 MiB,1,0,UNSUPPORTED,UNSUPPORTED,UNSUPPORTED
+softmax/large_dimension_1024,PASS,0.048 s,0.01 MiB,0.01 MiB,1,0,SKIP,SKIP,SKIP
-softmax/negative_axis,PASS,0.064 s,0.00 MiB,0.00 MiB,1,0,UNSUPPORTED,UNSUPPORTED,UNSUPPORTED
+softmax/negative_axis,PASS,0.053 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-split/basic,PASS,0.066 s,0.00 MiB,0.00 MiB,1,0,0.000403 ms,78.297767 mW,31554.000000 pJ
+split/basic,PASS,0.050 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-split/equal_three_way,PASS,0.065 s,0.00 MiB,0.00 MiB,1,0,0.000564 ms,78.297872 mW,44160.000000 pJ
+split/equal_three_way,PASS,0.049 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-split/negative_axis,PASS,0.058 s,0.00 MiB,0.00 MiB,1,0,0.001083 ms,78.288089 mW,84786.000000 pJ
+split/negative_axis,PASS,0.067 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-split/uneven_channel_axis_4d,PASS,0.062 s,0.00 MiB,0.00 MiB,1,0,0.000242 ms,78.297521 mW,18948.000000 pJ
+split/uneven_channel_axis_4d,PASS,0.053 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-sub/after_gemm,PASS,0.068 s,0.01 MiB,0.01 MiB,5,4,0.007784 ms,104.703618 mW,815012.960000 pJ
+sub/after_gemm,PASS,0.059 s,0.01 MiB,0.01 MiB,5,4,SKIP,SKIP,SKIP
-sub/basic,PASS,0.060 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ
+sub/basic,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-sub/broadcast_row,PASS,0.062 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ
+sub/broadcast_row,PASS,0.051 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-sub/channel_broadcast_1024,PASS,0.057 s,0.02 MiB,0.01 MiB,1,0,0.006913 ms,78.118038 mW,540030.000000 pJ
+sub/channel_broadcast_1024,PASS,0.056 s,0.02 MiB,0.01 MiB,1,0,SKIP,SKIP,SKIP
-sub/constant_lhs_broadcast,PASS,0.057 s,0.00 MiB,0.00 MiB,1,0,0.000322 ms,78.223602 mW,25188.000000 pJ
+sub/constant_lhs_broadcast,PASS,0.052 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP
-sub/leading_dimension_broadcast,PASS,0.058 s,0.00 MiB,0.00 MiB,1,0,0.000323 ms,78.222910 mW,25266.000000 pJ
+sub/leading_dimension_broadcast,PASS,0.066 s,0.00 MiB,0.00 MiB,1,0,SKIP,SKIP,SKIP