big refactor

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

@@ -1,49 +1,133 @@
 #include "ConvGeometry.hpp"
 
 #include <algorithm>
+#include <limits>
 
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 
 namespace onnx_mlir {
 
+namespace {
+
+static const ONNXToSpatialPlanningOptions& defaultPlanningOptions() {
+  static const ONNXToSpatialPlanningOptions options {
+    std::numeric_limits<uint64_t>::max(),
+    std::numeric_limits<uint64_t>::max(),
+    spatial::ConvLoweringStrategy::Auto,
+    false,
+  };
+  return options;
+}
+
+} // namespace
+
+const ONNXToSpatialPlanningOptions& ConvLoweringState::planningOptions() const {
+  return options ? *options : defaultPlanningOptions();
+}
+
 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::SpatialTargetResources& target) {
   ConvGeometry geo {
-    state.batchSize,
-    state.numChannelsIn,
-    state.xHeight,
-    state.xWidth,
-    state.numChannelsOut,
-    state.wHeight,
-    state.wWidth,
-    state.outHeight,
-    state.outWidth,
-    state.group,
-    state.numChannelsInPerGroup,
-    state.numChannelsOutPerGroup,
-    state.numChannelsInPerGroup * state.wHeight * state.wWidth,
-    state.numChannelsOutPerGroup,
-    state.batchSize * state.outHeight * state.outWidth,
-    static_cast<int64_t>(crossbarSize.getValue()),
+    problem.numChannelsInPerGroup * problem.wHeight * problem.wWidth,
+    problem.numChannelsOutPerGroup,
+    problem.batchSize * problem.outHeight * problem.outWidth,
+    static_cast<int64_t>(target.matrixShape.rows),
+    static_cast<int64_t>(target.matrixUnitsPerProcessor),
     1,
     0,
-    state.hasBias,
-    isDepthwiseConv(state.group, state.numChannelsIn, state.numChannelsOut, state.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(
+    spatial::ConvLoweringStrategy strategy) {
+  switch (strategy) {
+  case spatial::ConvLoweringStrategy::Depthwise:
+    return ConvMaterializationKind::StructuredDepthwise;
+  case spatial::ConvLoweringStrategy::Legacy:
+  case spatial::ConvLoweringStrategy::PackedIm2Col:
+    return ConvMaterializationKind::PackedIm2Col;
+  case spatial::ConvLoweringStrategy::StreamedPatch:
+  case spatial::ConvLoweringStrategy::OutputChannelTiled:
+  case spatial::ConvLoweringStrategy::Tiled2D:
+    return ConvMaterializationKind::StreamedPatch;
+  case spatial::ConvLoweringStrategy::StreamedPacked:
+    return ConvMaterializationKind::StreamedPacked;
+  case spatial::ConvLoweringStrategy::InputKTiled:
+    return ConvMaterializationKind::InputKTiled;
+  case spatial::ConvLoweringStrategy::Auto:
+    break;
+  }
+  llvm_unreachable("auto is not a Conv materialization kind");
+}
+
+static bool fitsSingleCrossbar(const ConvGeometry& geo) {
+  return geo.k <= geo.xbarSize && geo.c <= geo.xbarSize;
+}
+
+static bool fitsPackedIm2Col(const ConvGeometry& geo,
+                             const ONNXToSpatialPlanningOptions& options) {
+  return fitsSingleCrossbar(geo) && geo.pack >= 2
+    && geo.im2colElements <= options.convIm2colMaxElements;
+}
+
+mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
+                                       spatial::ConvLoweringStrategy strategy,
+                                       const spatial::SpatialTargetResources& target,
+                                       const ONNXToSpatialPlanningOptions& options) {
+  ConvGeometry geo = buildConvGeometry(problem, target);
+  auto plan = [&]() { return ConvPlan {getMaterializationKind(strategy)}; };
+  auto ifApplicable = [&](bool applicable) -> mlir::FailureOr<ConvPlan> {
+    return applicable ? mlir::FailureOr<ConvPlan>(plan()) : mlir::FailureOr<ConvPlan>(mlir::failure());
+  };
+  switch (strategy) {
+  case spatial::ConvLoweringStrategy::Auto:
+    return mlir::failure();
+  case spatial::ConvLoweringStrategy::Legacy:
+    return plan();
+  case spatial::ConvLoweringStrategy::Depthwise:
+    return ifApplicable(problem.isDepthwise);
+  case spatial::ConvLoweringStrategy::PackedIm2Col:
+    return ifApplicable(fitsPackedIm2Col(geo, options));
+  case spatial::ConvLoweringStrategy::StreamedPatch:
+    return ifApplicable(fitsSingleCrossbar(geo));
+  case spatial::ConvLoweringStrategy::StreamedPacked:
+    return ifApplicable(fitsSingleCrossbar(geo) && geo.pack >= 2);
+  case spatial::ConvLoweringStrategy::OutputChannelTiled:
+    return ifApplicable(geo.k <= geo.xbarSize && geo.c > geo.xbarSize);
+  case spatial::ConvLoweringStrategy::InputKTiled:
+    return ifApplicable(geo.k > geo.xbarSize && geo.c <= geo.xbarSize);
+  case spatial::ConvLoweringStrategy::Tiled2D:
+    return ifApplicable(geo.k > geo.xbarSize && geo.c > geo.xbarSize);
+  }
+  llvm_unreachable("unknown Conv lowering strategy");
+}
+
+uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
+                                    int64_t packFactor,
+                                    const ONNXToSpatialPlanningOptions& options) {
   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, options.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, options.convStreamChunkPositions));
 
   if (packFactor > 1 && chunkPositions > static_cast<uint64_t>(packFactor)) {
     chunkPositions -= chunkPositions % static_cast<uint64_t>(packFactor);
@@ -52,24 +136,26 @@ uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor)
   return std::max<uint64_t>(1, chunkPositions);
 }
 
-RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvLoweringState& state) {
-  const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin;
+RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvProblem& problem) {
+  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;
-  return {std::max<int64_t>(0, rawBegin), std::min<int64_t>(state.xHeight, rawEnd)};
+    (outputRows.end - 1) * problem.strideHeight - problem.padHeightBegin
+    + 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/ONNXToSpatialOptions.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
 #include <cstdint>
 
+namespace mlir {
+class Operation;
+} // namespace mlir
+
 namespace onnx_mlir {
 
-struct ConvLoweringState {
-  mlir::Value x;
-  mlir::Value w;
-  mlir::Value b;
+struct ConvProblem {
   mlir::RankedTensorType xType;
   mlir::RankedTensorType wType;
   mlir::RankedTensorType outType;
@@ -35,29 +40,32 @@ struct ConvLoweringState {
   int64_t dilationHeight;
   int64_t dilationWidth;
   bool hasBias;
+  bool isDepthwise = false;
+  bool isGrouped = false;
+  bool isPointwise = false;
+};
+
+struct ConvLoweringState {
+  ConvProblem problem;
+  mlir::Operation* diagnosticAnchor = nullptr;
+  mlir::Value x;
+  mlir::Value w;
+  mlir::Value b;
+  const spatial::SpatialTargetResources* target = nullptr;
+  const ONNXToSpatialPlanningOptions* options = nullptr;
+
+  const spatial::SpatialTargetResources& targetInfo() const { return *target; }
+  const ONNXToSpatialPlanningOptions& planningOptions() const;
 };
 
 struct ConvGeometry {
-  int64_t batchSize;
-  int64_t numChannelsIn;
-  int64_t xHeight;
-  int64_t xWidth;
-  int64_t numChannelsOut;
-  int64_t wHeight;
-  int64_t wWidth;
-  int64_t outHeight;
-  int64_t outWidth;
-  int64_t group;
-  int64_t numChannelsInPerGroup;
-  int64_t numChannelsOutPerGroup;
   int64_t k;
   int64_t c;
   int64_t p;
   int64_t xbarSize;
+  int64_t matrixUnitsPerProcessor;
   int64_t pack;
   uint64_t im2colElements;
-  bool hasBias;
-  bool isDepthwise;
 };
 
 struct RowInterval {
@@ -73,14 +81,36 @@ struct ConvRowDemand {
   int64_t bottomHaloRows = 0;
 };
 
+enum class ConvMaterializationKind : uint8_t {
+  StructuredDepthwise,
+  PackedIm2Col,
+  StreamedPatch,
+  StreamedPacked,
+  InputKTiled,
+};
+
+struct ConvPlan {
+  ConvMaterializationKind kind = ConvMaterializationKind::PackedIm2Col;
+};
+
 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::SpatialTargetResources& target);
 
-RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvLoweringState& state);
+mlir::FailureOr<ConvPlan> makeConvPlan(const ConvProblem& problem,
+                                       spatial::ConvLoweringStrategy strategy,
+                                       const spatial::SpatialTargetResources& target,
+                                       const ONNXToSpatialPlanningOptions& options);
 
-ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state);
+uint64_t chooseStreamChunkPositions(const ConvGeometry& geo,
+                                    int64_t packFactor,
+                                    const ONNXToSpatialPlanningOptions& options);
+
+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();
@@ -48,6 +48,56 @@ static DenseElementsAttr getDenseConstantAttr(Value value) {
   return nullptr;
 }
 
+struct BlueprintSplatMulToSpatial : OpConversionPattern<ONNXMulOp> {
+  explicit BlueprintSplatMulToSpatial(MLIRContext* ctx) : OpConversionPattern(ctx, 2) {}
+
+  LogicalResult
+  matchAndRewrite(ONNXMulOp op, ONNXMulOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override {
+    auto blueprint = op.getA().getDefiningOp<spatial::SpatBlueprintOp>();
+    Value scalar = adaptor.getB();
+    if (!blueprint) {
+      blueprint = op.getB().getDefiningOp<spatial::SpatBlueprintOp>();
+      scalar = adaptor.getA();
+    }
+    auto scalarAttr = getDenseConstantAttr(scalar);
+    auto resultType = dyn_cast<RankedTensorType>(op.getResult().getType());
+    auto storageType = blueprint ? dyn_cast<RankedTensorType>(blueprint.getInput().getType()) : RankedTensorType();
+    if (!blueprint || !blueprint.getFragments().empty() || !scalarAttr || !scalarAttr.isSplat() || !resultType
+        || resultType != blueprint.getOutput().getType() || !storageType)
+      return failure();
+
+    auto mapped = mapGraphBatchFragments(
+      blueprint.getInput(), storageType, rewriter, op.getLoc(), [&](Value fragment, RankedTensorType fragmentType) {
+        auto splat = DenseElementsAttr::get(fragmentType, scalarAttr.getSplatValue<Attribute>());
+        Value constant = arith::ConstantOp::create(rewriter, op.getLoc(), fragmentType, splat);
+        return FailureOr<Value>(
+          spatial::SpatVMulOp::create(rewriter, op.getLoc(), fragmentType, fragment, constant).getResult());
+      });
+    if (failed(mapped))
+      return failure();
+
+    auto result = spatial::SpatBlueprintOp::create(rewriter,
+                                                   op.getLoc(),
+                                                   resultType,
+                                                   *mapped,
+                                                   ValueRange {},
+                                                   blueprint.getLogicalLayoutAttr(),
+                                                   blueprint.getPhysicalLayoutAttr(),
+                                                   blueprint.getFragmentOffsetsAttr(),
+                                                   blueprint.getFragmentSizesAttr(),
+                                                   blueprint.getIndexMapAttr(),
+                                                   blueprint.getModeAttr(),
+                                                   blueprint.getFragmentOperandIndicesAttr(),
+                                                   blueprint.getFragmentSourceSlotsAttr(),
+                                                   blueprint.getFragmentSourceOffsetsAttr(),
+                                                   blueprint.getFragmentStridesAttr(),
+                                                   blueprint.getConflictPolicyAttr(),
+                                                   blueprint.getCoveragePolicyAttr());
+    rewriter.replaceOp(op, result.getOutput());
+    return success();
+  }
+};
+
 static FailureOr<Value> materializeBroadcastedConstantTensor(Value value,
                                                              RankedTensorType resultType,
                                                              ConversionPatternRewriter& rewriter,
@@ -210,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();
     }
@@ -246,6 +298,7 @@ void populateElementwiseFusionPatterns(RewritePatternSet& patterns, MLIRContext*
 }
 
 void populateElementwisePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<BlueprintSplatMulToSpatial>(ctx);
   patterns.add<AddToSpatialCompute>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXSubOp, spatial::SpatVSubOp>>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXMulOp, spatial::SpatVMulOp>>(ctx);

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,9 @@
 #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/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 +35,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 +61,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 +74,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 +83,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 +93,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 +131,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 +197,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 +213,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 +235,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 +249,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 hOffset = createGemmBatchHOffset(args.lane, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
+      Value kOffset = createGemmBatchKOffset(args.lane, numOutRows, numKSlices, xbarSize, 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());
-      Value aTile = extractATile(args.inputs.front(), row, kOffset, aTileType, rewriter, loc);
+        RankedTensorType::get({1, xbarSize}, partialPiecesType.getElementType());
+      Value aTile = extractATile(args.inputs.front(), row, kOffset, aTileType, xbarSize, rewriter, loc);
 
       SmallVector<OpFoldResult> bOffsets {kOffset, hOffset};
-      SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(crossbarSize.getValue()),
-                                        rewriter.getIndexAttr(crossbarSize.getValue())};
+      SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarSize), rewriter.getIndexAttr(xbarSize)};
       SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, 2);
       Value bTile = extractStaticSliceOrIdentity(
         rewriter, loc, args.weights.front(), bTileType, bOffsets, bSizes, unitStrides);
@@ -260,7 +277,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 +297,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)};
@@ -311,13 +328,15 @@ static FailureOr<RankedTensorType> verifyDynamicGemmBiasType(RankedTensorType cT
 }
 
 static bool hasGemmBias(Value c) {
+  if (!c)
+    return false;
   Operation* definingOp = c.getDefiningOp();
   return (!definingOp || !isa<ONNXNoneOp>(definingOp)) && !isZeroSplatHostConstant(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 +349,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 +384,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 +444,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 +529,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 +546,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 +566,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 +597,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());
-  auto pieceType = RankedTensorType::get({numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
+  const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(1), xbarSize);
+  auto pieceType = RankedTensorType::get({numOutRows, xbarSize},
                                          partialPiecesType.getElementType());
 
   if (bias && cast<RankedTensorType>(bias.getType()) != paddedOutType)
@@ -590,20 +614,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 +661,96 @@ static FailureOr<Value> createReductionOutput(Value partialPieces,
 }
 
 struct GemmToSpatialComputes : OpConversionPattern<ONNXGemmOp> {
-  using OpConversionPattern::OpConversionPattern;
+  explicit GemmToSpatialComputes(MLIRContext* ctx, const spatial::SpatialTargetResources& target)
+      : OpConversionPattern<ONNXGemmOp>(ctx), target(target) {}
 
   LogicalResult matchAndRewrite(ONNXGemmOp gemmOp,
                                 ONNXGemmOpAdaptor gemmOpAdaptor,
                                 ConversionPatternRewriter& rewriter) const override;
+
+  const spatial::SpatialTargetResources& target;
 };
 
 } // namespace
 
-LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
-                                                     ONNXGemmOpAdaptor gemmOpAdaptor,
-                                                     ConversionPatternRewriter& rewriter) const {
-  Location loc = gemmOp.getLoc();
-  Value a = gemmOpAdaptor.getA();
-  Value b = gemmOpAdaptor.getB();
-  Value c = gemmOpAdaptor.getC();
-
+FailureOr<Value> lowerGemmToSpatial(
+    Operation* diagnosticAnchor,
+    Value a,
+    Value b,
+    Value c,
+    RankedTensorType outType,
+    bool transA,
+    bool transB,
+    float alpha,
+    float beta,
+    const spatial::SpatialTargetResources& 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 (!aType || !bType || !outType)
+  if (!diagnosticAnchor || !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);
-  const int64_t numOutCols = outType.getDimSize(1);
-  const int64_t reductionSize = aType.getDimSize(1);
-  const bool transposeB = gemmOpAdaptor.getTransB();
+  ContractionProblem problem;
+  problem.lhsBatchShape = {};
+  problem.rhsBatchShape = {};
+  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.lhsElementType = aType.getElementType();
+  problem.rhsElementType = bType.getElementType();
+  problem.resultElementType = outType.getElementType();
+  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 +758,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) {
-      gemmOp.emitOpError("has inconsistent A, B, and output shapes");
+    if (aType.getDimSize(0) != problem.m || bReductionSize != problem.k || bOutputColumns != problem.n) {
+      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 scalarPiecesType = spatial::getGraphBatchPhysicalResultType(numOutCols, columnType);
+    auto columnType = RankedTensorType::get({problem.m, 1}, outType.getElementType());
+    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 +778,128 @@ 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 success();
+    return outputCompute->getResult(0);
   }
 
   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 denseAttr = getHostConstDenseElementsAttr(b);
+      auto inputType = denseAttr ? dyn_cast<RankedTensorType>(denseAttr.getType()) : nullptr;
+      auto transposedAttr = inputType && inputType.hasStaticShape() && transposedType.hasStaticShape()
+        ? transposeDenseElementsAttr(denseAttr, {1, 0})
+        : FailureOr<DenseElementsAttr>(failure());
+      if (failed(transposedAttr) || transposedAttr->getType() != transposedType) {
+        diagnosticAnchor->emitOpError("requires Gemm input B transpose to remain statically materializable");
+        return failure();
+      }
+      b = getOrCreateConstant(rewriter,
+                              rewriter.getInsertionBlock()->getParentOp(),
+                              *transposedAttr,
+                              transposedType);
+    } 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) {
-    gemmOp.emitOpError("has inconsistent A, B, and output shapes after transpose handling");
+  if (aType.getDimSize(0) != problem.m || bType.getDimSize(0) != problem.k || bType.getDimSize(1) != problem.n) {
+    diagnosticAnchor->emitOpError("has inconsistent A, B, and output shapes after transpose handling");
     return failure();
   }
 
-  const int64_t numKSlices = ceilIntegerDivide(reductionSize, crossbarSize.getValue());
-  const int64_t numOutHSlices = ceilIntegerDivide(numOutCols, crossbarSize.getValue());
-  const int64_t paddedReductionSize = numKSlices * static_cast<int64_t>(crossbarSize.getValue());
-  const int64_t paddedOutCols = numOutHSlices * static_cast<int64_t>(crossbarSize.getValue());
+  ContractionPlan plan = makeContractionPlan(
+      problem, target, ContractionPlanKind::StaticTiled);
+  const int64_t xbarSize = plan.tileK;
+  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);
   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) {
-  patterns.insert<GemmToSpatialComputes>(ctx);
+void populateGemmPatterns(RewritePatternSet& patterns,
+                          MLIRContext* ctx,
+                          const spatial::SpatialTargetResources& 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 SpatialTargetResources;
+}
+
+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::SpatialTargetResources& target,
+    mlir::PatternRewriter& rewriter,
+    mlir::Location loc);
+
+} // 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"),
-           rewriter.getStringAttr("fragmented"),
+           spatial::getNCHWLayout(rewriter.getContext()),
+           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),