#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.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/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"

#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/raw_ostream.h"

#include <algorithm>
#include <cctype>
#include <map>
#include <mutex>
#include <optional>
#include <string>
#include <vector>

#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
#include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.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"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"

using namespace mlir;

namespace onnx_mlir {
namespace {

struct ConvToGemm : OpConversionPattern<ONNXConvOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult matchAndRewrite(ONNXConvOp convOp,
                                ONNXConvOpAdaptor convOpAdaptor,
                                ConversionPatternRewriter& rewriter) const override;
};

struct ConvLoweringDecision {
  PimConvLoweringType strategy;
  std::string reason;
  bool isAuto = false;
  std::string fallbackReason;
  std::string rejectedAutoStrategy;
};

struct PreparedConvInput {
  Value value;
  RankedTensorType type;
};

struct ConvStrategyEstimate {
  uint64_t estimatedMvmCount = 0;
  uint64_t estimatedReductionVAddCount = 0;
  uint64_t estimatedOutputFragments = 1;
  bool perOutputPositionReduction = false;
  bool requiresFuncReturnMaterialization = false;
  bool giantCollectorConcatExpected = false;
  bool fullInputBroadcastExpected = false;
  uint64_t concatOperandCount = 0;
  std::string materializationKind = "none";
  std::string collectorCore = "-";
};

struct ConvReportEntry {
  uint64_t id;
  std::string where;
  std::string strategy;
  std::string mode;
  std::string inputShape;
  std::string weightShape;
  std::string outputShape;
  int64_t groups;
  int64_t k;
  int64_t c;
  int64_t p;
  int64_t xbarSize;
  int64_t pack;
  uint64_t im2colElements;
  uint64_t im2colBudget;
  std::string chunkText;
  int64_t batchSize;
  int64_t numberOfBatches;
  std::string spatialComputeBatch;
  std::string batchedInstructionEmission;
  std::string reason;
  std::string fallbackReason;
  std::string rejectedAutoStrategy;
  uint64_t estimatedMvmCount;
  uint64_t estimatedReductionVAddCount;
  uint64_t estimatedOutputFragments;
  std::string materializationRequiredAtReturn;
  std::string materializationKind;
  uint64_t concatOperandCount;
  std::string collectorCore;
  std::string giantCollectorConcatExpected;
  std::string fullInputBroadcastExpected;
};

enum class DistributedTensorOpKind {
  Relu,
  Sigmoid,
  Add,
  Sub,
  Mul,
  Div,
  Conv,
};

enum class DistributedConvBarrierKind {
  Return,
  UnsupportedConsumer,
  Fanout,
  DeadValue,
  GroupedConv,
  Depthwise,
};

enum class DistributedTensorConstantKind {
  None,
  Splat,
  PerChannel,
};

enum class DistributedTensorLayoutKind {
  NchwRowStrip,
};

struct DistributedFragmentInfo {
  SmallVector<int64_t, 4> offsets;
  SmallVector<int64_t, 4> sizes;
  SmallVector<int64_t, 4> strides;
  int64_t producerLane = 0;
};

struct DistributedTensorInfo {
  Value storage;
  RankedTensorType logicalType;
  DistributedTensorLayoutKind layoutKind = DistributedTensorLayoutKind::NchwRowStrip;
  SmallVector<DistributedFragmentInfo, 8> fragments;
  int64_t laneCount = 0;
  int64_t fragmentHeight = 1;
  int64_t channels = 0;
  int64_t height = 0;
  int64_t width = 0;

  bool isRowStripNchw() const { return layoutKind == DistributedTensorLayoutKind::NchwRowStrip; }
};

struct DistributedTensorRegistry {
  llvm::DenseMap<Value, DistributedTensorInfo> infos;

  void bind(Value value, const DistributedTensorInfo& info) { infos[value] = info; }

  const DistributedTensorInfo* lookup(Value value) const {
    auto it = infos.find(value);
    if (it == infos.end())
      return nullptr;
    return &it->second;
  }
};

struct DistributedTensorStep {
  Operation* op = nullptr;
  DistributedTensorOpKind kind;
  DenseElementsAttr constantAttr;
  DistributedTensorConstantKind constantKind = DistributedTensorConstantKind::None;
  bool fragmentOnLhs = true;
  std::optional<ConvLoweringState> convState;
};

struct DistributedConvAnalysis {
  SmallVector<DistributedTensorStep, 4> steps;
  Operation* replacementOp = nullptr;
  DistributedConvBarrierKind barrierKind = DistributedConvBarrierKind::UnsupportedConsumer;
  std::string barrierDetail;

  bool hasLocalConsumers() const { return !steps.empty(); }
  bool hasDistributedConvConsumer() const {
    return llvm::any_of(steps, [](const DistributedTensorStep& step) { return step.kind == DistributedTensorOpKind::Conv; });
  }
};

struct DistributedChainReportEntry {
  uint64_t chainId = 0;
  uint64_t chainLength = 0;
  std::string producerKind;
  std::string distributedOps;
  std::string materializationPoints;
  std::string firstMaterializationReason;
  uint64_t maxLiveFragments = 0;
  uint64_t maxFragmentFanout = 0;
  uint64_t patchBuilderCoreCount = 0;
  std::string centralJunctionDetected = "no";
  std::string convInputMaterializationKind = "dense_materialization_fallback";
  uint64_t localPatchFragments = 0;
  uint64_t remotePatchFragments = 0;
  uint64_t haloTransferCount = 0;
  uint64_t groupedTransferCount = 0;
  std::string fallbackReason;
};

struct DistributedConvReportTotals {
  uint64_t totalConvs = 0;
  uint64_t distributedTensorsCreated = 0;
  uint64_t distributedValuesPropagated = 0;
  uint64_t distributedConsumersHandled = 0;
  uint64_t distributedConvInputsSeen = 0;
  uint64_t distributedConvInputsConsumed = 0;
  uint64_t materializationBarriersInserted = 0;
  std::map<std::string, uint64_t> fallbackReasons;
  std::map<std::string, uint64_t> barrierReasons;
  SmallVector<DistributedChainReportEntry, 16> chains;
};

static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter& rewriter);
static FailureOr<Value> createRowStripPackedRows(Value rows,
                                                 const ConvLoweringState& state,
                                                 PatternRewriter& rewriter,
                                                 Location loc);

static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp, Value x, Value w, Value b);

static StringRef stringifyDistributedConvBarrierKind(DistributedConvBarrierKind kind) {
  switch (kind) {
  case DistributedConvBarrierKind::Return: return "func.return";
  case DistributedConvBarrierKind::UnsupportedConsumer: return "unsupported_consumer";
  case DistributedConvBarrierKind::Fanout: return "fanout";
  case DistributedConvBarrierKind::DeadValue: return "dead_value";
  case DistributedConvBarrierKind::GroupedConv: return "grouped_conv";
  case DistributedConvBarrierKind::Depthwise: return "depthwise_conv";
  }
  llvm_unreachable("unknown distributed conv barrier kind");
}

static StringRef stringifyConvLoweringStrategy(PimConvLoweringType strategy) {
  switch (strategy) {
  case PimConvLoweringAuto: return "auto";
  case PimConvLoweringLegacy: return "legacy";
  case PimConvLoweringDepthwise: return "depthwise";
  case PimConvLoweringPackedIm2Col: return "packed-im2col";
  case PimConvLoweringStreamedPatch: return "streamed-patch";
  case PimConvLoweringStreamedPacked: return "streamed-packed";
  case PimConvLoweringOutputChannelTiled: return "output-channel-tiled";
  case PimConvLoweringInputKTiled: return "input-k-tiled";
  case PimConvLoweringTiled2D: return "tiled-2d";
  }
  llvm_unreachable("unknown conv lowering strategy");
}

static bool requiresFuncReturnMaterialization(const DistributedConvAnalysis& analysis) {
  return !analysis.hasLocalConsumers() && analysis.barrierKind == DistributedConvBarrierKind::Return;
}

static ConvStrategyEstimate estimateConvStrategy(const ConvGeometry& geo,
                                                 PimConvLoweringType strategy,
                                                 const DistributedConvAnalysis& analysis) {
  ConvStrategyEstimate estimate;
  estimate.requiresFuncReturnMaterialization = requiresFuncReturnMaterialization(analysis);
  estimate.materializationKind = estimate.requiresFuncReturnMaterialization ? "func.return" : "none";

  switch (strategy) {
  case PimConvLoweringLegacy:
  case PimConvLoweringPackedIm2Col:
    estimate.estimatedMvmCount = static_cast<uint64_t>(std::max<int64_t>(1, geo.p));
    break;
  case PimConvLoweringStreamedPatch:
  case PimConvLoweringStreamedPacked:
  case PimConvLoweringOutputChannelTiled: {
    uint64_t chunkPositions = chooseStreamChunkPositions(geo, /*packFactor=*/1);
    estimate.estimatedMvmCount = static_cast<uint64_t>(std::max<int64_t>(1, geo.p));
    estimate.estimatedOutputFragments =
      std::max<uint64_t>(1, static_cast<uint64_t>(ceilIntegerDivide(geo.p, static_cast<int64_t>(chunkPositions))));
    break;
  }
  case PimConvLoweringInputKTiled: {
    const int64_t numKSlices = ceilIntegerDivide(geo.k, geo.xbarSize);
    const uint64_t maxLanesPerBatch =
      std::max<uint64_t>(1,
                         static_cast<uint64_t>(crossbarCountInCore.getValue())
                           / 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),
                          maxLanesPerBatch,
                          static_cast<uint64_t>(std::max<int64_t>(1, geo.outWidth))}));
    estimate.estimatedMvmCount =
      static_cast<uint64_t>(std::max<int64_t>(1, geo.p)) * static_cast<uint64_t>(std::max<int64_t>(1, numKSlices));
    estimate.estimatedReductionVAddCount =
      static_cast<uint64_t>(std::max<int64_t>(1, geo.p))
      * static_cast<uint64_t>(std::max<int64_t>(0, numKSlices - 1) + (geo.hasBias ? 1 : 0));
    estimate.estimatedOutputFragments = static_cast<uint64_t>(std::max<int64_t>(1, geo.batchSize))
                                        * static_cast<uint64_t>(std::max<int64_t>(1, geo.outHeight))
                                        * static_cast<uint64_t>(
                                          ceilIntegerDivide(geo.outWidth, static_cast<int64_t>(rowChunkWidth)));
    estimate.perOutputPositionReduction = numKSlices > 1;
    estimate.fullInputBroadcastExpected = estimate.estimatedOutputFragments > 1;
    if (estimate.requiresFuncReturnMaterialization && estimate.estimatedOutputFragments >= 128) {
      estimate.giantCollectorConcatExpected = true;
      estimate.materializationKind = "single_collector_concat";
      estimate.concatOperandCount = estimate.estimatedOutputFragments;
      estimate.collectorCore = "scheduled_post_merge";
    }
    break;
  }
  case PimConvLoweringTiled2D:
    estimate.estimatedMvmCount = static_cast<uint64_t>(std::max<int64_t>(1, geo.p));
    break;
  case PimConvLoweringDepthwise:
  case PimConvLoweringAuto:
    break;
  }

  if (estimate.requiresFuncReturnMaterialization && estimate.materializationKind == "none")
    estimate.materializationKind = "func.return";
  return estimate;
}

static std::string formatShape(ArrayRef<int64_t> dims) {
  std::string text;
  llvm::raw_string_ostream os(text);
  os << "[";
  for (size_t i = 0; i < dims.size(); ++i) {
    if (i != 0)
      os << "x";
    os << dims[i];
  }
  os << "]";
  return text;
}

static std::string collapseWhitespace(StringRef text) {
  std::string out;
  out.reserve(text.size());
  bool lastWasSpace = false;
  for (char c : text) {
    bool isSpace = std::isspace(static_cast<unsigned char>(c));
    if (isSpace) {
      if (!lastWasSpace && !out.empty())
        out.push_back(' ');
      lastWasSpace = true;
      continue;
    }
    out.push_back(c);
    lastWasSpace = false;
  }
  return out;
}

static std::string abbreviate(StringRef text, size_t maxLen) {
  if (text.size() <= maxLen)
    return text.str();
  return (text.take_front(maxLen - 3) + "...").str();
}

static std::string abbreviateFromEnd(StringRef text, size_t maxLen) {
  if (text.size() <= maxLen)
    return text.str();
  return ("..." + text.take_back(maxLen - 3)).str();
}

static std::string summarizeLocation(Location loc, size_t maxLen = 44) {
  std::string text;
  llvm::raw_string_ostream os(text);
  loc.print(os);
  os.flush();
  std::string collapsed = collapseWhitespace(text);
  if (collapsed.size() <= maxLen)
    return collapsed;
  if (collapsed.find('/') != std::string::npos || collapsed.find('#') != std::string::npos)
    return abbreviateFromEnd(collapsed, maxLen);
  return abbreviate(collapsed, maxLen);
}

static std::string alignCell(StringRef text, size_t width, bool rightAlign = false) {
  std::string cell = text.str();
  if (cell.size() < width) {
    size_t padding = width - cell.size();
    if (rightAlign)
      cell.insert(cell.begin(), padding, ' ');
    else
      cell.append(padding, ' ');
  }
  return cell;
}

static bool hasSameStaticTensorType(Value value, Type expectedType) {
  auto valueType = dyn_cast<RankedTensorType>(value.getType());
  auto expectedTensorType = dyn_cast<RankedTensorType>(expectedType);
  return valueType && expectedTensorType && valueType.hasStaticShape() && valueType == expectedTensorType;
}

static bool isSplatConstantValue(Value value, DenseElementsAttr& denseAttr) {
  denseAttr = getHostConstDenseElementsAttr(value);
  return static_cast<bool>(denseAttr) && denseAttr.isSplat();
}

static bool isPerChannelConstantValue(Value value, RankedTensorType currentType, DenseElementsAttr& denseAttr) {
  denseAttr = getHostConstDenseElementsAttr(value);
  if (!denseAttr || denseAttr.isSplat())
    return false;

  auto constantType = dyn_cast<RankedTensorType>(denseAttr.getType());
  if (!constantType || !constantType.hasStaticShape())
    return false;

  const int64_t channels = currentType.getDimSize(1);
  if (constantType.getRank() == 1)
    return constantType.getDimSize(0) == channels;
  if (constantType.getRank() == 2)
    return constantType.getDimSize(0) == 1 && constantType.getDimSize(1) == channels;
  if (constantType.getRank() != 4)
    return false;
  return constantType.getDimSize(0) == 1 && constantType.getDimSize(1) == channels
         && constantType.getDimSize(2) == 1 && constantType.getDimSize(3) == 1;
}

static std::optional<DistributedTensorStep>
classifyDistributedBinaryConsumer(Operation* user,
                                  Value currentValue,
                                  Value lhs,
                                  Value rhs,
                                  DistributedTensorOpKind kind,
                                  bool allowFragmentOnRhs,
                                  std::string& failureDetail) {
  if (user->getNumResults() != 1 || !hasSameStaticTensorType(user->getResult(0), currentValue.getType())) {
    failureDetail = "result type mismatch";
    return std::nullopt;
  }

  auto currentType = cast<RankedTensorType>(currentValue.getType());
  DenseElementsAttr constantAttr;
  if (lhs == currentValue) {
    DistributedTensorConstantKind constantKind = DistributedTensorConstantKind::None;
    if (isSplatConstantValue(rhs, constantAttr))
      constantKind = DistributedTensorConstantKind::Splat;
    else if (isPerChannelConstantValue(rhs, currentType, constantAttr))
      constantKind = DistributedTensorConstantKind::PerChannel;
    else
      failureDetail = "unsupported rhs broadcast";
    if (constantKind == DistributedTensorConstantKind::None)
      return std::nullopt;
    return DistributedTensorStep {user, kind, constantAttr, constantKind, /*fragmentOnLhs=*/true, std::nullopt};
  }

  if (rhs == currentValue && allowFragmentOnRhs) {
    DistributedTensorConstantKind constantKind = DistributedTensorConstantKind::None;
    if (isSplatConstantValue(lhs, constantAttr))
      constantKind = DistributedTensorConstantKind::Splat;
    else if (isPerChannelConstantValue(lhs, currentType, constantAttr))
      constantKind = DistributedTensorConstantKind::PerChannel;
    else
      failureDetail = "unsupported lhs broadcast";
    if (constantKind == DistributedTensorConstantKind::None)
      return std::nullopt;
    return DistributedTensorStep {user, kind, constantAttr, constantKind, /*fragmentOnLhs=*/false, std::nullopt};
  }

  failureDetail = "conv result is not the supported binary operand";
  return std::nullopt;
}

static bool covers(RowInterval acquired, RowInterval needed) {
  return acquired.begin <= needed.begin && acquired.end >= needed.end;
}

static bool canConsumeRowStripHwcInput(const ConvLoweringState& state, StringRef& failureReason) {
  if (state.batchSize != 1) {
    failureReason = "unsupported_batch";
    return false;
  }
  if (state.group != 1) {
    failureReason = "unsupported_groups";
    return false;
  }
  if (isDepthwiseConv(state.group, state.numChannelsIn, state.numChannelsOut, state.numChannelsInPerGroup)) {
    failureReason = "unsupported_depthwise";
    return false;
  }
  if (state.strideHeight != 1 || state.strideWidth != 1) {
    failureReason = "unsupported_stride";
    return false;
  }
  if (state.dilationHeight != 1 || state.dilationWidth != 1) {
    failureReason = "unsupported_dilation";
    return false;
  }
  if (state.padHeightBegin != state.padHeightEnd || state.padWidthBegin != state.padWidthEnd) {
    failureReason = "unsupported_padding";
    return false;
  }
  if (state.padHeightBegin != 1 || state.padWidthBegin != 1) {
    failureReason = "unsupported_padding";
    return false;
  }
  if (state.wHeight != 3 || state.wWidth != 3) {
    failureReason = "unsupported_kernel";
    return false;
  }
  if (state.outHeight != state.xHeight || state.outWidth != state.xWidth) {
    failureReason = "unsupported_output_shape";
    return false;
  }
  if (!getHostConstDenseElementsAttr(state.w)) {
    failureReason = "non_constant_weights";
    return false;
  }
  if (state.hasBias && !getHostConstDenseElementsAttr(state.b)) {
    failureReason = "non_constant_bias";
    return false;
  }
  failureReason = "";
  return true;
}

static std::string stringifyDistributedTensorOpKind(DistributedTensorOpKind kind) {
  switch (kind) {
  case DistributedTensorOpKind::Relu: return "Relu";
  case DistributedTensorOpKind::Sigmoid: return "Sigmoid";
  case DistributedTensorOpKind::Add: return "Add";
  case DistributedTensorOpKind::Sub: return "Sub";
  case DistributedTensorOpKind::Mul: return "Mul";
  case DistributedTensorOpKind::Div: return "Div";
  case DistributedTensorOpKind::Conv: return "Conv";
  }
  llvm_unreachable("unknown distributed tensor op kind");
}

[[maybe_unused]] static DistributedConvAnalysis analyzeDistributedConvConsumers(ONNXConvOp convOp) {
  DistributedConvAnalysis analysis;
  analysis.replacementOp = convOp;

  Value currentValue = convOp.getResult();
  while (true) {
    if (currentValue.use_empty()) {
      analysis.barrierKind = DistributedConvBarrierKind::DeadValue;
      analysis.barrierDetail = "result has no users";
      return analysis;
    }

    if (!currentValue.hasOneUse()) {
      analysis.barrierKind = DistributedConvBarrierKind::Fanout;
      analysis.barrierDetail = "value has multiple users";
      return analysis;
    }

    Operation* user = *currentValue.getUsers().begin();
    if (isa<func::ReturnOp>(user)) {
      analysis.barrierKind = DistributedConvBarrierKind::Return;
      analysis.barrierDetail = "materialize at func.return";
      return analysis;
    }

    std::optional<DistributedTensorStep> step;
    std::string failureDetail;
    if (auto reluOp = dyn_cast<ONNXReluOp>(user)) {
      if (hasSameStaticTensorType(reluOp.getResult(), currentValue.getType()))
        step = DistributedTensorStep {
          user, DistributedTensorOpKind::Relu, {}, DistributedTensorConstantKind::None, true, std::nullopt};
      else
        failureDetail = "relu result type mismatch";
    }
    else if (auto sigmoidOp = dyn_cast<ONNXSigmoidOp>(user)) {
      if (hasSameStaticTensorType(sigmoidOp.getResult(), currentValue.getType()))
        step = DistributedTensorStep {
          user, DistributedTensorOpKind::Sigmoid, {}, DistributedTensorConstantKind::None, true, std::nullopt};
      else
        failureDetail = "sigmoid result type mismatch";
    }
    else if (auto addOp = dyn_cast<ONNXAddOp>(user)) {
      step = classifyDistributedBinaryConsumer(
        user, currentValue, addOp.getA(), addOp.getB(), DistributedTensorOpKind::Add, /*allowFragmentOnRhs=*/true,
        failureDetail);
    }
    else if (auto subOp = dyn_cast<ONNXSubOp>(user)) {
      step = classifyDistributedBinaryConsumer(
        user, currentValue, subOp.getA(), subOp.getB(), DistributedTensorOpKind::Sub, /*allowFragmentOnRhs=*/true,
        failureDetail);
    }
    else if (auto mulOp = dyn_cast<ONNXMulOp>(user)) {
      step = classifyDistributedBinaryConsumer(
        user, currentValue, mulOp.getA(), mulOp.getB(), DistributedTensorOpKind::Mul, /*allowFragmentOnRhs=*/true,
        failureDetail);
    }
    else if (auto divOp = dyn_cast<ONNXDivOp>(user)) {
      step = classifyDistributedBinaryConsumer(
        user, currentValue, divOp.getA(), divOp.getB(), DistributedTensorOpKind::Div, /*allowFragmentOnRhs=*/false,
        failureDetail);
      if (step) {
        auto denseAttr = dyn_cast<DenseFPElementsAttr>(step->constantAttr);
        if (!denseAttr) {
          failureDetail = "div requires floating-point splat constant";
          step.reset();
        }
      }
    }
    else if (auto nextConv = dyn_cast<ONNXConvOp>(user)) {
      failureDetail = "onnx.Conv distributed consumer blocked by dim0-only whole-batch materialization in MergeComputeNodes";
    }
    else {
      failureDetail = (user->getName().getStringRef() + " is not distributed-aware yet").str();
    }

    if (!step) {
      analysis.barrierKind = DistributedConvBarrierKind::UnsupportedConsumer;
      analysis.barrierDetail = failureDetail;
      return analysis;
    }

    analysis.replacementOp = user;
    analysis.steps.push_back(*step);
    currentValue = user->getResult(0);
  }
}

static void rewriteDistributedConvReport(const DistributedConvReportTotals& totals) {
  std::fstream reportFile = openReportFile("conv_distributed_consumption_report");
  if (!reportFile.is_open())
    return;

  reportFile << "# PIM Conv Distributed Consumption Report\n\n";
  reportFile << "Totals:\n";
  reportFile << "- convs_seen: " << totals.totalConvs << "\n";
  reportFile << "- distributed_tensors_created: " << totals.distributedTensorsCreated << "\n";
  reportFile << "- distributed_values_propagated: " << totals.distributedValuesPropagated << "\n";
  reportFile << "- distributed_consumers_handled_locally: " << totals.distributedConsumersHandled << "\n";
  reportFile << "- distributed_conv_inputs_seen: " << totals.distributedConvInputsSeen << "\n";
  reportFile << "- distributed_conv_inputs_consumed: " << totals.distributedConvInputsConsumed << "\n";
  reportFile << "- materialization_barriers_inserted: " << totals.materializationBarriersInserted << "\n\n";

  if (!totals.barrierReasons.empty()) {
    reportFile << "Materialization barriers:\n";
    for (const auto& [reason, count] : totals.barrierReasons)
      reportFile << "- " << reason << ": " << count << "\n";
    reportFile << "\n";
  }

  if (!totals.fallbackReasons.empty()) {
    reportFile << "Fallback / no-distribution reasons:\n";
    for (const auto& [reason, count] : totals.fallbackReasons)
      reportFile << "- " << reason << ": " << count << "\n";
    reportFile << "\n";
  }

  if (!totals.chains.empty()) {
    reportFile << "Chains:\n";
    for (const DistributedChainReportEntry& chain : totals.chains) {
      reportFile << "- chain_id: " << chain.chainId << "\n";
      reportFile << "  chain_length: " << chain.chainLength << "\n";
      reportFile << "  producer_kind: " << chain.producerKind << "\n";
      reportFile << "  distributed_ops: " << chain.distributedOps << "\n";
      reportFile << "  materialization_points: " << chain.materializationPoints << "\n";
      reportFile << "  first_materialization_reason: " << chain.firstMaterializationReason << "\n";
      reportFile << "  max_live_fragments: " << chain.maxLiveFragments << "\n";
      reportFile << "  max_fragment_fanout: " << chain.maxFragmentFanout << "\n";
      reportFile << "  patch_builder_core_count: " << chain.patchBuilderCoreCount << "\n";
      reportFile << "  central_junction_detected: " << chain.centralJunctionDetected << "\n";
      reportFile << "  conv_input_materialization_kind: " << chain.convInputMaterializationKind << "\n";
      reportFile << "  local_patch_fragments: " << chain.localPatchFragments << "\n";
      reportFile << "  remote_patch_fragments: " << chain.remotePatchFragments << "\n";
      reportFile << "  halo_transfer_count: " << chain.haloTransferCount << "\n";
      reportFile << "  grouped_transfer_count: " << chain.groupedTransferCount << "\n";
      if (!chain.fallbackReason.empty())
        reportFile << "  fallback_reason: " << chain.fallbackReason << "\n";
    }
  }
}

[[maybe_unused]] static void recordDistributedConvOutcome(const DistributedConvAnalysis& analysis) {
  static std::mutex reportMutex;
  static DistributedConvReportTotals totals;

  std::string barrierKey = stringifyDistributedConvBarrierKind(analysis.barrierKind).str();
  if (!analysis.barrierDetail.empty())
    barrierKey += ": " + analysis.barrierDetail;

  std::lock_guard<std::mutex> guard(reportMutex);
  totals.totalConvs++;
  if (analysis.hasLocalConsumers()) {
    totals.distributedTensorsCreated++;
    totals.distributedValuesPropagated += analysis.steps.size();
    totals.distributedConsumersHandled += llvm::count_if(analysis.steps, [](const DistributedTensorStep& step) {
      return step.kind != DistributedTensorOpKind::Conv;
    });
    totals.distributedConvInputsSeen += llvm::count_if(analysis.steps, [](const DistributedTensorStep& step) {
      return step.kind == DistributedTensorOpKind::Conv;
    });
    totals.distributedConvInputsConsumed += llvm::count_if(analysis.steps, [](const DistributedTensorStep& step) {
      return step.kind == DistributedTensorOpKind::Conv;
    });
    totals.materializationBarriersInserted++;
    totals.barrierReasons[barrierKey]++;
  }
  else {
    totals.fallbackReasons[barrierKey]++;
  }
  DistributedChainReportEntry chain;
  chain.chainId = totals.totalConvs;
  chain.chainLength = analysis.steps.size() + 1;
  chain.producerKind = "Conv";
  chain.materializationPoints = stringifyDistributedConvBarrierKind(analysis.barrierKind).str();
  chain.firstMaterializationReason = analysis.barrierDetail;
  chain.maxFragmentFanout = analysis.barrierKind == DistributedConvBarrierKind::Fanout ? 2 : 1;
  std::string ops;
  for (size_t index = 0; index < analysis.steps.size(); ++index) {
    if (!ops.empty())
      ops += ", ";
    ops += stringifyDistributedTensorOpKind(analysis.steps[index].kind);
  }
  chain.distributedOps = ops;
  if (analysis.hasDistributedConvConsumer()) {
    chain.convInputMaterializationKind = "distributed_with_halo_exchange";
    chain.patchBuilderCoreCount = 1;
  }
  if (totals.chains.size() == 16)
    totals.chains.erase(totals.chains.begin());
  totals.chains.push_back(std::move(chain));
  rewriteDistributedConvReport(totals);
}

static std::string makeDivider(ArrayRef<size_t> widths) {
  std::string divider = "+";
  for (size_t width : widths) {
    divider.append(width + 2, '-');
    divider.push_back('+');
  }
  return divider;
}

static void printConvReportLegend(std::fstream& reportFile) {
  reportFile << "# PIM Conv Lowering Report\n\n";
  reportFile << "Legend:\n";
  reportFile << "- `id`: sequential Conv report entry index within this compiler invocation.\n";
  reportFile << "- `where`: summarized MLIR location of the Conv op.\n";
  reportFile << "- `mode`: whether the selected strategy came from `auto` policy or a forced compiler option.\n";
  reportFile << "- `strategy`: selected Conv lowering algorithm.\n";
  reportFile << "- `input`, `weight`, `output`: tensor shapes of the Conv operands/result.\n";
  reportFile << "- `groups`: ONNX Conv group count.\n";
  reportFile << "- `K`: logical reduction size, `CinPerGroup * Kh * Kw`.\n";
  reportFile << "- `C`: logical output-channel width handled by the selected strategy.\n";
  reportFile << "- `P`: total output positions, `N * Hout * Wout`.\n";
  reportFile << "- `X`: crossbar size.\n";
  reportFile << "- `pack`: packed spatial positions per MVM group, `floor(X / max(K, C))`.\n";
  reportFile << "- `im2col`: total explicit im2col element count, `P * K`.\n";
  reportFile << "- `im2col_budget`: maximum im2col element budget allowed by the compiler option.\n";
  reportFile << "- `stream_chunk`: output positions materialized per streamed chunk, or `-` when not applicable.\n";
  reportFile << "- `batch_size`: logical batch size passed into the compute lowering for this Conv form.\n";
  reportFile << "- `batches`: number of repeated compute batches emitted for the Conv.\n";
  reportFile << "- `spatial_compute_batch`: whether ONNX-to-Spatial used `spat.compute_batch` for this Conv.\n";
  reportFile << "- `batched_instruction_emission`: whether the lowering is expected to reach batched PIM emission.\n";
  reportFile << "- `reason`: strategy-selection reason.\n";
  reportFile << "- Per-conv details below the table include profitability estimates and materialization diagnostics.\n";
  reportFile << "- placeholders like `[7]`: value too long for the table cell; see the appendix at the end.\n\n";
}

struct ConvReportOverflow {
  uint64_t placeholderId;
  SmallVector<uint64_t, 4> rowIds;
  std::string column;
  std::string value;
};

static StringRef describeConvLoweringStrategy(PimConvLoweringType strategy) {
  switch (strategy) {
  case PimConvLoweringAuto: return "Automatic policy selection.";
  case PimConvLoweringLegacy: return "Legacy Conv lowering path kept for compatibility.";
  case PimConvLoweringDepthwise: return "Specialized depthwise lowering that avoids generic cross-channel GEMM mixing.";
  case PimConvLoweringPackedIm2Col: return "Explicit im2col plus packed GEMM lowering for Conv shapes that fit well in one crossbar.";
  case PimConvLoweringStreamedPatch: return "Chunked per-patch streaming Conv lowering without global im2col materialization.";
  case PimConvLoweringStreamedPacked: return "Chunked streamed Conv lowering that still packs multiple output positions per MVM group.";
  case PimConvLoweringOutputChannelTiled: return "Conv lowering that splits output channels across tiles when C exceeds one crossbar width.";
  case PimConvLoweringInputKTiled: return "Conv lowering that splits the reduction dimension K across tiles and accumulates partial sums.";
  case PimConvLoweringTiled2D: return "Conv lowering that tiles both K and output channels because neither dimension fits one crossbar.";
  }
  llvm_unreachable("unknown conv lowering strategy");
}

static std::string fitConvReportCell(StringRef text,
                                     size_t width,
                                     uint64_t rowId,
                                     StringRef column,
                                     std::vector<ConvReportOverflow>& overflows,
                                     uint64_t& nextPlaceholderId,
                                     bool rightAlign = false) {
  if (text.size() <= width)
    return alignCell(text, width, rightAlign);

  for (ConvReportOverflow& overflow : overflows) {
    if (overflow.column == column && overflow.value == text) {
      if (llvm::find(overflow.rowIds, rowId) == overflow.rowIds.end())
        overflow.rowIds.push_back(rowId);
      std::string placeholder = "[" + std::to_string(overflow.placeholderId) + "]";
      return alignCell(placeholder, width, rightAlign);
    }
  }

  std::string placeholder = "[" + std::to_string(nextPlaceholderId++) + "]";
  overflows.push_back({nextPlaceholderId - 1, {rowId}, column.str(), text.str()});
  return alignCell(placeholder, width, rightAlign);
}

static void writeConvReportTable(std::fstream& reportFile, ArrayRef<ConvReportEntry> entries) {
  static constexpr size_t kIdWidth = 4;
  static constexpr size_t kWhereWidth = 24;
  static constexpr size_t kModeWidth = 6;
  static constexpr size_t kStrategyWidth = 20;
  static constexpr size_t kShapeWidth = 14;
  static constexpr size_t kGroupsWidth = 3;
  static constexpr size_t kSmallWidth = 5;
  static constexpr size_t kPWidth = 10;
  static constexpr size_t kIm2colWidth = 10;
  static constexpr size_t kChunkWidth = 8;
  static constexpr size_t kFlagWidth = 3;
  static constexpr size_t kReasonWidth = 24;

  const SmallVector<size_t, 21> widths = {
    kIdWidth,        kWhereWidth,   kModeWidth,     kStrategyWidth, kShapeWidth, kShapeWidth, kShapeWidth, kGroupsWidth,
    kSmallWidth,     kSmallWidth,   kPWidth,        kSmallWidth, kSmallWidth, kIm2colWidth, kIm2colWidth,
    kChunkWidth,     kSmallWidth,   kSmallWidth,    kFlagWidth,  kFlagWidth,  kReasonWidth,
  };
  const std::string divider = makeDivider(widths);
  std::vector<ConvReportOverflow> overflows;
  uint64_t nextPlaceholderId = 1;

  auto printRow = [&](ArrayRef<std::string> cells) {
    reportFile << "|";
    for (size_t i = 0; i < cells.size(); ++i)
      reportFile << " " << cells[i] << " |";
    reportFile << "\n";
  };

  reportFile << divider << "\n";
  printRow({
    alignCell("id", kIdWidth, true),
    alignCell("where", kWhereWidth),
    alignCell("mode", kModeWidth),
    alignCell("strategy", kStrategyWidth),
    alignCell("input", kShapeWidth),
    alignCell("weight", kShapeWidth),
    alignCell("output", kShapeWidth),
    alignCell("grp", kGroupsWidth, true),
    alignCell("K", kSmallWidth, true),
    alignCell("C", kSmallWidth, true),
    alignCell("P", kPWidth, true),
    alignCell("X", kSmallWidth, true),
    alignCell("pack", kSmallWidth, true),
    alignCell("im2col", kIm2colWidth, true),
    alignCell("budget", kIm2colWidth, true),
    alignCell("chunk", kChunkWidth, true),
    alignCell("batch", kSmallWidth, true),
    alignCell("nbat", kSmallWidth, true),
    alignCell("scb", kFlagWidth),
    alignCell("bie", kFlagWidth),
    alignCell("reason", kReasonWidth),
  });
  reportFile << divider << "\n";

  for (const ConvReportEntry& entry : entries) {
    printRow({
      alignCell(std::to_string(entry.id), kIdWidth, true),
      fitConvReportCell(entry.where, kWhereWidth, entry.id, "where", overflows, nextPlaceholderId),
      alignCell(entry.mode, kModeWidth),
      fitConvReportCell(entry.strategy, kStrategyWidth, entry.id, "strategy", overflows, nextPlaceholderId),
      fitConvReportCell(entry.inputShape, kShapeWidth, entry.id, "input", overflows, nextPlaceholderId),
      fitConvReportCell(entry.weightShape, kShapeWidth, entry.id, "weight", overflows, nextPlaceholderId),
      fitConvReportCell(entry.outputShape, kShapeWidth, entry.id, "output", overflows, nextPlaceholderId),
      alignCell(std::to_string(entry.groups), kGroupsWidth, true),
      alignCell(std::to_string(entry.k), kSmallWidth, true),
      alignCell(std::to_string(entry.c), kSmallWidth, true),
      alignCell(std::to_string(entry.p), kPWidth, true),
      alignCell(std::to_string(entry.xbarSize), kSmallWidth, true),
      alignCell(std::to_string(entry.pack), kSmallWidth, true),
      alignCell(std::to_string(entry.im2colElements), kIm2colWidth, true),
      alignCell(std::to_string(entry.im2colBudget), kIm2colWidth, true),
      fitConvReportCell(entry.chunkText, kChunkWidth, entry.id, "stream_chunk", overflows, nextPlaceholderId, true),
      alignCell(std::to_string(entry.batchSize), kSmallWidth, true),
      alignCell(std::to_string(entry.numberOfBatches), kSmallWidth, true),
      alignCell(entry.spatialComputeBatch, kFlagWidth),
      alignCell(entry.batchedInstructionEmission, kFlagWidth),
      fitConvReportCell(entry.reason, kReasonWidth, entry.id, "reason", overflows, nextPlaceholderId),
    });
  }
  reportFile << divider << "\n";

  if (overflows.empty())
    reportFile << "\n";
  else {
    reportFile << "\nAppendix:\n";
    for (const ConvReportOverflow& overflow : overflows) {
      reportFile << "  [" << overflow.placeholderId << "] rows ";
      for (size_t i = 0; i < overflow.rowIds.size(); ++i) {
        if (i != 0)
          reportFile << ", ";
        reportFile << overflow.rowIds[i];
      }
      reportFile << ", " << overflow.column << ": " << overflow.value << "\n";
    }
    reportFile << "\n";
  }

  reportFile << "Per-Conv Details:\n";
  for (const ConvReportEntry& entry : entries) {
    reportFile << "- Conv " << entry.id << ": mode=" << entry.mode << ", strategy=" << entry.strategy
               << ", reason=" << entry.reason << "\n";
    reportFile << "  K=" << entry.k << ", Cout=" << entry.c << ", output_positions=" << entry.p
               << ", estimated_mvm_count=" << entry.estimatedMvmCount
               << ", estimated_reduction_vadd_count=" << entry.estimatedReductionVAddCount
               << ", estimated_output_fragments=" << entry.estimatedOutputFragments << "\n";
    reportFile << "  materialization_required_at_func_return=" << entry.materializationRequiredAtReturn
               << ", materialization_kind=" << entry.materializationKind
               << ", giant_collector_concat_expected=" << entry.giantCollectorConcatExpected
               << ", concat_operand_count=" << entry.concatOperandCount
               << ", collector_core=" << entry.collectorCore
               << ", full_input_broadcast_expected=" << entry.fullInputBroadcastExpected << "\n";
    if (!entry.rejectedAutoStrategy.empty())
      reportFile << "  rejected_auto_strategy=" << entry.rejectedAutoStrategy << "\n";
    if (!entry.fallbackReason.empty())
      reportFile << "  fallback_reason=" << entry.fallbackReason << "\n";
  }
  reportFile << "\n";

  llvm::SmallVector<PimConvLoweringType, 8> usedStrategies;
  for (const ConvReportEntry& entry : entries) {
    PimConvLoweringType strategy = PimConvLoweringAuto;
    for (PimConvLoweringType candidate : {
           PimConvLoweringAuto,
           PimConvLoweringLegacy,
           PimConvLoweringDepthwise,
           PimConvLoweringPackedIm2Col,
           PimConvLoweringStreamedPatch,
           PimConvLoweringStreamedPacked,
           PimConvLoweringOutputChannelTiled,
           PimConvLoweringInputKTiled,
           PimConvLoweringTiled2D,
         }) {
      if (entry.strategy == stringifyConvLoweringStrategy(candidate)) {
        strategy = candidate;
        break;
      }
    }
    if (llvm::find(usedStrategies, strategy) == usedStrategies.end())
      usedStrategies.push_back(strategy);
  }

  reportFile << "Strategies used in this report:\n";
  for (PimConvLoweringType strategy : usedStrategies)
    reportFile << "- `" << stringifyConvLoweringStrategy(strategy).str() << "`: "
               << describeConvLoweringStrategy(strategy).str() << "\n";
}

static void rewriteConvLoweringReport(ArrayRef<ConvReportEntry> entries) {
  std::fstream reportFile = openReportFile("conv_lowering_report");
  if (!reportFile.is_open())
    return;
  printConvReportLegend(reportFile);
  writeConvReportTable(reportFile, entries);
}

[[maybe_unused]] static FailureOr<PimConvLoweringType> resolveRequestedConvLoweringStrategy(ONNXConvOp convOp) {
  if (!useExperimentalConvImpl)
    return pimConvLowering.getValue();

  if (pimConvLowering != PimConvLoweringAuto && pimConvLowering != PimConvLoweringPackedIm2Col) {
    convOp.emitOpError() << "--use-experimental-conv-impl conflicts with --pim-conv-lowering="
                         << stringifyConvLoweringStrategy(pimConvLowering);
    return failure();
  }
  return PimConvLoweringPackedIm2Col;
}

static ConvLoweringDecision chooseConvLoweringStrategy(const ConvGeometry& geo,
                                                       PimConvLoweringType requested,
                                                       const DistributedConvAnalysis& analysis) {
  if (requested != PimConvLoweringAuto)
    return {requested, "forced by compiler option", /*isAuto=*/false, "", ""};

  // 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, "depthwise convolution", /*isAuto=*/true, "", ""};
  if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize && geo.pack >= 2 && geo.im2colElements <= pimConvIm2colMaxElements)
    return {PimConvLoweringPackedIm2Col,
            "fits crossbar, packing useful, and global im2col fits budget",
            /*isAuto=*/true,
            "",
            ""};
  if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize && geo.pack >= 2 && geo.im2colElements > pimConvIm2colMaxElements)
    return {PimConvLoweringStreamedPacked,
            "fits crossbar and packing useful, but global im2col exceeds budget",
            /*isAuto=*/true,
            "",
            ""};
  if (geo.k <= geo.xbarSize && geo.c <= geo.xbarSize)
    return {PimConvLoweringStreamedPatch, "fits crossbar but packing is not useful", /*isAuto=*/true, "", ""};
  if (geo.k <= geo.xbarSize && geo.c > geo.xbarSize)
    return {PimConvLoweringOutputChannelTiled,
            "output channels exceed one crossbar width",
            /*isAuto=*/true,
            "",
            ""};
  if (geo.k > geo.xbarSize && geo.c <= geo.xbarSize) {
    ConvStrategyEstimate estimate = estimateConvStrategy(geo, PimConvLoweringInputKTiled, analysis);
    std::string fallbackReason = "auto rejects input-k-tiled because the reduction-heavy path is force-only for now";
    if (estimate.requiresFuncReturnMaterialization && estimate.perOutputPositionReduction
        && estimate.giantCollectorConcatExpected) {
      fallbackReason += "; func.return would materialize " + std::to_string(estimate.concatOperandCount)
                        + " output fragments through a single collector concat after per-position reductions";
    }
    if (estimate.fullInputBroadcastExpected)
      fallbackReason += "; the current lowering also broadcasts the padded input to many workers";
    return {PimConvLoweringLegacy,
            "fall back to legacy explicit-im2col for the current auto policy",
            /*isAuto=*/true,
            fallbackReason,
            stringifyConvLoweringStrategy(PimConvLoweringInputKTiled).str()};
  }
  return {PimConvLoweringTiled2D, "both reduction K and output channels exceed one crossbar", /*isAuto=*/true, "", ""};
}

[[maybe_unused]] static LogicalResult verifyForcedConvLoweringStrategy(ONNXConvOp convOp,
                                                      const ConvGeometry& geo,
                                                      PimConvLoweringType strategy) {
  switch (strategy) {
  case PimConvLoweringAuto:
  case PimConvLoweringLegacy:
    return success();
  case PimConvLoweringDepthwise:
    if (geo.isDepthwise)
      return success();
    return convOp.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 convOp.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 convOp.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 convOp.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 convOp.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 convOp.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 convOp.emitOpError("forced tiled-2d Conv lowering requires K > X and C > X");
  }
  llvm_unreachable("unknown conv lowering strategy");
}

static void reportConvLoweringDecision(ONNXConvOp convOp,
                                       const ConvGeometry& geo,
                                       const ConvLoweringDecision& decision,
                                       const ConvStrategyEstimate& estimate,
                                       int64_t batchSize,
                                       int64_t numberOfBatches,
                                       bool usesComputeBatch,
                                       bool usesBatchedInstructionEmission,
                                       std::optional<uint64_t> streamChunkPositions = std::nullopt) {
  if (!pimReportConvLowering)
    return;

  const std::string location = summarizeLocation(convOp.getLoc());
  const std::string strategy = stringifyConvLoweringStrategy(decision.strategy).str();
  const std::string mode = decision.isAuto ? "auto" : "forced";
  const std::string inputShape = formatShape(cast<RankedTensorType>(convOp.getX().getType()).getShape());
  const std::string weightShape = formatShape(cast<RankedTensorType>(convOp.getW().getType()).getShape());
  const std::string outputShape = formatShape(cast<RankedTensorType>(convOp.getY().getType()).getShape());
  const std::string chunkText = streamChunkPositions ? std::to_string(*streamChunkPositions) : "-";
  const std::string scbText = usesComputeBatch ? "yes" : "no";
  const std::string bieText = usesBatchedInstructionEmission ? "yes" : "no";

  static uint64_t reportIndex = 0;
  const uint64_t currentIndex = ++reportIndex;
  static std::mutex reportMutex;
  static std::vector<ConvReportEntry> reportEntries;
  std::lock_guard<std::mutex> lock(reportMutex);
  reportEntries.push_back({
    currentIndex,
    location,
    strategy,
    mode,
    inputShape,
    weightShape,
    outputShape,
    geo.group,
    geo.k,
    geo.c,
    geo.p,
    geo.xbarSize,
    geo.pack,
    geo.im2colElements,
    pimConvIm2colMaxElements,
    chunkText,
    batchSize,
    numberOfBatches,
    scbText,
    bieText,
    decision.reason,
    decision.fallbackReason,
    decision.rejectedAutoStrategy,
    estimate.estimatedMvmCount,
    estimate.estimatedReductionVAddCount,
    estimate.estimatedOutputFragments,
    estimate.requiresFuncReturnMaterialization ? "yes" : "no",
    estimate.materializationKind,
    estimate.concatOperandCount,
    estimate.collectorCore,
    estimate.giantCollectorConcatExpected ? "yes" : "no",
    estimate.fullInputBroadcastExpected ? "yes" : "no",
  });
  rewriteConvLoweringReport(reportEntries);
}

static Value expandBiasIfNeeded(Value bias, PatternRewriter& rewriter, Location loc) {
  auto biasType = cast<RankedTensorType>(bias.getType());
  if (biasType.getRank() != 1)
    return bias;

  auto expandedBiasType = RankedTensorType::get({1, biasType.getDimSize(0)}, biasType.getElementType());
  return tensor::ExpandShapeOp::create(rewriter,
                                       loc,
                                       expandedBiasType,
                                       bias,
                                       SmallVector<ReassociationIndices> {
                                         {0, 1}
  });
}

static int64_t findLargestDivisorAtMost(int64_t value, int64_t limit) {
  assert(value > 0 && "expected positive value");
  limit = std::min(value, limit);
  for (int64_t candidate = limit; candidate >= 1; --candidate)
    if (value % candidate == 0)
      return candidate;
  return 1;
}

static Value createZeroPaddedTensor(Value value,
                                    RankedTensorType resultType,
                                    ArrayRef<int64_t> lowPadValues,
                                    ArrayRef<int64_t> highPadValues,
                                    PatternRewriter& rewriter,
                                    Location loc) {
  auto valueType = cast<RankedTensorType>(value.getType());
  if (valueType == resultType)
    return value;

  SmallVector<OpFoldResult> lowPads;
  SmallVector<OpFoldResult> highPads;
  lowPads.reserve(lowPadValues.size());
  highPads.reserve(highPadValues.size());
  for (auto lowPad : lowPadValues)
    lowPads.push_back(rewriter.getIndexAttr(lowPad));
  for (auto highPad : highPadValues)
    highPads.push_back(rewriter.getIndexAttr(highPad));

  auto padOp = tensor::PadOp::create(rewriter, loc, resultType, value, lowPads, highPads);
  auto* padBlock = new Block();
  for (int64_t dim = 0, rank = resultType.getRank(); dim < rank; ++dim)
    padBlock->addArgument(rewriter.getIndexType(), loc);
  padOp.getRegion().push_back(padBlock);
  rewriter.setInsertionPointToStart(padBlock);
  auto zero = getOrCreateConstant(
    rewriter, padOp.getOperation(), rewriter.getZeroAttr(resultType.getElementType()), resultType.getElementType());
  tensor::YieldOp::create(rewriter, loc, zero);
  rewriter.setInsertionPointAfter(padOp);
  return padOp.getResult();
}

static Value createConvInputPatch(Value input,
                                  RankedTensorType patchType,
                                  Value batchIndex,
                                  Value channelOffset,
                                  Value inputHeightOffset,
                                  Value inputWidthOffset,
                                  int64_t dilationHeight,
                                  int64_t dilationWidth,
                                  PatternRewriter& rewriter,
                                  Location loc) {
  const int64_t patchChannels = patchType.getDimSize(1);
  const int64_t kernelHeight = patchType.getDimSize(2);
  const int64_t kernelWidth = patchType.getDimSize(3);
  if (dilationHeight == 1 && dilationWidth == 1) {
    SmallVector<OpFoldResult> offsets {batchIndex, channelOffset, inputHeightOffset, inputWidthOffset};
    SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
                                     rewriter.getIndexAttr(patchChannels),
                                     rewriter.getIndexAttr(kernelHeight),
                                     rewriter.getIndexAttr(kernelWidth)};
    return tensor::ExtractSliceOp::create(rewriter, loc, patchType, input, offsets, sizes, getUnitStrides(rewriter, 4));
  }

  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
  auto elementType = patchType.getElementType();
  auto pixelType = RankedTensorType::get({1, patchChannels, 1, 1}, elementType, patchType.getEncoding());
  Value patch = tensor::EmptyOp::create(rewriter, loc, patchType.getShape(), elementType);
  for (int64_t kernelH = 0; kernelH < kernelHeight; ++kernelH) {
    Value sourceHeightOffset = affineAddConst(rewriter, loc, inputHeightOffset, kernelH * dilationHeight, anchorOp);
    for (int64_t kernelW = 0; kernelW < kernelWidth; ++kernelW) {
      Value sourceWidthOffset = affineAddConst(rewriter, loc, inputWidthOffset, kernelW * dilationWidth, anchorOp);
      SmallVector<OpFoldResult> sourceOffsets {batchIndex, channelOffset, sourceHeightOffset, sourceWidthOffset};
      SmallVector<OpFoldResult> sourceSizes {rewriter.getIndexAttr(1),
                                             rewriter.getIndexAttr(patchChannels),
                                             rewriter.getIndexAttr(1),
                                             rewriter.getIndexAttr(1)};
      Value sourcePixel = tensor::ExtractSliceOp::create(
        rewriter, loc, pixelType, input, sourceOffsets, sourceSizes, getUnitStrides(rewriter, 4));
      SmallVector<OpFoldResult> targetOffsets {
        rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(kernelH), rewriter.getIndexAttr(kernelW)};
      patch = tensor::InsertSliceOp::create(
        rewriter, loc, sourcePixel, patch, targetOffsets, sourceSizes, getUnitStrides(rewriter, 4));
    }
  }
  return patch;
}

static Value createCollectedConvOutput(ValueRange gemmRows,
                                       Type convType,
                                       RankedTensorType gemmOutType,
                                       RankedTensorType nhwcType,
                                       RankedTensorType outType,
                                       int64_t numPatches,
                                       int64_t numChannelsOut,
                                       int64_t packFactor,
                                       ArrayRef<DistributedTensorStep> distributedConsumers,
                                       PatternRewriter& rewriter,
                                       Location loc);
static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp, Value x, Value w, Value b);

namespace depthwise {

struct Tiling {
  int64_t outputMultiplier;
  int64_t kernelElements;
  int64_t channelsPerTile;
  int64_t tileInputRows;
  int64_t tileOutputChannels;
  int64_t numChannelTiles;
  int64_t spatialPatchesPerBatch;
  int64_t totalPatches;
};

static std::optional<Tiling> computeTiling(int64_t batchSize,
                                           int64_t numChannelsIn,
                                           int64_t numChannelsOut,
                                           int64_t wHeight,
                                           int64_t wWidth,
                                           int64_t outHeight,
                                           int64_t outWidth) {
  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;

  const int64_t maxChannelsPerTile = std::min(xbarDim / kernelElements, xbarDim / outputMultiplier);
  if (maxChannelsPerTile <= 0)
    return std::nullopt;

  const int64_t channelsPerTile = findLargestDivisorAtMost(numChannelsIn, maxChannelsPerTile);
  const int64_t tileInputRows = channelsPerTile * kernelElements;
  const int64_t tileOutputChannels = channelsPerTile * outputMultiplier;
  if (tileInputRows > xbarDim || tileOutputChannels > xbarDim)
    return std::nullopt;

  return Tiling {
    outputMultiplier,
    kernelElements,
    channelsPerTile,
    tileInputRows,
    tileOutputChannels,
    numChannelsIn / channelsPerTile,
    outHeight * outWidth,
    batchSize * outHeight * outWidth,
  };
}

static Value buildPackedWeights(DenseElementsAttr wDenseAttr,
                                RankedTensorType wType,
                                const Tiling& tiling,
                                PatternRewriter& rewriter,
                                Location loc) {
  auto packedWeightType = RankedTensorType::get(
    {tiling.numChannelTiles, tiling.tileInputRows, tiling.tileOutputChannels}, wType.getElementType());
  SmallVector<Attribute> packedValues(packedWeightType.getNumElements(),
                                      cast<Attribute>(rewriter.getZeroAttr(wType.getElementType())));
  SmallVector<Attribute> sourceValues(wDenseAttr.getValues<Attribute>());

  for (int64_t tileIndex = 0; tileIndex < tiling.numChannelTiles; ++tileIndex) {
    const int64_t channelBase = tileIndex * tiling.channelsPerTile;
    for (int64_t localChannel = 0; localChannel < tiling.channelsPerTile; ++localChannel) {
      const int64_t globalChannel = channelBase + localChannel;
      for (int64_t kernelIndex = 0; kernelIndex < tiling.kernelElements; ++kernelIndex) {
        const int64_t kernelH = kernelIndex / wType.getDimSize(3);
        const int64_t kernelW = kernelIndex % wType.getDimSize(3);
        const int64_t targetRow = localChannel * tiling.kernelElements + kernelIndex;
        for (int64_t multiplierIndex = 0; multiplierIndex < tiling.outputMultiplier; ++multiplierIndex) {
          const int64_t globalOutChannel = globalChannel * tiling.outputMultiplier + multiplierIndex;
          const int64_t sourceFlatIndex =
            ((globalOutChannel * wType.getDimSize(1) * wType.getDimSize(2)) + kernelH) * wType.getDimSize(3) + kernelW;
          const int64_t targetCol = localChannel * tiling.outputMultiplier + multiplierIndex;
          const int64_t targetFlatIndex =
            ((tileIndex * tiling.tileInputRows) + targetRow) * tiling.tileOutputChannels + targetCol;
          packedValues[targetFlatIndex] = sourceValues[sourceFlatIndex];
        }
      }
    }
  }

  auto packedAttr = DenseElementsAttr::get(packedWeightType, packedValues);
  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), packedAttr, packedWeightType);
}

static Value createPaddedInput(Value input,
                               RankedTensorType inputType,
                               int64_t padHeightBegin,
                               int64_t padHeightEnd,
                               int64_t padWidthBegin,
                               int64_t padWidthEnd,
                               PatternRewriter& rewriter,
                               Location loc) {
  if (padHeightBegin == 0 && padHeightEnd == 0 && padWidthBegin == 0 && padWidthEnd == 0)
    return input;

  auto paddedInputType = RankedTensorType::get({inputType.getDimSize(0),
                                                inputType.getDimSize(1),
                                                inputType.getDimSize(2) + padHeightBegin + padHeightEnd,
                                                inputType.getDimSize(3) + padWidthBegin + padWidthEnd},
                                               inputType.getElementType());
  auto computeOp = createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, input, [&](Value computeInput) {
    Value padded = createZeroPaddedTensor(computeInput,
                                          paddedInputType,
                                          {0, 0, padHeightBegin, padWidthBegin},
                                          {0, 0, padHeightEnd, padWidthEnd},
                                          rewriter,
                                          loc);
    spatial::SpatYieldOp::create(rewriter, loc, padded);
  });
  return computeOp.getResult(0);
}

static Value createInputTile(Value input,
                             Value patchIndex,
                             Value channelTileIndex,
                             RankedTensorType inputTileType,
                             const Tiling& tiling,
                             int64_t strideHeight,
                             int64_t strideWidth,
                             int64_t dilationHeight,
                             int64_t dilationWidth,
                             int64_t outWidth,
                             PatternRewriter& rewriter,
                             Location loc) {
  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
  Value batchIndex = affineFloorDivConst(rewriter, loc, patchIndex, tiling.spatialPatchesPerBatch, anchorOp);
  Value batchPatchIndex = affineModConst(rewriter, loc, patchIndex, tiling.spatialPatchesPerBatch, anchorOp);
  Value outHeightIndex = affineFloorDivConst(rewriter, loc, batchPatchIndex, outWidth, anchorOp);
  Value outWidthIndex = affineModConst(rewriter, loc, batchPatchIndex, outWidth, anchorOp);
  Value inputHeightOffset =
    strideHeight == 1 ? outHeightIndex : affineMulConst(rewriter, loc, outHeightIndex, strideHeight, anchorOp);
  Value inputWidthOffset =
    strideWidth == 1 ? outWidthIndex : affineMulConst(rewriter, loc, outWidthIndex, strideWidth, anchorOp);
  Value channelOffset = tiling.channelsPerTile == 1
                        ? channelTileIndex
                        : affineMulConst(rewriter, loc, channelTileIndex, tiling.channelsPerTile, anchorOp);
  Value tile4D = createConvInputPatch(input,
                                      inputTileType,
                                      batchIndex,
                                      channelOffset,
                                      inputHeightOffset,
                                      inputWidthOffset,
                                      dilationHeight,
                                      dilationWidth,
                                      rewriter,
                                      loc);
  auto collapsedType = RankedTensorType::get({1, tiling.tileInputRows}, inputTileType.getElementType());
  return tensor::CollapseShapeOp::create(rewriter,
                                         loc,
                                         collapsedType,
                                         tile4D,
                                         SmallVector<ReassociationIndices> {
                                           {0},
                                           {1, 2, 3}
  });
}

static Value createWeightTile(Value packedWeights,
                              Value channelTileIndex,
                              RankedTensorType packedWeightType,
                              const Tiling& tiling,
                              PatternRewriter& rewriter,
                              Location loc) {
  SmallVector<OpFoldResult> offsets {channelTileIndex, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
                                   rewriter.getIndexAttr(tiling.tileInputRows),
                                   rewriter.getIndexAttr(tiling.tileOutputChannels)};
  auto sliceType =
    RankedTensorType::get({1, tiling.tileInputRows, tiling.tileOutputChannels}, packedWeightType.getElementType());
  Value slice = tensor::ExtractSliceOp::create(
    rewriter, loc, sliceType, packedWeights, offsets, sizes, getUnitStrides(rewriter, 3));
  auto collapsedType =
    RankedTensorType::get({tiling.tileInputRows, tiling.tileOutputChannels}, packedWeightType.getElementType());
  return tensor::CollapseShapeOp::create(rewriter,
                                         loc,
                                         collapsedType,
                                         slice,
                                         SmallVector<ReassociationIndices> {
                                           {0, 1},
                                           {2}
  });
}

static Value createBiasTile(
  Value bias, Value channelTileIndex, const Tiling& tiling, PatternRewriter& rewriter, Location loc) {
  auto biasType = cast<RankedTensorType>(bias.getType());
  auto biasTileType = RankedTensorType::get({1, tiling.tileOutputChannels}, biasType.getElementType());
  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
  Value channelOffset = tiling.tileOutputChannels == 1
                        ? channelTileIndex
                        : affineMulConst(rewriter, loc, channelTileIndex, tiling.tileOutputChannels, anchorOp);
  SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), channelOffset};
  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(tiling.tileOutputChannels)};
  return tensor::ExtractSliceOp::create(rewriter, loc, biasTileType, bias, offsets, sizes, getUnitStrides(rewriter, 2));
}

static Value insertOutputTile(Value rowTile,
                              Value rowAccumulator,
                              Value channelTileIndex,
                              const Tiling& tiling,
                              PatternRewriter& rewriter,
                              Location loc) {
  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
  Value channelOffset = tiling.tileOutputChannels == 1
                        ? channelTileIndex
                        : affineMulConst(rewriter, loc, channelTileIndex, tiling.tileOutputChannels, anchorOp);
  SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), channelOffset};
  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(tiling.tileOutputChannels)};
  return tensor::InsertSliceOp::create(
    rewriter, loc, rowTile, rowAccumulator, offsets, sizes, getUnitStrides(rewriter, 2));
}

static FailureOr<Value> reconstructDepthwiseGemmRows(Value pieces,
                                                     RankedTensorType piecesType,
                                                     RankedTensorType gemmOutType,
                                                     const Tiling& tiling,
                                                     PatternRewriter& rewriter,
                                                     Location loc) {
  auto collectedOp = createSpatCompute<1>(rewriter, loc, TypeRange {gemmOutType}, {}, pieces, [&](Value piecesArg) {
    auto rowType = RankedTensorType::get({1, gemmOutType.getDimSize(1)}, gemmOutType.getElementType());
    Value outputInit = tensor::EmptyOp::create(rewriter, loc, gemmOutType.getShape(), gemmOutType.getElementType());
    Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
    Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
    Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
    Value cNumPatches = getOrCreateIndexConstant(rewriter, anchorOp, tiling.totalPatches);
    Value cNumChannelTiles = getOrCreateIndexConstant(rewriter, anchorOp, tiling.numChannelTiles);

    auto patchLoop = buildNormalizedScfFor(
      rewriter,
      loc,
      c0,
      cNumPatches,
      c1,
      ValueRange {outputInit},
      [&](OpBuilder&,
          Location nestedLoc,
          Value patchIndex,
          ValueRange patchIterArgs,
          SmallVectorImpl<Value>& patchYielded) {
        Value outputAcc = patchIterArgs.front();
        Value rowInit = tensor::EmptyOp::create(rewriter, nestedLoc, rowType.getShape(), rowType.getElementType());
        auto tileLoop = buildNormalizedScfFor(
          rewriter,
          nestedLoc,
          c0,
          cNumChannelTiles,
          c1,
          ValueRange {rowInit},
          [&](OpBuilder&,
              Location tileLoc,
              Value channelTileIndex,
              ValueRange tileIterArgs,
              SmallVectorImpl<Value>& tileYielded) {
            Value rowAcc = tileIterArgs.front();
            MLIRContext* context = rewriter.getContext();
            AffineExpr d0 = getAffineDimExpr(0, context);
            AffineExpr d1 = getAffineDimExpr(1, context);
            Value laneIndex = createOrFoldAffineApply(
              rewriter, tileLoc, (d0 * tiling.totalPatches) + d1, ValueRange {channelTileIndex, patchIndex}, anchorOp);
            auto rowTileType = RankedTensorType::get({1, tiling.tileOutputChannels}, piecesType.getElementType());
            SmallVector<OpFoldResult> pieceOffsets {laneIndex, rewriter.getIndexAttr(0)};
            SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(1),
                                                  rewriter.getIndexAttr(tiling.tileOutputChannels)};
            Value rowTile = tensor::ExtractSliceOp::create(
              rewriter, tileLoc, rowTileType, piecesArg, pieceOffsets, pieceSizes, getUnitStrides(rewriter, 2));
            Value rowNext = insertOutputTile(rowTile, rowAcc, channelTileIndex, tiling, rewriter, tileLoc);
            tileYielded.push_back(rowNext);
            return success();
          });
        if (failed(tileLoop))
          return failure();

        SmallVector<OpFoldResult> rowOffsets {patchIndex, rewriter.getIndexAttr(0)};
        SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(gemmOutType.getDimSize(1))};
        Value outputNext = tensor::InsertSliceOp::create(rewriter,
                                                         nestedLoc,
                                                         tileLoop->results.front(),
                                                         outputAcc,
                                                         rowOffsets,
                                                         rowSizes,
                                                         getUnitStrides(rewriter, 2))
                             .getResult();
        patchYielded.push_back(outputNext);
        return success();
      });
    if (failed(patchLoop))
      return failure();

    spatial::SpatYieldOp::create(rewriter, loc, patchLoop->results.front());
    return success();
  });
  if (failed(collectedOp))
    return failure();
  return collectedOp->getResult(0);
}

static bool canUseStructuredRewrite(const ConvLoweringState& state) {
  if (!getHostConstDenseElementsAttr(state.w))
    return false;

  auto tiling = computeTiling(state.batchSize,
                              state.numChannelsIn,
                              state.numChannelsOut,
                              state.wHeight,
                              state.wWidth,
                              state.outHeight,
                              state.outWidth);
  if (!tiling)
    return false;

  if (!state.hasBias)
    return true;

  auto biasType = dyn_cast<RankedTensorType>(state.b.getType());
  if (!biasType)
    return false;
  if (biasType.getRank() == 1)
    return biasType.getDimSize(0) == state.numChannelsOut;
  if (biasType.getRank() != 2)
    return false;
  return biasType.getDimSize(0) == 1 && biasType.getDimSize(1) == state.numChannelsOut;
}

static FailureOr<Value>
rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter& rewriter, Location loc) {
  auto wDenseAttr = getHostConstDenseElementsAttr(state.w);
  if (!wDenseAttr) {
    convOp->emitOpError("requires constant-derived weights for structured depthwise Spatial lowering");
    return failure();
  }

  auto tiling = computeTiling(state.xType.getDimSize(0),
                              state.xType.getDimSize(1),
                              state.outType.getDimSize(1),
                              state.wType.getDimSize(2),
                              state.wType.getDimSize(3),
                              state.outType.getDimSize(2),
                              state.outType.getDimSize(3));
  if (!tiling) {
    convOp->emitOpError("failed to derive a structured depthwise tiling that fits Spatial weighted VMM lowering");
    return failure();
  }

  Value paddedInput = createPaddedInput(state.x,
                                        state.xType,
                                        state.padHeightBegin,
                                        state.padHeightEnd,
                                        state.padWidthBegin,
                                        state.padWidthEnd,
                                        rewriter,
                                        loc);
  Value packedWeights = buildPackedWeights(wDenseAttr, state.wType, *tiling, rewriter, loc);

  Value expandedBias;
  SmallVector<Value> batchInputs {paddedInput};
  if (state.hasBias) {
    expandedBias = expandBiasIfNeeded(state.b, rewriter, loc);
    auto biasType = dyn_cast<RankedTensorType>(expandedBias.getType());
    if (!biasType || biasType.getRank() != 2 || biasType.getDimSize(0) != 1
        || biasType.getDimSize(1) != state.outType.getDimSize(1)) {
      convOp->emitOpError("requires bias sliceable as tensor<1xCout> for structured depthwise Spatial lowering");
      return failure();
    }
    batchInputs.push_back(expandedBias);
  }

  auto gemmOutType =
    RankedTensorType::get({tiling->totalPatches, state.outType.getDimSize(1)}, state.outType.getElementType());
  auto piecesType = RankedTensorType::get({tiling->totalPatches * tiling->numChannelTiles, tiling->tileOutputChannels},
                                          state.outType.getElementType());
  auto paddedInputType = cast<RankedTensorType>(paddedInput.getType());
  auto inputTileType =
    RankedTensorType::get({1, tiling->channelsPerTile, state.wType.getDimSize(2), state.wType.getDimSize(3)},
                          paddedInputType.getElementType());
  SmallVector<Value> batchWeights;
  if (tiling->numChannelTiles == 1) {
    Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
    batchWeights.push_back(createWeightTile(packedWeights,
                                            c0,
                                            cast<RankedTensorType>(packedWeights.getType()),
                                            *tiling,
                                            rewriter,
                                            loc));
  }
  else {
    batchWeights.push_back(packedWeights);
  }

  auto batchOp = createSpatComputeBatch(
    rewriter,
    loc,
    TypeRange {piecesType},
    tiling->totalPatches * tiling->numChannelTiles,
    batchWeights,
    batchInputs,
    [&](detail::SpatComputeBatchBodyArgs args) {
      auto pickInputByRank = [&](int64_t rank) -> Value {
        for (Value input : args.inputs) {
          auto inputType = dyn_cast<RankedTensorType>(input.getType());
          if (inputType && inputType.getRank() == rank)
            return input;
        }
        return Value();
      };

      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
      Value patchIndex = tiling->numChannelTiles == 1
                         ? args.lane
                         : affineModConst(rewriter, loc, args.lane, tiling->totalPatches, anchorOp);
      Value channelTileIndex = tiling->numChannelTiles == 1
                               ? getOrCreateIndexConstant(rewriter, anchorOp, 0)
                               : affineFloorDivConst(rewriter, loc, args.lane, tiling->totalPatches, anchorOp);
      Value paddedInputArg = pickInputByRank(/*rank=*/4);
      if (!paddedInputArg) {
        convOp->emitOpError("structured depthwise batch body requires a rank-4 padded input block argument");
        return failure();
      }

      Value inputTile = createInputTile(paddedInputArg,
                                        patchIndex,
                                        channelTileIndex,
                                        inputTileType,
                                        *tiling,
                                        state.strideHeight,
                                        state.strideWidth,
                                        state.dilationHeight,
                                        state.dilationWidth,
                                        state.outType.getDimSize(3),
                                        rewriter,
                                        loc);
      Value weightTile = tiling->numChannelTiles == 1
                         ? args.weights.front()
                         : createWeightTile(args.weights.front(),
                                            channelTileIndex,
                                            cast<RankedTensorType>(args.weights.front().getType()),
                                            *tiling,
                                            rewriter,
                                            loc);
      auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.outType.getElementType());
      Value rowTile = spatial::SpatVMMOp::create(rewriter, loc, rowTileType, weightTile, inputTile).getResult();
      if (args.inputs.size() > 1) {
        Value biasArg = pickInputByRank(/*rank=*/2);
        if (!biasArg) {
          convOp->emitOpError("structured depthwise batch body requires a rank-2 bias block argument when bias is present");
          return failure();
        }
        Value biasTile = tiling->numChannelTiles == 1 ? biasArg : createBiasTile(biasArg, channelTileIndex, *tiling, rewriter, loc);
        rowTile = spatial::SpatVAddOp::create(rewriter, loc, rowTileType, rowTile, biasTile).getResult();
      }
      SmallVector<OpFoldResult> outputOffsets {args.lane, rewriter.getIndexAttr(0)};
      SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1),
                                             rewriter.getIndexAttr(tiling->tileOutputChannels)};
      createParallelInsertSliceIntoBatchOutput(
        rewriter, loc, rowTile, args.outputs.front(), outputOffsets, outputSizes, getUnitStrides(rewriter, 2));
      return success();
    });
  if (failed(batchOp))
    return failure();

  auto nhwcType = RankedTensorType::get(
    {state.xType.getDimSize(0), state.outType.getDimSize(2), state.outType.getDimSize(3), state.outType.getDimSize(1)},
    state.outType.getElementType());
  Value collectedRows = batchOp->getResult(0);
  if (tiling->numChannelTiles != 1) {
    auto reconstructedRows =
      reconstructDepthwiseGemmRows(batchOp->getResult(0), piecesType, gemmOutType, *tiling, rewriter, loc);
    if (failed(reconstructedRows))
      return failure();
    collectedRows = *reconstructedRows;
  }

  return createCollectedConvOutput(ValueRange {collectedRows},
                                   state.outType,
                                   gemmOutType,
                                   nhwcType,
                                   state.outType,
                                   tiling->totalPatches,
                                   state.outType.getDimSize(1),
                                   /*packFactor=*/1,
                                   {},
                                   rewriter,
                                   loc);
}

} // namespace depthwise

namespace standard {

struct ConvGemmPlan {
  int64_t patchSize;
  int64_t numPatchesPerBatch;
  int64_t globalNumPatches;
  int64_t chunkStart;
  int64_t chunkNumPatches;
  int64_t maxParallelPixels;
  int64_t effectiveMaxParallelPixels;
  int64_t packedNumRows;

  RankedTensorType im2colType;
  RankedTensorType im2colRowType;
  RankedTensorType gemmInputRowsType;
  RankedTensorType wFlatType;
  RankedTensorType wTransType;
  RankedTensorType gemmOutType;
  RankedTensorType gemmOutputRowsType;
  RankedTensorType nhwcType;
};

static ConvGemmPlan
buildConvGemmPlan(const ConvLoweringState& state,
                  bool canPackWeightsAsConstants,
                  bool canPackBiasAsConstants,
                  int64_t chunkStart,
                  int64_t chunkNumPatches,
                  std::optional<int64_t> forcedPackFactor = std::nullopt);

static PreparedConvInput prepareInputForIm2Col(const ConvLoweringState& state,
                                               PatternRewriter& rewriter,
                                               Location loc) {
  if (state.padHeightBegin == 0 && state.padHeightEnd == 0 && state.padWidthBegin == 0 && state.padWidthEnd == 0)
    return {state.x, state.xType};

  auto paddedType = RankedTensorType::get({state.batchSize,
                                           state.numChannelsIn,
                                           state.xHeight + state.padHeightBegin + state.padHeightEnd,
                                           state.xWidth + state.padWidthBegin + state.padWidthEnd},
                                          state.xType.getElementType());
  auto paddedInputOp =
    createSpatCompute<1>(rewriter, loc, TypeRange {paddedType}, {}, state.x, [&](Value inputArg) {
      Value paddedInput = createZeroPaddedTensor(inputArg,
                                                 paddedType,
                                                 {0, 0, state.padHeightBegin, state.padWidthBegin},
                                                 {0, 0, state.padHeightEnd, state.padWidthEnd},
                                                 rewriter,
                                                 loc);
      spatial::SpatYieldOp::create(rewriter, loc, paddedInput);
    });
  return {paddedInputOp.getResult(0), paddedType};
}

static Value createPaddedRows(Value rows,
                              RankedTensorType rowsType,
                              int64_t paddedRows,
                              PatternRewriter& rewriter,
                              Location loc) {
  if (rowsType.getDimSize(0) == paddedRows)
    return rows;

  auto paddedType =
    RankedTensorType::get({paddedRows, rowsType.getDimSize(1)}, rowsType.getElementType(), rowsType.getEncoding());
  return createZeroPaddedTensor(
    rows, paddedType, {0, 0}, {paddedRows - rowsType.getDimSize(0), 0}, rewriter, loc);
}

static Value packRowsForParallelGemm(
  Value rows, RankedTensorType rowsType, int64_t packFactor, PatternRewriter& rewriter, Location loc) {
  if (packFactor == 1)
    return rows;

  const int64_t paddedNumRows = ceilIntegerDivide(rowsType.getDimSize(0), packFactor) * packFactor;
  const int64_t packedNumRows = paddedNumRows / packFactor;
  const int64_t rowWidth = rowsType.getDimSize(1);
  auto groupedType =
    RankedTensorType::get({packedNumRows, packFactor, rowWidth}, rowsType.getElementType(), rowsType.getEncoding());
  auto packedType =
    RankedTensorType::get({packedNumRows, packFactor * rowWidth}, rowsType.getElementType(), rowsType.getEncoding());

  Value padded = createPaddedRows(rows, rowsType, paddedNumRows, rewriter, loc);
  Value grouped = tensor::ExpandShapeOp::create(rewriter,
                                                loc,
                                                groupedType,
                                                padded,
                                                SmallVector<ReassociationIndices> {
                                                  {0, 1},
                                                  {2}
  });
  return tensor::CollapseShapeOp::create(rewriter,
                                         loc,
                                         packedType,
                                         grouped,
                                         SmallVector<ReassociationIndices> {
                                           {0},
                                           {1, 2}
  });
}

static Value unpackRowsFromParallelGemm(Value packedRows,
                                        RankedTensorType packedRowsType,
                                        int64_t unpackedRows,
                                        int64_t rowWidth,
                                        int64_t packFactor,
                                        PatternRewriter& rewriter,
                                        Location loc) {
  if (packFactor == 1)
    return packedRows;

  const int64_t packedNumRows = packedRowsType.getDimSize(0);
  const int64_t paddedNumRows = packedNumRows * packFactor;
  auto expandedType = RankedTensorType::get(
    {packedNumRows, packFactor, rowWidth}, packedRowsType.getElementType(), packedRowsType.getEncoding());
  auto paddedType =
    RankedTensorType::get({paddedNumRows, rowWidth}, packedRowsType.getElementType(), packedRowsType.getEncoding());
  auto unpackedType =
    RankedTensorType::get({unpackedRows, rowWidth}, packedRowsType.getElementType(), packedRowsType.getEncoding());

  Value expanded = tensor::ExpandShapeOp::create(rewriter,
                                                 loc,
                                                 expandedType,
                                                 packedRows,
                                                 SmallVector<ReassociationIndices> {
                                                   {0},
                                                   {1, 2}
  });
  Value padded = tensor::CollapseShapeOp::create(rewriter,
                                                 loc,
                                                 paddedType,
                                                 expanded,
                                                 SmallVector<ReassociationIndices> {
                                                   {0, 1},
                                                   {2}
  });
  if (paddedNumRows == unpackedRows)
    return padded;

  SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(unpackedRows), rewriter.getIndexAttr(rowWidth)};
  return tensor::ExtractSliceOp::create(rewriter, loc, unpackedType, padded, offsets, sizes, getUnitStrides(rewriter, 2));
}

static Value createWeightMatrix(
  Value weights, const ConvGemmPlan& plan, PatternRewriter& rewriter, Location loc) {
  auto buildWeightMatrix = [&](Value weight) -> Value {
    Value flattened = tensor::CollapseShapeOp::create(rewriter,
                                                      loc,
                                                      plan.wFlatType,
                                                      weight,
                                                      SmallVector<ReassociationIndices> {
                                                        {0},
                                                        {1, 2, 3}
    });
    return ONNXTransposeOp::create(rewriter, loc, plan.wTransType, flattened, rewriter.getI64ArrayAttr({1, 0}))
      .getResult();
  };

  if (isCompileTimeComputable(weights))
    return buildWeightMatrix(weights);

  auto computeOp =
    createSpatCompute<1>(rewriter, loc, TypeRange {plan.wTransType}, {}, ValueRange {weights}, [&](Value weight) {
      spatial::SpatYieldOp::create(rewriter, loc, buildWeightMatrix(weight));
    });
  return computeOp.getResult(0);
}

static Value createPaddedConvMatrix(Value matrix,
                                    RankedTensorType sourceType,
                                    RankedTensorType paddedType,
                                    PatternRewriter& rewriter,
                                    Location loc) {
  if (sourceType == paddedType)
    return matrix;
  return createZeroPaddedTensor(matrix,
                                paddedType,
                                {0, 0},
                                {paddedType.getDimSize(0) - sourceType.getDimSize(0),
                                 paddedType.getDimSize(1) - sourceType.getDimSize(1)},
                                rewriter,
                                loc);
}

static Value createPaddedConstantMatrix(DenseElementsAttr sourceAttr,
                                        RankedTensorType sourceType,
                                        RankedTensorType paddedType,
                                        PatternRewriter& rewriter) {
  SmallVector<Attribute> paddedValues(
    paddedType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(paddedType.getElementType())));
  SmallVector<Attribute> sourceValues(sourceAttr.getValues<Attribute>());
  const int64_t sourceRows = sourceType.getDimSize(0);
  const int64_t sourceCols = sourceType.getDimSize(1);
  const int64_t paddedCols = paddedType.getDimSize(1);
  for (int64_t row = 0; row < sourceRows; ++row)
    for (int64_t col = 0; col < sourceCols; ++col)
      paddedValues[row * paddedCols + col] = sourceValues[row * sourceCols + col];
  auto paddedAttr = DenseElementsAttr::get(paddedType, paddedValues);
  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), paddedAttr, paddedType);
}

static Value createPaddedInputKTiledWeightConstant(DenseElementsAttr sourceAttr,
                                                   const ConvLoweringState& state,
                                                   int64_t paddedK,
                                                   int64_t paddedC,
                                                   PatternRewriter& rewriter) {
  auto paddedType = RankedTensorType::get({paddedK, paddedC}, state.wType.getElementType());
  SmallVector<Attribute> sourceValues(sourceAttr.getValues<Attribute>());
  SmallVector<Attribute> paddedValues(
    paddedType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(paddedType.getElementType())));
  for (int64_t outChannel = 0; outChannel < state.numChannelsOut; ++outChannel) {
    for (int64_t inChannel = 0; inChannel < state.numChannelsIn; ++inChannel) {
      for (int64_t kernelH = 0; kernelH < state.wHeight; ++kernelH) {
        for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW) {
          const int64_t sourceFlatIndex =
            (((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW;
          const int64_t patchIndex = ((inChannel * state.wHeight) + kernelH) * state.wWidth + kernelW;
          paddedValues[patchIndex * paddedC + outChannel] = sourceValues[sourceFlatIndex];
        }
      }
    }
  }
  auto paddedAttr = DenseElementsAttr::get(paddedType, paddedValues);
  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), paddedAttr, paddedType);
}

static FailureOr<Value> rewriteInputKTiledConv(const ConvLoweringState& state,
                                               ArrayRef<DistributedTensorStep> distributedConsumers,
                                               PatternRewriter& rewriter,
                                               Location loc) {
  PreparedConvInput preparedInput = prepareInputForIm2Col(state, rewriter, loc);
  ConvGeometry geo = buildConvGeometry(state);
  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>(std::max<int64_t>(1, numKSlices * 4)));
  const uint64_t rowChunkWidth = std::max<uint64_t>(
    1,
    std::min<uint64_t>({chooseStreamChunkPositions(geo, /*packFactor=*/1),
                        maxLanesPerBatch,
                        static_cast<uint64_t>(state.outWidth)}));
  const auto elementType = state.outType.getElementType();
  auto wDenseAttr = getHostConstDenseElementsAttr(state.w);
  if (!wDenseAttr)
    return failure();

  Value paddedWeight = createPaddedInputKTiledWeightConstant(wDenseAttr, state, paddedK, xbarDim, rewriter);

  Value paddedBias;
  RankedTensorType paddedBiasType;
  if (state.hasBias) {
    Value biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc);
    auto biasMatrixType = cast<RankedTensorType>(biasMatrix.getType());
    paddedBiasType = RankedTensorType::get({1, xbarDim}, elementType);
    if (auto biasDenseAttr = getHostConstDenseElementsAttr(state.b))
      paddedBias = createPaddedConstantMatrix(biasDenseAttr, biasMatrixType, paddedBiasType, rewriter);
    else
      paddedBias = materializeOrComputeUnary(
        biasMatrix, paddedBiasType, rewriter, loc, [&](Value biasValue) {
          return createPaddedConvMatrix(biasValue, biasMatrixType, paddedBiasType, rewriter, loc);
        });
  }

  SmallVector<Value> chunkRows;
  const int64_t totalPatches = state.batchSize * state.outHeight * state.outWidth;
  chunkRows.reserve(
    state.batchSize * state.outHeight * ceilIntegerDivide(state.outWidth, static_cast<int64_t>(rowChunkWidth)));
  for (int64_t batchIndex = 0; batchIndex < state.batchSize; ++batchIndex) {
    for (int64_t outHeightIndex = 0; outHeightIndex < state.outHeight; ++outHeightIndex) {
      for (int64_t outWidthChunkStart = 0; outWidthChunkStart < state.outWidth;
           outWidthChunkStart += static_cast<int64_t>(rowChunkWidth)) {
        const int64_t chunkNumPatches =
          std::min<int64_t>(static_cast<int64_t>(rowChunkWidth), state.outWidth - outWidthChunkStart);
        auto chunkRowsType = RankedTensorType::get({chunkNumPatches, state.numChannelsOut}, elementType);
        auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
        auto paddedChunkRowType = RankedTensorType::get({1, paddedK}, elementType);
        auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elementType);
        auto collapsedPatchType = RankedTensorType::get({1, geo.k}, elementType);
        auto weightTileType = RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType());
        auto rowType = RankedTensorType::get({1, state.numChannelsOut}, elementType);
        SmallVector<Value> inputsStorage {preparedInput.value};
        if (state.hasBias)
          inputsStorage.push_back(paddedBias);
        ValueRange inputs(inputsStorage);
        auto chunkCompute = spatial::SpatCompute::create(rewriter, loc, TypeRange {chunkRowsType}, ValueRange {paddedWeight}, inputs);
        auto* block = new Block();
        block->addArgument(paddedWeight.getType(), loc);
        for (Value input : inputs)
          block->addArgument(input.getType(), loc);
        chunkCompute.getBody().push_back(block);
        rewriter.setInsertionPointToStart(block);

        auto buildChunk = [&]() -> LogicalResult {
          Value weightArg = block->getArgument(0);
          Value inputArg = block->getArgument(1);
          Value biasArg = state.hasBias ? block->getArgument(2) : Value();
          Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
          Value cBatchIndex = getOrCreateIndexConstant(rewriter, anchorOp, batchIndex);
          Value cZero = getOrCreateIndexConstant(rewriter, anchorOp, 0);
          Value cKSlices = getOrCreateIndexConstant(rewriter, anchorOp, numKSlices);
          Value cOne = getOrCreateIndexConstant(rewriter, anchorOp, 1);
          Value cXbar = getOrCreateIndexConstant(rewriter, anchorOp, xbarDim);
          Value cInputHeightOffset =
            getOrCreateIndexConstant(rewriter, anchorOp, outHeightIndex * state.strideHeight);
          Value chunkRowsValue = tensor::EmptyOp::create(rewriter, loc, chunkRowsType.getShape(), elementType);

          auto widthLoop = buildNormalizedScfFor(
            rewriter,
            loc,
            cZero,
            getOrCreateIndexConstant(rewriter, anchorOp, chunkNumPatches),
            cOne,
            ValueRange {chunkRowsValue},
            [&](OpBuilder&, Location nestedLoc, Value widthIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
              Value laneWithChunkOffset = affineAddConst(rewriter, nestedLoc, widthIndex, outWidthChunkStart, anchorOp);
              Value inputWidthOffset = createOrFoldAffineApply(rewriter,
                                                               nestedLoc,
                                                               getAffineDimExpr(0, rewriter.getContext()) * state.strideWidth,
                                                               ValueRange {laneWithChunkOffset},
                                                               anchorOp);
              Value patch = createConvInputPatch(inputArg,
                                                 patchType,
                                                 cBatchIndex,
                                                 cZero,
                                                 cInputHeightOffset,
                                                 inputWidthOffset,
                                                 state.dilationHeight,
                                                 state.dilationWidth,
                                                 rewriter,
                                                 nestedLoc);
              Value patchRow = tensor::CollapseShapeOp::create(rewriter,
                                                               nestedLoc,
                                                               collapsedPatchType,
                                                               patch,
                                                               SmallVector<ReassociationIndices> {
                                                                 {0},
                                                                 {1, 2, 3}
              });
              Value paddedPatchRow = createZeroPaddedTensor(
                patchRow, paddedChunkRowType, {0, 0}, {0, paddedK - geo.k}, rewriter, nestedLoc);

              auto zeroAttr = DenseElementsAttr::get(paddedRowType, rewriter.getZeroAttr(elementType));
              Value zeroRow = getOrCreateConstant(rewriter, anchorOp, zeroAttr, paddedRowType);
              auto kLoop = buildNormalizedScfFor(
                rewriter,
                nestedLoc,
                cZero,
                cKSlices,
                cOne,
                ValueRange {zeroRow},
                [&](OpBuilder&, Location reduceLoc, Value kSlice, ValueRange reduceIterArgs, SmallVectorImpl<Value>& reduceYielded) {
                  Value acc = reduceIterArgs.front();
                  Value kOffset = arith::MulIOp::create(rewriter, reduceLoc, kSlice, cXbar);
                  SmallVector<OpFoldResult> aOffsets {rewriter.getIndexAttr(0), kOffset};
                  SmallVector<OpFoldResult> aSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)};
                  SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, 2);
                  Value aTile = tensor::ExtractSliceOp::create(
                    rewriter, reduceLoc, paddedRowType, paddedPatchRow, aOffsets, aSizes, unitStrides);
                  SmallVector<OpFoldResult> bOffsets {kOffset, rewriter.getIndexAttr(0)};
                  SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)};
                  Value bTile = extractStaticSliceOrIdentity(
                    rewriter, reduceLoc, weightArg, weightTileType, bOffsets, bSizes, unitStrides);
                  Value piece = spatial::SpatVMMOp::create(rewriter, reduceLoc, paddedRowType, bTile, aTile).getResult();
                  reduceYielded.push_back(
                    spatial::SpatVAddOp::create(rewriter, reduceLoc, paddedRowType, acc, piece).getResult());
                  return success();
                });
              if (failed(kLoop))
                return failure();

              Value reduced = kLoop->results.front();
              if (state.hasBias)
                reduced = spatial::SpatVAddOp::create(rewriter, nestedLoc, paddedRowType, reduced, biasArg).getResult();

              Value row = reduced;
              if (state.numChannelsOut != xbarDim) {
                SmallVector<OpFoldResult> rowOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
                SmallVector<OpFoldResult> rowSizes {
                  rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
                row = tensor::ExtractSliceOp::create(
                  rewriter, nestedLoc, rowType, reduced, rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
              }

              SmallVector<OpFoldResult> outputOffsets {widthIndex, rewriter.getIndexAttr(0)};
              SmallVector<OpFoldResult> outputSizes {
                rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
              Value updatedRows = tensor::InsertSliceOp::create(
                rewriter, nestedLoc, row, iterArgs.front(), outputOffsets, outputSizes, getUnitStrides(rewriter, 2));
              yielded.push_back(updatedRows);
              return success();
            });
          if (failed(widthLoop))
            return failure();
          spatial::SpatYieldOp::create(rewriter, loc, widthLoop->results.front());
          return success();
        };
        if (failed(buildChunk())) {
          rewriter.setInsertionPointAfter(chunkCompute);
          rewriter.eraseOp(chunkCompute);
          return failure();
        }
        rewriter.setInsertionPointAfter(chunkCompute);
        chunkRows.push_back(chunkCompute.getResult(0));
      }
    }
  }

  auto nhwcType = RankedTensorType::get({state.batchSize, state.outHeight, state.outWidth, state.numChannelsOut},
                                        elementType);
  return createCollectedConvOutput(
    chunkRows, state.outType, cast<RankedTensorType>(chunkRows.front().getType()), nhwcType, state.outType, totalPatches,
    state.numChannelsOut, /*packFactor=*/1, distributedConsumers, rewriter, loc);
}

static Value buildPackedWeights(DenseElementsAttr wDenseAttr,
                                Value wTrans,
                                const ConvLoweringState& state,
                                const ConvGemmPlan& plan,
                                PatternRewriter& rewriter,
                                Location loc) {
  if (plan.effectiveMaxParallelPixels == 1)
    return wTrans;

  auto packedWeightType = RankedTensorType::get(
    {plan.effectiveMaxParallelPixels * plan.patchSize, plan.effectiveMaxParallelPixels * state.numChannelsOut},
    state.wType.getElementType());
  SmallVector<Attribute> sourceValues(wDenseAttr.getValues<Attribute>());
  SmallVector<Attribute> packedValues(packedWeightType.getNumElements(),
                                      cast<Attribute>(rewriter.getZeroAttr(state.wType.getElementType())));

  for (int64_t copyId = 0; copyId < plan.effectiveMaxParallelPixels; ++copyId) {
    for (int64_t outChannel = 0; outChannel < state.numChannelsOut; ++outChannel) {
      for (int64_t inChannel = 0; inChannel < state.numChannelsIn; ++inChannel) {
        for (int64_t kernelH = 0; kernelH < state.wHeight; ++kernelH) {
          for (int64_t kernelW = 0; kernelW < state.wWidth; ++kernelW) {
            const int64_t sourceFlatIndex =
              (((outChannel * state.numChannelsIn) + inChannel) * state.wHeight + kernelH) * state.wWidth + kernelW;
            const int64_t patchIndex = ((inChannel * state.wHeight) + kernelH) * state.wWidth + kernelW;
            const int64_t targetRow = copyId * plan.patchSize + patchIndex;
            const int64_t targetCol = copyId * state.numChannelsOut + outChannel;
            packedValues[targetRow * packedWeightType.getDimSize(1) + targetCol] = sourceValues[sourceFlatIndex];
          }
        }
      }
    }
  }

  auto packedAttr = DenseElementsAttr::get(packedWeightType, packedValues);
  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), packedAttr, packedWeightType);
}

static Value buildPackedBias(Value gemmBias,
                             Value biasMatrix,
                             DenseElementsAttr biasDenseAttr,
                             const ConvLoweringState& state,
                             const ConvGemmPlan& plan,
                             PatternRewriter& rewriter,
                             Location loc) {
  if (!state.hasBias)
    return gemmBias;

  if (plan.effectiveMaxParallelPixels == 1)
    return biasMatrix;

  SmallVector<Attribute> sourceValues(biasDenseAttr.getValues<Attribute>());
  SmallVector<Attribute> packedValues;
  packedValues.reserve(plan.effectiveMaxParallelPixels * state.numChannelsOut);
  for (int64_t copyId = 0; copyId < plan.effectiveMaxParallelPixels; ++copyId)
    packedValues.append(sourceValues.begin(), sourceValues.end());

  auto packedBiasType =
    RankedTensorType::get({1, plan.effectiveMaxParallelPixels * state.numChannelsOut}, state.outType.getElementType());
  auto packedBiasAttr = DenseElementsAttr::get(packedBiasType, packedValues);
  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), packedBiasAttr, packedBiasType);
}

static ConvGemmPlan
buildConvGemmPlan(const ConvLoweringState& state,
                  bool canPackWeightsAsConstants,
                  bool canPackBiasAsConstants,
                  int64_t chunkStart,
                  int64_t chunkNumPatches,
                  std::optional<int64_t> forcedPackFactor) {
  ConvGemmPlan plan;
  plan.patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
  plan.numPatchesPerBatch = state.outHeight * state.outWidth;
  plan.globalNumPatches = state.batchSize * plan.numPatchesPerBatch;
  plan.chunkStart = chunkStart;
  plan.chunkNumPatches = chunkNumPatches;
  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);
  plan.effectiveMaxParallelPixels =
    (canPackWeightsAsConstants && canPackBiasAsConstants) ? plan.maxParallelPixels : 1;
  plan.packedNumRows = ceilIntegerDivide(plan.chunkNumPatches, plan.effectiveMaxParallelPixels);

  auto elemType = state.xType.getElementType();
  auto outElemType = state.outType.getElementType();
  plan.im2colType = RankedTensorType::get({plan.chunkNumPatches, plan.patchSize}, elemType);
  plan.im2colRowType = RankedTensorType::get({1, plan.patchSize}, elemType);
  plan.gemmInputRowsType =
    RankedTensorType::get({plan.packedNumRows, plan.effectiveMaxParallelPixels * plan.patchSize}, elemType);
  plan.wFlatType = RankedTensorType::get({state.numChannelsOut, plan.patchSize}, state.wType.getElementType());
  plan.wTransType = RankedTensorType::get({plan.patchSize, state.numChannelsOut}, state.wType.getElementType());
  plan.gemmOutType = RankedTensorType::get({plan.chunkNumPatches, state.numChannelsOut}, outElemType);
  plan.gemmOutputRowsType =
    RankedTensorType::get({plan.packedNumRows, plan.effectiveMaxParallelPixels * state.numChannelsOut}, outElemType);
  plan.nhwcType =
    RankedTensorType::get({state.batchSize, state.outHeight, state.outWidth, state.numChannelsOut}, outElemType);
  return plan;
}

static Value createIm2colRows(const ConvLoweringState& state,
                              const PreparedConvInput& preparedInput,
                              const ConvGemmPlan& plan,
                              PatternRewriter& rewriter,
                              Location loc) {
  constexpr size_t numInputs = 1;
  auto im2colComputeOp =
    createSpatCompute<numInputs>(rewriter, loc, TypeRange {plan.gemmInputRowsType}, {}, preparedInput.value, [&](Value xArg) {
      auto elemType = preparedInput.type.getElementType();
      // Keep the standard im2col view of convolution, flipped so filters sit in
      // B / crossbar columns:
      //   A (im2col):  [numPatches, patchSize]  -- one row per output spatial position
      //   B (weights): [patchSize, cOut]
      //   Gemm output: [numPatches, cOut]
      Value im2colInit = tensor::EmptyOp::create(rewriter, loc, plan.im2colType.getShape(), elemType);
      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
      Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
      Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
      Value cNumPatches = getOrCreateIndexConstant(rewriter, anchorOp, plan.chunkNumPatches);

      auto im2colLoop = buildNormalizedScfFor(
        rewriter,
        loc,
        c0,
        cNumPatches,
        c1,
        ValueRange {im2colInit},
        [&](OpBuilder&, Location nestedLoc, Value patchIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
          Value im2colAcc = iterArgs.front();
          Value batchIndex =
            affineAddFloorDivConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp);
          Value batchPatchIndex =
            affineAddModConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp);
          Value outHeightIndex = affineFloorDivConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp);
          Value outWidthIndex = affineModConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp);
          Value inputHeightOffset =
            affineMulConst(rewriter, nestedLoc, outHeightIndex, state.strideHeight, anchorOp);
          Value inputWidthOffset =
            affineMulConst(rewriter, nestedLoc, outWidthIndex, state.strideWidth, anchorOp);

          auto patchType =
            RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
          Value patch = createConvInputPatch(xArg,
                                             patchType,
                                             batchIndex,
                                             c0,
                                             inputHeightOffset,
                                             inputWidthOffset,
                                             state.dilationHeight,
                                             state.dilationWidth,
                                             rewriter,
                                             nestedLoc);
          Value row = tensor::CollapseShapeOp::create(rewriter,
                                                      nestedLoc,
                                                      plan.im2colRowType,
                                                      patch,
                                                      SmallVector<ReassociationIndices> {
                                                        {0},
                                                        {1, 2, 3}
          });

          SmallVector<OpFoldResult> rowOffsets {patchIndex, rewriter.getIndexAttr(0)};
          SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(plan.patchSize)};
          Value next = tensor::InsertSliceOp::create(
            rewriter, nestedLoc, row, im2colAcc, rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
          yielded.push_back(next);
          return success();
        });
      if (failed(im2colLoop))
        return failure();

      Value gemmInputRows = im2colLoop->results.front();
      // Pack N old im2col rows into one longer row so one GEMM can cover N
      // pixels in parallel. The corresponding packed weight matrix contains N
      // block-diagonal copies of W^T, and the packed output must be unpacked
      // back to one row per spatial patch.
      if (plan.effectiveMaxParallelPixels != 1)
        gemmInputRows = packRowsForParallelGemm(gemmInputRows, plan.im2colType, plan.effectiveMaxParallelPixels, rewriter, loc);

      spatial::SpatYieldOp::create(rewriter, loc, gemmInputRows);
      return success();
    });

  assert(succeeded(im2colComputeOp) && "Conv im2col compute construction must succeed");
  return im2colComputeOp->getResult(0);
}

static Value maybeUnpackChunkRows(Value gemmRows,
                                  const ConvGemmPlan& plan,
                                  PatternRewriter& rewriter,
                                  Location loc) {
  if (plan.effectiveMaxParallelPixels == 1)
    return gemmRows;
  auto unpackedType = RankedTensorType::get(
    {plan.chunkNumPatches, plan.gemmOutType.getDimSize(1)}, plan.gemmOutType.getElementType(), plan.gemmOutType.getEncoding());
  auto unpackCompute = createSpatCompute<1>(rewriter, loc, TypeRange {unpackedType}, {}, gemmRows, [&](Value rowsArg) {
    Value unpacked = unpackRowsFromParallelGemm(rowsArg,
                                                cast<RankedTensorType>(rowsArg.getType()),
                                                plan.chunkNumPatches,
                                                plan.gemmOutType.getDimSize(1),
                                                plan.effectiveMaxParallelPixels,
                                                rewriter,
                                                loc);
    spatial::SpatYieldOp::create(rewriter, loc, unpacked);
  });
  return unpackCompute.getResult(0);
}

static Value createChunkedConvRows(const ConvLoweringState& state,
                                   const PreparedConvInput& preparedInput,
                                   Value weightMatrix,
                                   Value biasMatrix,
                                   DenseElementsAttr wDenseAttr,
                                   DenseElementsAttr biasDenseAttr,
                                   int64_t forcedPackFactor,
                                   uint64_t chunkPositions,
                                   PatternRewriter& rewriter,
                                   Location loc) {
  SmallVector<Value> chunkRows;
  const int64_t totalPatches = state.batchSize * state.outHeight * state.outWidth;
  for (int64_t chunkStart = 0; chunkStart < totalPatches; chunkStart += static_cast<int64_t>(chunkPositions)) {
    const int64_t chunkNumPatches = std::min<int64_t>(static_cast<int64_t>(chunkPositions), totalPatches - chunkStart);
    ConvGemmPlan chunkPlan = buildConvGemmPlan(state,
                                               static_cast<bool>(wDenseAttr),
                                               !state.hasBias || static_cast<bool>(biasDenseAttr),
                                               chunkStart,
                                               chunkNumPatches,
                                               forcedPackFactor);
    Value chunkInputRows = createIm2colRows(state, preparedInput, chunkPlan, rewriter, loc);
    Value chunkB = buildPackedWeights(wDenseAttr, weightMatrix, state, chunkPlan, rewriter, loc);
    Value gemmBias = createZeroGemmBias(chunkPlan.gemmOutputRowsType, rewriter);
    if (state.hasBias)
      gemmBias = state.b;
    Value chunkC = buildPackedBias(gemmBias, biasMatrix, biasDenseAttr, state, chunkPlan, rewriter, loc);
    Value chunkGemmRows = ONNXGemmOp::create(rewriter,
                                             loc,
                                             chunkPlan.gemmOutputRowsType,
                                             chunkInputRows,
                                             chunkB,
                                             chunkC,
                                             APFloat(1.0f),
                                             APFloat(1.0f),
                                             /*transA=*/0,
                                             /*transB=*/0)
                           .getY();
    chunkRows.push_back(maybeUnpackChunkRows(chunkGemmRows, chunkPlan, rewriter, loc));
  }

  if (chunkRows.size() == 1)
    return chunkRows.front();

  auto rowType = RankedTensorType::get({totalPatches, state.numChannelsOut}, state.outType.getElementType());
  auto collectRows = createSpatCompute(rewriter, loc, TypeRange {rowType}, {}, chunkRows, [&](ValueRange rows) {
    spatial::SpatYieldOp::create(rewriter, loc, createSpatConcat(rewriter, loc, /*axis=*/0, rows));
  });
  return collectRows.getResult(0);
}

static Value rewritePackedIm2ColConv(const ConvLoweringState& state,
                                     ArrayRef<DistributedTensorStep> distributedConsumers,
                                     PatternRewriter& rewriter,
                                     Location loc) {
  auto wDenseAttr = getHostConstDenseElementsAttr(state.w);
  PreparedConvInput preparedInput = prepareInputForIm2Col(state, rewriter, loc);
  Value biasMatrix;
  DenseElementsAttr biasDenseAttr;
  if (state.hasBias) {
    biasDenseAttr = getHostConstDenseElementsAttr(state.b);
    biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc);
  }

  ConvGemmPlan plan =
    buildConvGemmPlan(state, static_cast<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(biasDenseAttr), 0,
                      state.batchSize * state.outHeight * state.outWidth);
  // Prepare weight matrix W for crossbar storage:
  //   W: [Cout, Cin, KH, KW] -> [Cout, patchSize] -> [patchSize, Cout]
  Value weightMatrix = createWeightMatrix(state.w, plan, rewriter, loc);
  Value gemmInputRows = createIm2colRows(state, preparedInput, plan, rewriter, loc);
  Value gemmB = buildPackedWeights(wDenseAttr, weightMatrix, state, plan, rewriter, loc);
  Value gemmBias = createZeroGemmBias(plan.gemmOutputRowsType, rewriter);
  if (state.hasBias)
    gemmBias = state.b;
  Value gemmC = buildPackedBias(gemmBias, biasMatrix, biasDenseAttr, state, plan, rewriter, loc);

  Value gemmRows = ONNXGemmOp::create(rewriter,
                                      loc,
                                      plan.gemmOutputRowsType,
                                      gemmInputRows,
                                      gemmB,
                                      gemmC,
                                      APFloat(1.0f),
                                      APFloat(1.0f),
                                      /*transA=*/0,
                                      /*transB=*/0)
                     .getY();

  return createCollectedConvOutput(ValueRange {gemmRows},
                                   state.outType,
                                   plan.gemmOutType,
                                   plan.nhwcType,
                                   state.outType,
                                   plan.chunkNumPatches,
                                   state.numChannelsOut,
                                   plan.effectiveMaxParallelPixels,
                                   distributedConsumers,
                                   rewriter,
                                   loc);
}

static Value rewriteStreamedConv(const ConvLoweringState& state,
                                 ArrayRef<DistributedTensorStep> distributedConsumers,
                                 PatternRewriter& rewriter,
                                 Location loc,
                                 int64_t forcedPackFactor) {
  auto wDenseAttr = getHostConstDenseElementsAttr(state.w);
  PreparedConvInput preparedInput = prepareInputForIm2Col(state, rewriter, loc);
  Value biasMatrix;
  DenseElementsAttr biasDenseAttr;
  if (state.hasBias) {
    biasDenseAttr = getHostConstDenseElementsAttr(state.b);
    biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc);
  }

  ConvGemmPlan seedPlan = buildConvGemmPlan(
    state, static_cast<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(biasDenseAttr), 0, 1, forcedPackFactor);
  Value weightMatrix = createWeightMatrix(state.w, seedPlan, rewriter, loc);
  ConvGeometry geo = buildConvGeometry(state);
  uint64_t chunkPositions = chooseStreamChunkPositions(geo, forcedPackFactor);
  Value collectedRows = createChunkedConvRows(state,
                                              preparedInput,
                                              weightMatrix,
                                              biasMatrix,
                                              wDenseAttr,
                                              biasDenseAttr,
                                              forcedPackFactor,
                                              chunkPositions,
                                              rewriter,
                                              loc);
  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),
    state.numChannelsOut, /*packFactor=*/1, distributedConsumers, rewriter, loc);
}

} // namespace standard

static RankedTensorType getRowStripFragmentType(RankedTensorType tensorType, int64_t width) {
  return RankedTensorType::get(
    {tensorType.getDimSize(0), tensorType.getDimSize(1), 1, width}, tensorType.getElementType(), tensorType.getEncoding());
}

static SmallVector<DistributedFragmentInfo, 8> buildRowStripFragments(RankedTensorType tensorType) {
  SmallVector<DistributedFragmentInfo, 8> fragments;
  const int64_t height = tensorType.getDimSize(2);
  const int64_t width = tensorType.getDimSize(3);
  const int64_t channels = tensorType.getDimSize(1);
  fragments.reserve(height);
  for (int64_t row = 0; row < height; ++row) {
    fragments.push_back(DistributedFragmentInfo {
      {0, 0, row, 0},
      {1, channels, 1, width},
      {1, 1, 1, 1},
      row,
    });
  }
  return fragments;
}

static DistributedTensorInfo makeDistributedTensorInfo(Value storage, RankedTensorType logicalType) {
  DistributedTensorInfo info;
  info.storage = storage;
  info.logicalType = logicalType;
  info.fragments = buildRowStripFragments(logicalType);
  info.laneCount = logicalType.getDimSize(2);
  info.channels = logicalType.getDimSize(1);
  info.height = logicalType.getDimSize(2);
  info.width = logicalType.getDimSize(3);
  return info;
}

static Value createPerChannelConstantFragment(DenseElementsAttr denseAttr,
                                              RankedTensorType fragmentType,
                                              PatternRewriter& rewriter) {
  auto denseType = cast<RankedTensorType>(denseAttr.getType());
  SmallVector<Attribute> channelValues;
  channelValues.reserve(fragmentType.getDimSize(1));
  SmallVector<Attribute> flattened(denseAttr.getValues<Attribute>());
  if (denseType.getRank() == 1) {
    channelValues = flattened;
  }
  else if (denseType.getRank() == 2) {
    channelValues = flattened;
  }
  else {
    for (int64_t channel = 0; channel < denseType.getDimSize(1); ++channel)
      channelValues.push_back(flattened[channel]);
  }

  SmallVector<Attribute> values;
  values.reserve(fragmentType.getNumElements());
  for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n)
    for (int64_t channel = 0; channel < fragmentType.getDimSize(1); ++channel)
      for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h)
        for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w)
          values.push_back(channelValues[channel]);

  auto attr = DenseElementsAttr::get(fragmentType, values);
  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), attr, fragmentType);
}

static Value createZeroGemmBias(RankedTensorType gemmResultType, PatternRewriter& rewriter) {
  auto zeroAttr = DenseElementsAttr::get(gemmResultType, rewriter.getZeroAttr(gemmResultType.getElementType()));
  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, gemmResultType);
}

static bool canDirectLowerRowStripConv(const ConvLoweringState& state, StringRef& failureReason) {
  if (!canConsumeRowStripHwcInput(state, failureReason))
    return false;

  ConvGeometry geometry = buildConvGeometry(state);
  if (state.numChannelsOut > geometry.xbarSize) {
    failureReason = "unsupported_output_channels";
    return false;
  }

  failureReason = "";
  return true;
}

static FailureOr<Value> createRowStripPackedRows(Value rows,
                                                 const ConvLoweringState& state,
                                                 PatternRewriter& rewriter,
                                                 Location loc) {
  auto rowsType = dyn_cast<RankedTensorType>(rows.getType());
  if (!rowsType || !rowsType.hasStaticShape() || rowsType.getRank() != 2)
    return failure();

  if (state.batchSize != 1)
    return failure();
  if (state.outType.getRank() != 4 || !state.outType.hasStaticShape())
    return failure();

  const int64_t outHeight = state.outType.getDimSize(2);
  const int64_t outWidth = state.outType.getDimSize(3);
  const int64_t outChannels = state.outType.getDimSize(1);
  if (rowsType.getDimSize(0) != outHeight * outWidth || rowsType.getDimSize(1) != outChannels)
    return failure();

  auto packedType = RankedTensorType::get({outHeight, outWidth, outChannels}, rowsType.getElementType(), rowsType.getEncoding());
  auto packedRows =
    createSpatCompute<1>(rewriter, loc, TypeRange {packedType}, {}, rows, [&](Value rowValues) {
      Value packed = tensor::ExpandShapeOp::create(
        rewriter, loc, packedType, rowValues, SmallVector<ReassociationIndices> {{0, 1}, {2}});
      spatial::SpatYieldOp::create(rewriter, loc, packed);
    });
  return packedRows.getResult(0);
}

static FailureOr<Value> createConvOutputFromRowStripHwc(Value inputHwc,
                                                        const ConvLoweringState& state,
                                                        PatternRewriter& rewriter,
                                                        Location loc) {
  auto inputType = dyn_cast<RankedTensorType>(inputHwc.getType());
  if (!inputType || !inputType.hasStaticShape() || inputType.getRank() != 3)
    return failure();
  if (inputType.getDimSize(0) != state.xHeight || inputType.getDimSize(1) != state.xWidth
      || inputType.getDimSize(2) != state.numChannelsIn)
    return failure();

  StringRef failureReason;
  if (!canDirectLowerRowStripConv(state, failureReason))
    return failure();

  ConvRowDemand demand = buildConvRowDemand(RowInterval {0, state.outHeight}, state);
  if (!covers(demand.acquiredInputRows, demand.neededInputRows))
    return failure();

  ConvGeometry geometry = buildConvGeometry(state);
  const int64_t xbarDim = geometry.xbarSize;
  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
  const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
  const int64_t paddedK = numKSlices * xbarDim;
  auto elementType = inputType.getElementType();
  auto paddedInputType = RankedTensorType::get({state.xHeight + state.padHeightBegin + state.padHeightEnd,
                                                state.xWidth + state.padWidthBegin + state.padWidthEnd,
                                                state.numChannelsIn},
                                               elementType,
                                               inputType.getEncoding());
  auto paddedPatchType =
    RankedTensorType::get({state.wHeight, state.wWidth, 1}, elementType, inputType.getEncoding());
  auto flatPatchType = RankedTensorType::get({state.wHeight * state.wWidth}, elementType, inputType.getEncoding());
  auto rowChunkType = RankedTensorType::get({1, state.wHeight * state.wWidth}, elementType, inputType.getEncoding());
  auto rowType = RankedTensorType::get({1, state.numChannelsOut}, state.outType.getElementType());
  auto packedOutputType =
    RankedTensorType::get({state.outHeight, state.outWidth, state.numChannelsOut}, state.outType.getElementType());
  auto packedOutputSliceType =
    RankedTensorType::get({1, 1, state.numChannelsOut}, state.outType.getElementType());
  auto paddedRowType = RankedTensorType::get({1, xbarDim}, state.outType.getElementType());
  auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType, inputType.getEncoding());
  auto paddedWeightTileType = RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType());
  auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
  if (!weightDenseAttr)
    return failure();
  Value paddedWeights = standard::createPaddedInputKTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);

  Value paddedBias;
  if (state.hasBias) {
    Value biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc);
    auto biasMatrixType = cast<RankedTensorType>(biasMatrix.getType());
    auto paddedBiasType = RankedTensorType::get({1, xbarDim}, state.outType.getElementType());
    if (auto biasDenseAttr = getHostConstDenseElementsAttr(state.b))
      paddedBias = standard::createPaddedConstantMatrix(biasDenseAttr, biasMatrixType, paddedBiasType, rewriter);
    else
      paddedBias = materializeOrComputeUnary(
        biasMatrix, paddedBiasType, rewriter, loc, [&](Value biasValue) {
          return standard::createPaddedConvMatrix(biasValue, biasMatrixType, paddedBiasType, rewriter, loc);
        });
  }

  auto paddedInputOp =
    createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, inputHwc, [&](Value hwcInputArg) {
      Value paddedInput = createZeroPaddedTensor(hwcInputArg,
                                                 paddedInputType,
                                                 {state.padHeightBegin, state.padWidthBegin, 0},
                                                 {state.padHeightEnd, state.padWidthEnd, 0},
                                                 rewriter,
                                                 loc);
      spatial::SpatYieldOp::create(rewriter, loc, paddedInput);
    });

  SmallVector<Value> batchInputs {paddedInputOp.getResult(0)};
  if (state.hasBias)
    batchInputs.push_back(paddedBias);
  auto batchOp = createSpatComputeBatch(
    rewriter,
    loc,
    TypeRange {packedOutputType},
    state.outHeight,
    ValueRange {paddedWeights},
    batchInputs,
    [&](detail::SpatComputeBatchBodyArgs args) {
      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
      Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
      Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
      Value cNumKSlices = getOrCreateIndexConstant(rewriter, anchorOp, numKSlices);
      Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth);
      Value cNumChannels = getOrCreateIndexConstant(rewriter, anchorOp, state.numChannelsIn);
      Value localHeightOffset = args.lane;
      Value packedRowInit =
        tensor::EmptyOp::create(rewriter, loc, ArrayRef<int64_t> {1, state.outWidth, state.numChannelsOut}, elementType);
      auto widthLoop = buildNormalizedScfFor(
        rewriter,
        loc,
        c0,
        cOutWidth,
        c1,
        ValueRange {packedRowInit},
        [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl<Value>& widthYielded) {
          Value localWidthOffset = widthIndex;
          Value rowInit = tensor::EmptyOp::create(rewriter, widthLoc, ArrayRef<int64_t> {1, patchSize}, elementType);
          auto rowLoop = buildNormalizedScfFor(
            rewriter,
            widthLoc,
            c0,
            cNumChannels,
            c1,
            ValueRange {rowInit},
            [&](OpBuilder&, Location rowLoc, Value channel, ValueRange rowIterArgs, SmallVectorImpl<Value>& rowYielded) {
              SmallVector<OpFoldResult> patchOffsets {localHeightOffset, localWidthOffset, channel};
              SmallVector<OpFoldResult> patchSizes {
                rewriter.getIndexAttr(state.wHeight), rewriter.getIndexAttr(state.wWidth), rewriter.getIndexAttr(1)};
              Value channelPatch = tensor::ExtractSliceOp::create(
                rewriter, rowLoc, paddedPatchType, args.inputs.front(), patchOffsets, patchSizes, getUnitStrides(rewriter, 3));
              Value flatPatch = tensor::CollapseShapeOp::create(
                rewriter, rowLoc, flatPatchType, channelPatch, SmallVector<ReassociationIndices> {{0, 1, 2}});
              Value rowChunk = tensor::ExpandShapeOp::create(
                rewriter, rowLoc, rowChunkType, flatPatch, SmallVector<ReassociationIndices> {{0, 1}});
              Value flatOffset = affineMulConst(
                rewriter, rowLoc, channel, state.wHeight * state.wWidth, anchorOp);
              SmallVector<OpFoldResult> rowOffsets {rewriter.getIndexAttr(0), flatOffset};
              SmallVector<OpFoldResult> rowSizes {
                rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.wHeight * state.wWidth)};
              Value nextRow = tensor::InsertSliceOp::create(
                rewriter, rowLoc, rowChunk, rowIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
              rowYielded.push_back(nextRow);
              return success();
            });
          if (failed(rowLoop))
            return failure();

          Value paddedRow = rowLoop->results.front();
          if (patchSize != paddedK)
            paddedRow = createZeroPaddedTensor(
              paddedRow, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, widthLoc);

          auto zeroAttr = DenseElementsAttr::get(paddedRowType, rewriter.getZeroAttr(state.outType.getElementType()));
          Value zeroRow = getOrCreateConstant(rewriter, anchorOp, zeroAttr, paddedRowType);
          auto kLoop = buildNormalizedScfFor(
            rewriter,
            widthLoc,
            c0,
            cNumKSlices,
            c1,
            ValueRange {zeroRow},
            [&](OpBuilder&, Location reduceLoc, Value kSlice, ValueRange reduceIterArgs, SmallVectorImpl<Value>& reduceYielded) {
              Value kOffset = affineMulConst(rewriter, reduceLoc, kSlice, xbarDim, anchorOp);
              SmallVector<OpFoldResult> aOffsets {rewriter.getIndexAttr(0), kOffset};
              SmallVector<OpFoldResult> aSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)};
              Value aTile = tensor::ExtractSliceOp::create(
                rewriter, reduceLoc, paddedRowType, paddedRow, aOffsets, aSizes, getUnitStrides(rewriter, 2));
              SmallVector<OpFoldResult> bOffsets {kOffset, rewriter.getIndexAttr(0)};
              SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)};
              Value bTile = extractStaticSliceOrIdentity(rewriter,
                                                         reduceLoc,
                                                         args.weights.front(),
                                                         paddedWeightTileType,
                                                         bOffsets,
                                                         bSizes,
                                                         getUnitStrides(rewriter, 2));
              Value piece = spatial::SpatVMMOp::create(rewriter, reduceLoc, paddedRowType, bTile, aTile).getResult();
              reduceYielded.push_back(
                spatial::SpatVAddOp::create(rewriter, reduceLoc, paddedRowType, reduceIterArgs.front(), piece).getResult());
              return success();
            });
          if (failed(kLoop))
            return failure();

          Value rowResult = kLoop->results.front();
          if (state.hasBias)
            rowResult =
              spatial::SpatVAddOp::create(rewriter, widthLoc, paddedRowType, rowResult, args.inputs[1]).getResult();

          Value outputRow = rowResult;
          if (state.numChannelsOut != xbarDim) {
            SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
            SmallVector<OpFoldResult> outputSizes {
              rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
            outputRow = tensor::ExtractSliceOp::create(
              rewriter, widthLoc, rowType, rowResult, outputOffsets, outputSizes, getUnitStrides(rewriter, 2));
          }

          Value outputFragment = tensor::ExpandShapeOp::create(rewriter,
                                                               widthLoc,
                                                               packedOutputSliceType,
                                                               outputRow,
                                                               SmallVector<ReassociationIndices> {{0}, {1, 2}});
          SmallVector<OpFoldResult> rowOffsets {rewriter.getIndexAttr(0), widthIndex, rewriter.getIndexAttr(0)};
          SmallVector<OpFoldResult> rowSizes {
            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
          Value nextPackedRow = tensor::InsertSliceOp::create(
            rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 3));
          widthYielded.push_back(nextPackedRow);
          return success();
        });
      if (failed(widthLoop))
        return failure();

      SmallVector<OpFoldResult> batchOffsets {args.lane, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
      SmallVector<OpFoldResult> batchSizes {
        rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.outWidth), rewriter.getIndexAttr(state.numChannelsOut)};
      createParallelInsertSliceIntoBatchOutput(
        rewriter, loc, widthLoop->results.front(), args.outputs.front(), batchOffsets, batchSizes, getUnitStrides(rewriter, 3));
      return success();
    });
  if (failed(batchOp))
    return failure();
  return batchOp->getResult(0);
}

static FailureOr<Value> createConvRowsFromRowStripInput(const ConvLoweringState& state,
                                                        [[maybe_unused]] const ConvLoweringDecision& decision,
                                                        Value rowStripInput,
                                                        PatternRewriter& rewriter,
                                                        Location loc) {
  return createConvOutputFromRowStripHwc(rowStripInput, state, rewriter, loc);
}

static Value createFragmentConstant(const DistributedTensorStep& step,
                                    RankedTensorType fragmentType,
                                    PatternRewriter& rewriter) {
  if (step.constantKind == DistributedTensorConstantKind::PerChannel)
    return createPerChannelConstantFragment(step.constantAttr, fragmentType, rewriter);

  Attribute splatValue = step.constantAttr.getSplatValue<Attribute>();
  return getOrCreateConstant(rewriter,
                             rewriter.getInsertionBlock()->getParentOp(),
                             DenseElementsAttr::get(fragmentType, splatValue),
                             fragmentType);
}

static Value createFragmentReciprocalConstant(const DistributedTensorStep& step,
                                              RankedTensorType fragmentType,
                                              PatternRewriter& rewriter) {
  SmallVector<APFloat> values;
  if (step.constantKind == DistributedTensorConstantKind::PerChannel) {
    auto denseType = cast<RankedTensorType>(step.constantAttr.getType());
    SmallVector<APFloat> channelValues;
    for (const APFloat& value : step.constantAttr.getValues<APFloat>())
      channelValues.push_back(value);
    values.reserve(fragmentType.getNumElements());
    for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n)
      for (int64_t channel = 0; channel < fragmentType.getDimSize(1); ++channel)
        for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h)
          for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w) {
            APFloat reciprocal = channelValues[channel];
            APFloat one(reciprocal.getSemantics(), 1);
            [[maybe_unused]] APFloat::opStatus status = one.divide(reciprocal, APFloat::rmNearestTiesToEven);
            assert(!(status & APFloat::opInvalidOp) && "distributed conv div requires finite non-zero constant");
            values.push_back(one);
          }
    (void)denseType;
  }
  else {
    APFloat reciprocal = cast<DenseFPElementsAttr>(step.constantAttr).getSplatValue<APFloat>();
    APFloat one(reciprocal.getSemantics(), 1);
    [[maybe_unused]] APFloat::opStatus status = one.divide(reciprocal, APFloat::rmNearestTiesToEven);
    assert(!(status & APFloat::opInvalidOp) && "distributed conv div requires finite non-zero constant");
    values.assign(fragmentType.getNumElements(), one);
  }
  return getOrCreateConstant(rewriter,
                             rewriter.getInsertionBlock()->getParentOp(),
                             DenseFPElementsAttr::get(fragmentType, values),
                             fragmentType);
}

[[maybe_unused]] static FailureOr<Value> createConvRowsForStrategy(const ConvLoweringState& state,
                                                                   const ConvLoweringDecision& decision,
                                                                   PatternRewriter& rewriter,
                                                                   Location loc) {
  auto wDenseAttr = getHostConstDenseElementsAttr(state.w);
  PreparedConvInput preparedInput = standard::prepareInputForIm2Col(state, rewriter, loc);
  Value biasMatrix;
  DenseElementsAttr biasDenseAttr;
  if (state.hasBias) {
    biasDenseAttr = getHostConstDenseElementsAttr(state.b);
    biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc);
  }

  switch (decision.strategy) {
  case PimConvLoweringLegacy:
  case PimConvLoweringPackedIm2Col: {
    standard::ConvGemmPlan plan = standard::buildConvGemmPlan(
      state, static_cast<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(biasDenseAttr), 0,
      state.batchSize * state.outHeight * state.outWidth);
    Value weightMatrix = standard::createWeightMatrix(state.w, plan, rewriter, loc);
    Value gemmInputRows = standard::createIm2colRows(state, preparedInput, plan, rewriter, loc);
    Value gemmB = standard::buildPackedWeights(wDenseAttr, weightMatrix, state, plan, rewriter, loc);
    Value gemmBias = createZeroGemmBias(plan.gemmOutputRowsType, rewriter);
    if (state.hasBias)
      gemmBias = state.b;
    Value gemmC = standard::buildPackedBias(gemmBias, biasMatrix, biasDenseAttr, state, plan, rewriter, loc);
    Value gemmRows = ONNXGemmOp::create(rewriter,
                                        loc,
                                        plan.gemmOutputRowsType,
                                        gemmInputRows,
                                        gemmB,
                                        gemmC,
                                        APFloat(1.0f),
                                        APFloat(1.0f),
                                        /*transA=*/0,
                                        /*transB=*/0)
                       .getY();
    return standard::maybeUnpackChunkRows(gemmRows, plan, rewriter, loc);
  }
  case PimConvLoweringStreamedPatch:
  case PimConvLoweringOutputChannelTiled:
  case PimConvLoweringTiled2D:
  case PimConvLoweringStreamedPacked: {
    standard::ConvGemmPlan seedPlan = standard::buildConvGemmPlan(
      state, static_cast<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(biasDenseAttr), 0, 1,
      decision.strategy == PimConvLoweringStreamedPacked ? buildConvGeometry(state).pack : 1);
    Value weightMatrix = standard::createWeightMatrix(state.w, seedPlan, rewriter, loc);
    ConvGeometry geo = buildConvGeometry(state);
    int64_t packFactor = decision.strategy == PimConvLoweringStreamedPacked ? geo.pack : 1;
    uint64_t chunkPositions = chooseStreamChunkPositions(geo, packFactor);
    return standard::createChunkedConvRows(state,
                                           preparedInput,
                                           weightMatrix,
                                           biasMatrix,
                                           wDenseAttr,
                                           biasDenseAttr,
                                           packFactor,
                                           chunkPositions,
                                           rewriter,
                                           loc);
  }
  default:
    return failure();
  }
}

[[maybe_unused]] static FailureOr<DistributedTensorInfo> createDistributedTensorFromRows(Value rows,
                                                                                         RankedTensorType logicalType,
                                                                                         PatternRewriter& rewriter,
                                                                                         Location loc) {
  const int64_t width = logicalType.getDimSize(3);
  const int64_t height = logicalType.getDimSize(2);
  auto rowsType = cast<RankedTensorType>(rows.getType());
  auto rowSliceType =
    RankedTensorType::get({width, logicalType.getDimSize(1)}, logicalType.getElementType(), rowsType.getEncoding());
  auto channelWidthType =
    RankedTensorType::get({logicalType.getDimSize(1), width}, logicalType.getElementType(), rowsType.getEncoding());
  auto fragmentType = getRowStripFragmentType(logicalType, width);
  auto batchOp = createSpatComputeBatch(
    rewriter, loc, TypeRange {logicalType}, height, {}, ValueRange {rows}, [&](detail::SpatComputeBatchBodyArgs args) {
      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
      Value rowStart = affineMulConst(rewriter, loc, args.lane, width, anchorOp);
      SmallVector<OpFoldResult> rowOffsets {rowStart, rewriter.getIndexAttr(0)};
      SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(width), rewriter.getIndexAttr(logicalType.getDimSize(1))};
      Value rowSlice = tensor::ExtractSliceOp::create(
        rewriter, loc, rowSliceType, args.inputs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
      Value channelWidth = ONNXTransposeOp::create(
        rewriter, loc, channelWidthType, rowSlice, rewriter.getI64ArrayAttr({1, 0})).getResult();
      Value fragment = tensor::ExpandShapeOp::create(rewriter,
                                                     loc,
                                                     fragmentType,
                                                     channelWidth,
                                                     SmallVector<ReassociationIndices> {{0, 1}, {2, 3}});
      SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), args.lane,
                                               rewriter.getIndexAttr(0)};
      SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(logicalType.getDimSize(1)),
                                             rewriter.getIndexAttr(1), rewriter.getIndexAttr(width)};
      createParallelInsertSliceIntoBatchOutput(
        rewriter, loc, fragment, args.outputs.front(), outputOffsets, outputSizes, getUnitStrides(rewriter, 4));
      return success();
    });
  if (failed(batchOp))
    return failure();
  return makeDistributedTensorInfo(batchOp->getResult(0), logicalType);
}

[[maybe_unused]] static FailureOr<DistributedTensorInfo> applyDistributedPreservingStep(const DistributedTensorInfo& inputInfo,
                                                                                        const DistributedTensorStep& step,
                                                                                        PatternRewriter& rewriter,
                                                                                        Location loc) {
  auto logicalType = inputInfo.logicalType;
  const int64_t width = logicalType.getDimSize(3);
  auto fragmentType = getRowStripFragmentType(logicalType, width);
  auto batchOp = createSpatComputeBatch(rewriter,
                                        loc,
                                        TypeRange {logicalType},
                                        inputInfo.laneCount,
                                        {},
                                        ValueRange {inputInfo.storage},
                                        [&](detail::SpatComputeBatchBodyArgs args) {
                                          SmallVector<OpFoldResult> offsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0),
                                                                             args.lane, rewriter.getIndexAttr(0)};
                                          SmallVector<OpFoldResult> sizes {
                                            rewriter.getIndexAttr(1), rewriter.getIndexAttr(logicalType.getDimSize(1)),
                                            rewriter.getIndexAttr(1), rewriter.getIndexAttr(width)};
                                          Value fragment = tensor::ExtractSliceOp::create(
                                            rewriter, loc, fragmentType, args.inputs.front(), offsets, sizes, getUnitStrides(rewriter, 4));
                                          switch (step.kind) {
                                          case DistributedTensorOpKind::Relu:
                                            fragment = spatial::SpatReluOp::create(rewriter, loc, fragmentType, fragment).getResult();
                                            break;
                                          case DistributedTensorOpKind::Sigmoid:
                                            fragment = spatial::SpatSigmoidOp::create(rewriter, loc, fragmentType, fragment).getResult();
                                            break;
                                          case DistributedTensorOpKind::Add: {
                                            Value constant = createFragmentConstant(step, fragmentType, rewriter);
                                            fragment =
                                              spatial::SpatVAddOp::create(rewriter, loc, fragmentType, fragment, constant).getResult();
                                            break;
                                          }
                                          case DistributedTensorOpKind::Sub: {
                                            Value constant = createFragmentConstant(step, fragmentType, rewriter);
                                            Value lhs = step.fragmentOnLhs ? fragment : constant;
                                            Value rhs = step.fragmentOnLhs ? constant : fragment;
                                            fragment = spatial::SpatVSubOp::create(rewriter, loc, fragmentType, lhs, rhs).getResult();
                                            break;
                                          }
                                          case DistributedTensorOpKind::Mul: {
                                            Value constant = createFragmentConstant(step, fragmentType, rewriter);
                                            fragment =
                                              spatial::SpatVMulOp::create(rewriter, loc, fragmentType, fragment, constant).getResult();
                                            break;
                                          }
                                          case DistributedTensorOpKind::Div: {
                                            Value constant = createFragmentReciprocalConstant(step, fragmentType, rewriter);
                                            fragment =
                                              spatial::SpatVMulOp::create(rewriter, loc, fragmentType, fragment, constant).getResult();
                                            break;
                                          }
                                          case DistributedTensorOpKind::Conv:
                                            return failure();
                                          }
                                          createParallelInsertSliceIntoBatchOutput(
                                            rewriter, loc, fragment, args.outputs.front(), offsets, sizes, getUnitStrides(rewriter, 4));
                                          return success();
                                        });
  if (failed(batchOp))
    return failure();
  return makeDistributedTensorInfo(batchOp->getResult(0), logicalType);
}

static Value createCollectedConvOutput(ValueRange gemmRows,
                                       Type convType,
                                       RankedTensorType gemmOutType,
                                       RankedTensorType nhwcType,
                                       RankedTensorType outType,
                                       int64_t numPatches,
                                       int64_t numChannelsOut,
                                       int64_t packFactor,
                                       ArrayRef<DistributedTensorStep> distributedConsumers,
                                       PatternRewriter& rewriter,
                                       Location loc) {
  auto materializeSplatTensor = [&](DenseElementsAttr denseAttr, RankedTensorType targetType) {
    Attribute splatValue = denseAttr.getSplatValue<Attribute>();
    auto targetAttr = DenseElementsAttr::get(targetType, splatValue);
    return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), targetAttr, targetType);
  };

  auto materializeReciprocalSplatTensor = [&](DenseFPElementsAttr denseAttr, RankedTensorType targetType) {
    APFloat reciprocal = denseAttr.getSplatValue<APFloat>();
    APFloat one(reciprocal.getSemantics(), 1);
    [[maybe_unused]] APFloat::opStatus status = one.divide(reciprocal, APFloat::rmNearestTiesToEven);
    assert(!(status & APFloat::opInvalidOp) && "distributed conv div consumer requires finite non-zero scalar");
    return getOrCreateConstant(
      rewriter, rewriter.getInsertionBlock()->getParentOp(), DenseFPElementsAttr::get(targetType, one), targetType);
  };

  auto applyDistributedConsumers = [&](Value fragment) {
    Value current = fragment;
    for (const DistributedTensorStep& step : distributedConsumers) {
      auto fragmentType = cast<RankedTensorType>(current.getType());
      switch (step.kind) {
      case DistributedTensorOpKind::Relu:
        current = spatial::SpatReluOp::create(rewriter, loc, fragmentType, current).getResult();
        break;
      case DistributedTensorOpKind::Sigmoid:
        current = spatial::SpatSigmoidOp::create(rewriter, loc, fragmentType, current).getResult();
        break;
      case DistributedTensorOpKind::Add: {
        Value splat = materializeSplatTensor(step.constantAttr, fragmentType);
        current = spatial::SpatVAddOp::create(rewriter, loc, fragmentType, current, splat).getResult();
        break;
      }
      case DistributedTensorOpKind::Sub: {
        Value splat = materializeSplatTensor(step.constantAttr, fragmentType);
        Value lhs = step.fragmentOnLhs ? current : splat;
        Value rhs = step.fragmentOnLhs ? splat : current;
        current = spatial::SpatVSubOp::create(rewriter, loc, fragmentType, lhs, rhs).getResult();
        break;
      }
      case DistributedTensorOpKind::Mul: {
        Value splat = materializeSplatTensor(step.constantAttr, fragmentType);
        current = spatial::SpatVMulOp::create(rewriter, loc, fragmentType, current, splat).getResult();
        break;
      }
      case DistributedTensorOpKind::Div: {
        auto reciprocalAttr = cast<DenseFPElementsAttr>(step.constantAttr);
        Value reciprocal = materializeReciprocalSplatTensor(reciprocalAttr, fragmentType);
        current = spatial::SpatVMulOp::create(rewriter, loc, fragmentType, current, reciprocal).getResult();
        break;
      }
      case DistributedTensorOpKind::Conv:
        llvm_unreachable("conv-consuming distributed chains should not materialize through createCollectedConvOutput");
      }
    }
    return current;
  };

  auto collectComputeOp = createSpatCompute(rewriter, loc, convType, {}, gemmRows, [&](ValueRange gemmRowArgs) {
    SmallVector<Value> transformedRows;
    transformedRows.reserve(gemmRowArgs.size());
    for (Value row : gemmRowArgs)
      transformedRows.push_back(applyDistributedConsumers(row));

    Value gemmOut;
    if (packFactor == 1) {
      gemmOut = createSpatConcat(rewriter, loc, /*axis=*/0, transformedRows);
    }
    else {
      Value packedOutput = createSpatConcat(rewriter, loc, /*axis=*/0, transformedRows);
      gemmOut = standard::unpackRowsFromParallelGemm(
        packedOutput, cast<RankedTensorType>(packedOutput.getType()), numPatches, numChannelsOut, packFactor, rewriter, loc);
    }

    // Restore output layout:
    // [numPatches, numChannelsOut]
    // -> [N, Hout, Wout, Cout]
    // -> [N, Cout, Hout, Wout]
    Value nhwcOut = tensor::ExpandShapeOp::create(rewriter,
                                                  loc,
                                                  nhwcType,
                                                  gemmOut,
                                                  SmallVector<ReassociationIndices> {
                                                    {0, 1, 2},
                                                    {3}
    });
    Value nchwOut = ONNXTransposeOp::create(rewriter, loc, outType, nhwcOut, rewriter.getI64ArrayAttr({0, 3, 1, 2}));
    spatial::SpatYieldOp::create(rewriter, loc, nchwOut);
  });
  return collectComputeOp.getResult(0);
}

static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp, Value x, Value w, Value b) {
  ConvLoweringState state;
  state.x = x;
  state.w = w;
  state.b = b;
  state.xType = cast<RankedTensorType>(state.x.getType());
  state.wType = cast<RankedTensorType>(state.w.getType());
  state.outType = cast<RankedTensorType>(convOp.getY().getType());

  if (!state.xType.hasStaticShape()) {
    pim::emitUnsupportedStaticShapeDiagnostic(convOp, "conv input");
    return failure();
  }
  if (!state.wType.hasStaticShape()) {
    pim::emitUnsupportedStaticShapeDiagnostic(convOp, "conv weight");
    return failure();
  }
  if (!state.outType.hasStaticShape()) {
    pim::emitUnsupportedStaticShapeDiagnostic(convOp, "conv result");
    return failure();
  }
  if (state.xType.getRank() != 4) {
    pim::emitUnsupportedRankDiagnostic(convOp, "conv input", state.xType.getRank(), {4});
    return failure();
  }
  if (state.wType.getRank() != 4) {
    pim::emitUnsupportedRankDiagnostic(convOp, "conv weight", state.wType.getRank(), {4});
    return failure();
  }
  if (state.outType.getRank() != 4) {
    pim::emitUnsupportedRankDiagnostic(convOp, "conv result", state.outType.getRank(), {4});
    return failure();
  }

  state.group = convOp.getGroup();
  if (state.group < 1) {
    convOp.emitOpError("requires group >= 1 for Spatial lowering");
    return failure();
  }

  state.batchSize = state.xType.getDimSize(0);
  state.numChannelsIn = state.xType.getDimSize(1);
  state.xHeight = state.xType.getDimSize(2);
  state.xWidth = state.xType.getDimSize(3);
  state.numChannelsOut = state.wType.getDimSize(0);
  state.wHeight = state.wType.getDimSize(2);
  state.wWidth = state.wType.getDimSize(3);
  state.outHeight = state.outType.getDimSize(2);
  state.outWidth = state.outType.getDimSize(3);
  state.hasBias = state.b && !isa<ONNXNoneOp>(state.b.getDefiningOp());

  if (state.numChannelsIn % state.group != 0) {
    convOp.emitOpError() << "requires input channels " << state.numChannelsIn << " to be divisible by group "
                         << state.group << " for Spatial lowering";
    return failure();
  }
  if (state.numChannelsOut % state.group != 0) {
    convOp.emitOpError() << "requires output channels " << state.numChannelsOut << " to be divisible by group "
                         << state.group << " for Spatial lowering";
    return failure();
  }

  state.numChannelsInPerGroup = state.numChannelsIn / state.group;
  state.numChannelsOutPerGroup = state.numChannelsOut / state.group;
  if (state.wType.getDimSize(1) != state.numChannelsInPerGroup) {
    convOp.emitOpError() << "requires grouped conv weight input channels " << state.wType.getDimSize(1)
                         << " to match input channels per group " << state.numChannelsInPerGroup
                         << " for Spatial lowering";
    return failure();
  }
  if (state.wType.getDimSize(0) != state.numChannelsOut) {
    convOp.emitOpError() << "requires weight output channels " << state.wType.getDimSize(0)
                         << " to match result channels " << state.numChannelsOut << " for Spatial lowering";
    return failure();
  }

  const auto stridesAttr = convOp.getStrides();
  const auto dilationsAttr = convOp.getDilations();
  const auto padsAttr = convOp.getPads();

  if (stridesAttr && stridesAttr->size() != 2) {
    convOp.emitOpError("requires exactly two stride values for Spatial lowering");
    return failure();
  }
  if (dilationsAttr && dilationsAttr->size() != 2) {
    convOp.emitOpError("requires exactly two dilation values for Spatial lowering");
    return failure();
  }
  if (padsAttr && padsAttr->size() != 4) {
    convOp.emitOpError("requires exactly four pad values for 2D Spatial lowering");
    return failure();
  }

  state.strideHeight = getOptionalI64Attr(stridesAttr, 0, 1);
  state.strideWidth = getOptionalI64Attr(stridesAttr, 1, 1);
  state.dilationHeight = getOptionalI64Attr(dilationsAttr, 0, 1);
  state.dilationWidth = getOptionalI64Attr(dilationsAttr, 1, 1);
  state.padHeightBegin = 0;
  state.padHeightEnd = 0;
  state.padWidthBegin = 0;
  state.padWidthEnd = 0;

  if (padsAttr) {
    state.padHeightBegin = getI64Attr(*padsAttr, 0);
    state.padWidthBegin = getI64Attr(*padsAttr, 1);
    state.padHeightEnd = getI64Attr(*padsAttr, 2);
    state.padWidthEnd = getI64Attr(*padsAttr, 3);
    return state;
  }

  const auto autoPad = convOp.getAutoPad();
  if (autoPad == "SAME_UPPER" || autoPad == "SAME_LOWER") {
    const int64_t effectiveKernelH = (state.wHeight - 1) * state.dilationHeight + 1;
    const int64_t effectiveKernelW = (state.wWidth - 1) * state.dilationWidth + 1;
    const int64_t totalPadH =
      std::max(static_cast<int64_t>(0), (state.outHeight - 1) * state.strideHeight + effectiveKernelH - state.xHeight);
    const int64_t totalPadW =
      std::max(static_cast<int64_t>(0), (state.outWidth - 1) * state.strideWidth + effectiveKernelW - state.xWidth);

    if (autoPad == "SAME_UPPER") {
      state.padHeightBegin = totalPadH / 2;
      state.padHeightEnd = totalPadH - state.padHeightBegin;
      state.padWidthBegin = totalPadW / 2;
      state.padWidthEnd = totalPadW - state.padWidthBegin;
    }
    else {
      state.padHeightEnd = totalPadH / 2;
      state.padHeightBegin = totalPadH - state.padHeightEnd;
      state.padWidthEnd = totalPadW / 2;
      state.padWidthBegin = totalPadW - state.padWidthEnd;
    }
    return state;
  }

  if (autoPad != "NOTSET" && autoPad != "VALID") {
    convOp.emitOpError() << "unsupported auto_pad value `" << autoPad << "` for Spatial lowering";
    return failure();
  }

  return state;
}

static FailureOr<ConvLoweringState> analyzeConvLoweringState(ONNXConvOp convOp, ONNXConvOpAdaptor convOpAdaptor) {
  return analyzeConvLoweringState(convOp, convOpAdaptor.getX(), convOpAdaptor.getW(), convOpAdaptor.getB());
}

static FailureOr<ConvLoweringState> analyzeConvLoweringState(spatial::SpatConv2DPlanOp planOp) {
  ConvLoweringState state;
  state.x = planOp.getInput();
  state.w = planOp.getWeight();
  state.b = planOp.getBias() ? planOp.getBias() : Value();
  state.xType = dyn_cast<RankedTensorType>(state.x.getType());
  state.wType = dyn_cast<RankedTensorType>(state.w.getType());
  state.outType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());

  if (!state.xType || !state.wType || !state.outType)
    return planOp.emitOpError("requires ranked tensor input, weight, and output"), failure();
  if (!state.xType.hasStaticShape() || !state.wType.hasStaticShape() || !state.outType.hasStaticShape())
    return planOp.emitOpError("requires static input, weight, and output shapes"), failure();
  if (state.xType.getRank() != 4 || state.wType.getRank() != 4 || state.outType.getRank() != 4)
    return planOp.emitOpError("requires rank-4 input, weight, and output tensors"), failure();

  state.group = planOp.getGroup();
  if (state.group < 1)
    return planOp.emitOpError("requires group >= 1"), failure();

  state.batchSize = state.xType.getDimSize(0);
  state.numChannelsIn = state.xType.getDimSize(1);
  state.xHeight = state.xType.getDimSize(2);
  state.xWidth = state.xType.getDimSize(3);
  state.numChannelsOut = state.wType.getDimSize(0);
  state.wHeight = state.wType.getDimSize(2);
  state.wWidth = state.wType.getDimSize(3);
  state.outHeight = state.outType.getDimSize(2);
  state.outWidth = state.outType.getDimSize(3);
  state.hasBias = static_cast<bool>(planOp.getBias());

  if (state.numChannelsIn % state.group != 0 || state.numChannelsOut % state.group != 0)
    return planOp.emitOpError("requires input and output channels divisible by group"), failure();

  state.numChannelsInPerGroup = state.numChannelsIn / state.group;
  state.numChannelsOutPerGroup = state.numChannelsOut / state.group;
  if (state.wType.getDimSize(1) != state.numChannelsInPerGroup)
    return planOp.emitOpError("requires grouped conv weight channels to match input channels per group"), failure();

  auto pads = planOp.getPads();
  auto strides = planOp.getStrides();
  auto dilations = planOp.getDilations();
  if (pads.size() != 4 || strides.size() != 2 || dilations.size() != 2)
    return planOp.emitOpError("requires 4 pads, 2 strides, and 2 dilations"), failure();

  state.padHeightBegin = pads[0];
  state.padWidthBegin = pads[1];
  state.padHeightEnd = pads[2];
  state.padWidthEnd = pads[3];
  state.strideHeight = strides[0];
  state.strideWidth = strides[1];
  state.dilationHeight = dilations[0];
  state.dilationWidth = dilations[1];
  return state;
}

static FailureOr<PimConvLoweringType> resolveRequestedConvLoweringStrategy(Operation* op) {
  if (!useExperimentalConvImpl)
    return pimConvLowering.getValue();

  if (pimConvLowering != PimConvLoweringAuto && pimConvLowering != PimConvLoweringPackedIm2Col) {
    op->emitOpError() << "--use-experimental-conv-impl conflicts with --pim-conv-lowering="
                      << stringifyConvLoweringStrategy(pimConvLowering);
    return failure();
  }
  return PimConvLoweringPackedIm2Col;
}

static LogicalResult verifyForcedConvLoweringStrategy(Operation* op,
                                                      const ConvGeometry& geo,
                                                      PimConvLoweringType strategy) {
  switch (strategy) {
  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<Value> lowerDenseSelectedConvPlan(Operation* op,
                                                   const ConvLoweringState& state,
                                                   PimConvLoweringType strategy,
                                                   PatternRewriter& rewriter,
                                                   Location loc);

static ConvLoweringState makeGroupedConvLoweringState(const ConvLoweringState& parent,
                                                      Value groupX,
                                                      Value groupW,
                                                      Value groupB,
                                                      RankedTensorType groupOutType);

static FailureOr<Value> buildConvValueForStrategy(Operation* op,
                                                  Location loc,
                                                  const ConvLoweringState& state,
                                                  const ConvLoweringDecision& decision,
                                                  const DistributedConvAnalysis& analysis,
                                                  ArrayRef<DistributedTensorStep> distributedConsumers,
                                                  PatternRewriter& rewriter);

static FailureOr<Value> buildGroupedConvValue(Operation* op,
                                              Location loc,
                                              const ConvLoweringState& state,
                                              const ConvLoweringDecision& decision,
                                              PatternRewriter& rewriter);

static FailureOr<Value> lowerGroupedSelectedConvPlan(Operation* op,
                                                     const ConvLoweringState& state,
                                                     PimConvLoweringType strategy,
                                                     PatternRewriter& rewriter,
                                                     Location loc) {
  ConvLoweringDecision decision {strategy, "", false, "", ""};
  return buildGroupedConvValue(op, loc, state, decision, rewriter);
}

static FailureOr<Value> lowerDenseSelectedConvPlan(Operation* op,
                                                   const ConvLoweringState& state,
                                                   PimConvLoweringType strategy,
                                                   PatternRewriter& rewriter,
                                                   Location loc) {
  DistributedConvAnalysis analysis;
  analysis.barrierKind = DistributedConvBarrierKind::UnsupportedConsumer;
  analysis.barrierDetail = "selected dense layout";
  ConvLoweringDecision decision {strategy, "", false, "", ""};
  return buildConvValueForStrategy(op, loc, state, decision, analysis, {}, rewriter);
}

static FailureOr<Value> buildConvValueForStrategy(Operation* op,
                                                  Location loc,
                                                  const ConvLoweringState& state,
                                                  const ConvLoweringDecision& decision,
                                                  const DistributedConvAnalysis& analysis,
                                                  ArrayRef<DistributedTensorStep> distributedConsumers,
                                                  PatternRewriter& rewriter) {
  (void)analysis;
  const ConvGeometry geo = buildConvGeometry(state);
  switch (decision.strategy) {
  case PimConvLoweringDepthwise: {
    return depthwise::rewriteConv(op, state, rewriter, loc);
  }
  case PimConvLoweringLegacy:
  case PimConvLoweringPackedIm2Col: {
    return standard::rewritePackedIm2ColConv(state, distributedConsumers, rewriter, loc);
  }
  case PimConvLoweringStreamedPatch:
  case PimConvLoweringOutputChannelTiled:
  case PimConvLoweringTiled2D: {
    return standard::rewriteStreamedConv(state, distributedConsumers, rewriter, loc, /*forcedPackFactor=*/1);
  }
  case PimConvLoweringInputKTiled: {
    return standard::rewriteInputKTiledConv(state, distributedConsumers, rewriter, loc);
  }
  case PimConvLoweringStreamedPacked: {
    return standard::rewriteStreamedConv(state, distributedConsumers, rewriter, loc, geo.pack);
  }
  case PimConvLoweringAuto:
    break;
  }
  op->emitOpError("unexpected auto strategy at Conv lowering dispatch");
  return failure();
}

static LogicalResult
createConvValueForStrategy(ONNXConvOp convOp,
                           const ConvLoweringState& state,
                           const ConvLoweringDecision& decision,
                           const DistributedConvAnalysis& analysis,
                           ArrayRef<DistributedTensorStep> distributedConsumers,
                           PatternRewriter& rewriter,
                           FailureOr<Value>& result) {
  result = buildConvValueForStrategy(convOp, convOp.getLoc(), state, decision, analysis, distributedConsumers, rewriter);
  if (failed(result))
    return failure();

  const ConvGeometry geo = buildConvGeometry(state);
  const ConvStrategyEstimate estimate = estimateConvStrategy(geo, decision.strategy, analysis);
  switch (decision.strategy) {
  case PimConvLoweringDepthwise:
    reportConvLoweringDecision(
      convOp, geo, decision, estimate, /*batchSize=*/geo.p, /*numberOfBatches=*/1, /*usesComputeBatch=*/true,
      /*usesBatchedInstructionEmission=*/true, std::nullopt);
    return success();
  case PimConvLoweringLegacy:
  case PimConvLoweringPackedIm2Col:
    reportConvLoweringDecision(
      convOp, geo, decision, estimate, /*batchSize=*/geo.pack, /*numberOfBatches=*/1, /*usesComputeBatch=*/true,
      /*usesBatchedInstructionEmission=*/true, std::nullopt);
    return success();
  case PimConvLoweringStreamedPatch:
  case PimConvLoweringOutputChannelTiled:
  case PimConvLoweringTiled2D: {
    uint64_t chunkPositions = chooseStreamChunkPositions(geo, /*packFactor=*/1);
    const int64_t batches = ceilIntegerDivide(geo.p, static_cast<int64_t>(chunkPositions));
    reportConvLoweringDecision(convOp,
                               geo,
                               decision,
                               estimate,
                               /*batchSize=*/1,
                               batches,
                               /*usesComputeBatch=*/true,
                               /*usesBatchedInstructionEmission=*/true,
                               chunkPositions);
    return success();
  }
  case PimConvLoweringInputKTiled: {
    const int64_t numKSlices = ceilIntegerDivide(geo.k, geo.xbarSize);
    const uint64_t maxLanesPerBatch =
      std::max<uint64_t>(1,
                         static_cast<uint64_t>(crossbarCountInCore.getValue())
                           / 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),
                          maxLanesPerBatch,
                          static_cast<uint64_t>(state.outWidth)}));
    const int64_t batches =
      state.batchSize * state.outHeight * ceilIntegerDivide(state.outWidth, static_cast<int64_t>(rowChunkWidth));
    reportConvLoweringDecision(convOp,
                               geo,
                               decision,
                               estimate,
                               /*batchSize=*/1,
                               batches,
                               /*usesComputeBatch=*/false,
                               /*usesBatchedInstructionEmission=*/false,
                               rowChunkWidth);
    return success();
  }
  case PimConvLoweringStreamedPacked: {
    uint64_t chunkPositions = chooseStreamChunkPositions(geo, geo.pack);
    const int64_t batches = ceilIntegerDivide(geo.p, static_cast<int64_t>(chunkPositions));
    reportConvLoweringDecision(convOp,
                               geo,
                               decision,
                               estimate,
                               /*batchSize=*/geo.pack,
                               batches,
                               /*usesComputeBatch=*/true,
                               /*usesBatchedInstructionEmission=*/true,
                               chunkPositions);
    return success();
  }
  case PimConvLoweringAuto:
    break;
  }
  return convOp.emitOpError("unexpected auto strategy at Conv lowering dispatch");
}

static LogicalResult
rewriteSelectedConv(ONNXConvOp convOp,
                    const ConvLoweringState& state,
                    const ConvLoweringDecision& decision,
                    const DistributedConvAnalysis& analysis,
                    PatternRewriter& rewriter) {
  FailureOr<Value> result = failure();
  if (failed(createConvValueForStrategy(convOp, state, decision, analysis, analysis.steps, rewriter, result)))
    return failure();

  if (!analysis.hasLocalConsumers()) {
    rewriter.replaceOp(convOp, *result);
    return success();
  }

  assert(analysis.replacementOp && "conv rewrite expects a replacement op");
  rewriter.replaceOp(analysis.replacementOp, *result);
  for (auto it = analysis.steps.rbegin(); it != analysis.steps.rend(); ++it)
    if (it->op != analysis.replacementOp)
      rewriter.eraseOp(it->op);
  rewriter.eraseOp(convOp);
  return success();
}

[[maybe_unused]] static LogicalResult
rewriteUngroupedConv(ONNXConvOp convOp,
                     const ConvLoweringState& state,
                     const ConvLoweringDecision& decision,
                     const DistributedConvAnalysis& analysis,
                     PatternRewriter& rewriter) {
  return rewriteSelectedConv(convOp, state, decision, analysis, rewriter);
}

static LogicalResult
rewriteGroupedConv(ONNXConvOp convOp,
                   const ConvLoweringState& state,
                   const ConvLoweringDecision& decision,
                   PatternRewriter& rewriter);

static ConvLoweringState makeGroupedConvLoweringState(const ConvLoweringState& parent,
                                                      Value groupX,
                                                      Value groupW,
                                                      Value groupB,
                                                      RankedTensorType groupOutType);

static ConvLoweringState makeGroupedConvLoweringState(
  const ConvLoweringState& parent, Value groupX, Value groupW, Value groupB, RankedTensorType groupOutType) {
  ConvLoweringState state = parent;
  state.x = groupX;
  state.w = groupW;
  state.b = groupB;
  state.xType = cast<RankedTensorType>(groupX.getType());
  state.wType = cast<RankedTensorType>(groupW.getType());
  state.outType = groupOutType;
  state.batchSize = state.xType.getDimSize(0);
  state.numChannelsIn = state.xType.getDimSize(1);
  state.xHeight = state.xType.getDimSize(2);
  state.xWidth = state.xType.getDimSize(3);
  state.numChannelsOut = state.wType.getDimSize(0);
  state.wHeight = state.wType.getDimSize(2);
  state.wWidth = state.wType.getDimSize(3);
  state.outHeight = state.outType.getDimSize(2);
  state.outWidth = state.outType.getDimSize(3);
  state.group = 1;
  state.numChannelsInPerGroup = state.numChannelsIn;
  state.numChannelsOutPerGroup = state.numChannelsOut;
  state.hasBias = static_cast<bool>(groupB);
  return state;
}

static FailureOr<Value> buildGroupedConvValue(Operation* op,
                                              Location loc,
                                              const ConvLoweringState& state,
                                              const ConvLoweringDecision& decision,
                                              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);
  SmallVector<Value> bSlices;
  if (state.hasBias) {
    auto biasType = cast<RankedTensorType>(state.b.getType());
    int64_t biasAxis = -1;
    if (biasType.getRank() == 1)
      biasAxis = 0;
    else if (biasType.getRank() == 2)
      biasAxis = biasType.getDimSize(0) != 1 ? 0 : 1;
    else {
      op->emitOpError() << "requires rank-1 or rank-2 bias for grouped convolution Spatial lowering, but got rank "
                        << biasType.getRank();
      return failure();
    }
    bSlices = sliceTensor(state.b, biasAxis, state.numChannelsOutPerGroup, rewriter, loc);
  }

  if (xSlices.size() != static_cast<size_t>(state.group) || wSlices.size() != static_cast<size_t>(state.group)
      || (state.hasBias && bSlices.size() != static_cast<size_t>(state.group))) {
    op->emitOpError("failed to partition grouped convolution operands for Spatial lowering");
    return failure();
  }

  SmallVector<Value> groupResults;
  groupResults.reserve(state.group);
  auto groupOutType = RankedTensorType::get(
    {state.batchSize, state.numChannelsOutPerGroup, state.outHeight, state.outWidth}, state.outType.getElementType());
  for (int64_t groupId = 0; groupId < state.group; groupId++) {
    Value groupX = xSlices[groupId];
    Value groupW = wSlices[groupId];
    Value groupB = state.hasBias ? bSlices[groupId] : Value();
    ConvLoweringState groupState = makeGroupedConvLoweringState(state, groupX, groupW, groupB, groupOutType);
    DistributedConvAnalysis groupAnalysis;
    groupAnalysis.barrierKind = DistributedConvBarrierKind::GroupedConv;
    groupAnalysis.barrierDetail = "grouped convolution still materializes densely";
    FailureOr<Value> groupResult =
      buildConvValueForStrategy(op, loc, groupState, decision, groupAnalysis, {}, rewriter);
    if (failed(groupResult))
      return failure();
    groupResults.push_back(*groupResult);
  }

  if (llvm::all_of(groupResults, isCompileTimeComputable))
    return createSpatConcat(rewriter, loc, /*axis=*/1, groupResults);

  auto concatCompute = createSpatCompute(rewriter, loc, TypeRange {state.outType}, {}, groupResults, [&](ValueRange args) {
    spatial::SpatYieldOp::create(rewriter, loc, createSpatConcat(rewriter, loc, /*axis=*/1, args));
  });
  return concatCompute.getResult(0);
}

[[maybe_unused]] static LogicalResult
rewriteGroupedConv(ONNXConvOp convOp,
                   const ConvLoweringState& state,
                   const ConvLoweringDecision& decision,
                   PatternRewriter& rewriter) {
  FailureOr<Value> result = buildGroupedConvValue(convOp.getOperation(), convOp.getLoc(), state, decision, rewriter);
  if (failed(result))
    return failure();
  rewriter.replaceOp(convOp, *result);
  return success();
}

} // namespace

LogicalResult ConvToGemm::matchAndRewrite(ONNXConvOp convOp,
                                          ONNXConvOpAdaptor convOpAdaptor,
                                          ConversionPatternRewriter& rewriter) const {
  FailureOr<ConvLoweringState> state = analyzeConvLoweringState(convOp, convOpAdaptor);
  if (failed(state))
    return failure();
  SmallVector<int64_t> pads {
    state->padHeightBegin, state->padWidthBegin, state->padHeightEnd, state->padWidthEnd};
  SmallVector<int64_t> strides {state->strideHeight, state->strideWidth};
  SmallVector<int64_t> dilations {state->dilationHeight, state->dilationWidth};
  Value bias = state->hasBias ? convOpAdaptor.getB() : Value();
  auto convPlan = spatial::SpatConv2DPlanOp::create(rewriter,
                                                    convOp.getLoc(),
                                                    convOp.getY().getType(),
                                                    convOpAdaptor.getX(),
                                                    convOpAdaptor.getW(),
                                                    bias,
                                                    rewriter.getDenseI64ArrayAttr(pads),
                                                    rewriter.getDenseI64ArrayAttr(strides),
                                                    rewriter.getDenseI64ArrayAttr(dilations),
                                                    rewriter.getI64IntegerAttr(state->group),
                                                    rewriter.getStringAttr("nchw"));
  rewriter.replaceOp(convOp, convPlan.getResult());
  return success();
}

void populateConvPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.insert<ConvToGemm>(ctx); }

LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) {
  FailureOr<ConvLoweringState> state = analyzeConvLoweringState(planOp);
  if (failed(state))
    return failure();

  if (state->group != 1 || state->batchSize != 1)
    return failure();
  if (state->outType.getRank() != 4 || !state->outType.hasStaticShape())
    return failure();

  FailureOr<PimConvLoweringType> requestedStrategy = resolveRequestedConvLoweringStrategy(planOp.getOperation());
  if (failed(requestedStrategy))
    return failure();

  DistributedConvAnalysis analysis;
  analysis.barrierKind = DistributedConvBarrierKind::UnsupportedConsumer;
  analysis.barrierDetail = "selected row-strip layout";
  ConvGeometry geometry = buildConvGeometry(*state);
  ConvLoweringDecision decision = chooseConvLoweringStrategy(geometry, *requestedStrategy, analysis);
  if (decision.strategy == PimConvLoweringDepthwise && !depthwise::canUseStructuredRewrite(*state)
      && *requestedStrategy == PimConvLoweringAuto) {
    decision = {PimConvLoweringLegacy,
                "depthwise auto fallback when structured depthwise lowering is not representable",
                /*isAuto=*/true,
                "",
                ""};
  }
  if (failed(verifyForcedConvLoweringStrategy(planOp.getOperation(), geometry, decision.strategy)))
    return failure();

  switch (decision.strategy) {
  case PimConvLoweringLegacy:
  case PimConvLoweringPackedIm2Col:
  case PimConvLoweringStreamedPatch:
  case PimConvLoweringOutputChannelTiled:
  case PimConvLoweringTiled2D:
  case PimConvLoweringStreamedPacked:
    return success();
  case PimConvLoweringAuto:
  case PimConvLoweringDepthwise:
  case PimConvLoweringInputKTiled:
    return failure();
  }
  llvm_unreachable("unknown conv lowering strategy");
}

LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp) {
  FailureOr<ConvLoweringState> state = analyzeConvLoweringState(planOp);
  if (failed(state))
    return failure();

  StringRef failureReason;
  return canDirectLowerRowStripConv(*state, failureReason) ? success() : failure();
}

FailureOr<Value>
lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
                        std::optional<Value> rowStripInput,
                        bool emitRowStripLayout,
                        PatternRewriter& rewriter) {
  FailureOr<ConvLoweringState> state = analyzeConvLoweringState(planOp);
  if (failed(state))
    return failure();

  FailureOr<PimConvLoweringType> requestedStrategy = resolveRequestedConvLoweringStrategy(planOp.getOperation());
  if (failed(requestedStrategy))
    return failure();

  DistributedConvAnalysis analysis;
  analysis.barrierKind = DistributedConvBarrierKind::UnsupportedConsumer;
  analysis.barrierDetail = emitRowStripLayout ? "selected row-strip layout" : "selected dense layout";
  ConvGeometry geometry = buildConvGeometry(*state);
  ConvLoweringDecision decision = chooseConvLoweringStrategy(geometry, *requestedStrategy, analysis);
  if (decision.strategy == PimConvLoweringDepthwise && !depthwise::canUseStructuredRewrite(*state)
      && *requestedStrategy == PimConvLoweringAuto) {
    decision = {PimConvLoweringLegacy,
                "depthwise auto fallback when structured depthwise lowering is not representable",
                /*isAuto=*/true,
                "",
                ""};
  }
  if (failed(verifyForcedConvLoweringStrategy(planOp.getOperation(), geometry, decision.strategy)))
    return failure();

  if (emitRowStripLayout) {
    if (rowStripInput) {
      if (failed(canConsumeAndProduceRowStrip(planOp)))
        return planOp.emitOpError("selected row-strip input/output layout is not supported for this Conv plan"), failure();
      return createConvRowsFromRowStripInput(*state, decision, *rowStripInput, rewriter, planOp.getLoc());
    }
    if (failed(canLowerConvPlanToRowStrip(planOp)))
      return planOp.emitOpError("selected row-strip layout is not supported for this Conv plan"), failure();
    FailureOr<Value> rows = createConvRowsForStrategy(*state, decision, rewriter, planOp.getLoc());
    if (failed(rows))
      return failure();
    FailureOr<Value> packedRows = createRowStripPackedRows(*rows, *state, rewriter, planOp.getLoc());
    if (failed(packedRows))
      return planOp.emitOpError("failed to pack Conv rows into the selected row-strip physical layout"), failure();
    return *packedRows;
  }

  if (decision.strategy == PimConvLoweringDepthwise)
    return lowerDenseSelectedConvPlan(planOp.getOperation(), *state, decision.strategy, rewriter, planOp.getLoc());
  if (state->group != 1)
    return lowerGroupedSelectedConvPlan(planOp.getOperation(), *state, decision.strategy, rewriter, planOp.getLoc());
  return lowerDenseSelectedConvPlan(planOp.getOperation(), *state, decision.strategy, rewriter, planOp.getLoc());
}

} // namespace onnx_mlir