Raptor/src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

#include "mlir/Dialect/Affine/IR/AffineOps.h"
#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/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/DialectConversion.h"

#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
#include "mlir/Transforms/Passes.h"
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
#include "src/Accelerators/PIM/Pass/PIMPasses.h"

using namespace mlir;

namespace onnx_mlir {
namespace {

static FailureOr<RowStripPhysicalValue> getRowStripValue(Value value) {
  return getRowStripPhysicalValue(value);
}

static FailureOr<Value> publishRowStripValue(Operation* planOp,
                                             Value storage,
                                             PatternRewriter& rewriter) {
  auto logicalType = dyn_cast<RankedTensorType>(planOp->getResult(0).getType());
  if (!logicalType)
    return planOp->emitOpError("requires ranked logical output type"), failure();
  FailureOr<RowStripPhysicalValue> value = describeRowStripPhysicalValue(storage, logicalType);
  if (failed(value))
    return planOp->emitOpError("lowering produced invalid row-strip physical storage"), failure();
  FailureOr<Value> blueprint = createRowStripStorageBlueprint(
      storage, logicalType, rewriter, planOp->getLoc());
  if (failed(blueprint))
    return planOp->emitOpError("failed to create row-strip storage Blueprint"), failure();
  rewriter.replaceOp(planOp, *blueprint);
  return *blueprint;
}

static bool isRowStripSelected(Operation* op) {
  auto selected = spatial::getSelectedPhysicalLayout(op);
  return selected && *selected == spatial::PhysicalLayout::NHWCRowStrip;
}

static bool isDenseSelected(Operation* op) {
  auto selected = spatial::getSelectedPhysicalLayout(op);
  return selected && *selected == spatial::PhysicalLayout::DenseNCHW;
}

static spatial::PhysicalLayout getKnownPhysicalLayout(Value value) {
  if (auto materialize = value.getDefiningOp<spatial::SpatMaterializeLayoutOp>())
    return materialize.getTargetPhysicalLayout();
  if (auto blueprint = value.getDefiningOp<spatial::SpatBlueprintOp>())
    return blueprint.getPhysicalLayout();
  if (Operation* producer = value.getDefiningOp()) {
    if (auto selected = spatial::getSelectedPhysicalLayout(producer))
      return *selected;
  }
  return spatial::PhysicalLayout::DenseNCHW;
}

static LogicalResult verifySelectedLayouts(
    func::FuncOp funcOp, const spatial::SpatialTargetInfo& target) {
  LogicalResult result = success();
  funcOp.walk([&](Operation* op) {
    auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(op);
    if (!capability)
      return;
    auto selected = spatial::getSelectedPhysicalLayout(op);
    if (!selected) {
      op->emitOpError("requires a selected physical layout from SpatialLayoutPlanning");
      result = failure();
      return;
    }
    if (*selected != spatial::PhysicalLayout::DenseNCHW
        && *selected != spatial::PhysicalLayout::NHWCRowStrip) {
      op->emitOpError("has an unsupported selected physical layout");
      result = failure();
      return;
    }
    SmallVector<spatial::PhysicalLayout> operandLayouts;
    operandLayouts.reserve(op->getNumOperands());
    for (Value operand : op->getOperands())
      operandLayouts.push_back(getKnownPhysicalLayout(operand));
    auto alternatives = capability.getLayoutAlternatives(target, operandLayouts);
    if (llvm::none_of(alternatives, [&](const spatial::LayoutAlternative& alternative) {
          return alternative.resultLayout == *selected
              && alternative.operandLayouts == operandLayouts;
        })) {
      op->emitOpError("selected physical layout is not lowerable for its explicit operand layouts");
      result = failure();
    }
  });
  return result;
}

static FailureOr<Value>
lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) {
  return applyRowStripRelu(input, rewriter, planOp.getLoc());
}

static FailureOr<Value>
lowerRowStripSilu(const RowStripPhysicalValue& input, spatial::SpatSiluPlanOp planOp, PatternRewriter& rewriter) {
  return applyRowStripSilu(input, rewriter, planOp.getLoc());
}

static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
                                             spatial::SpatBiasAddPlanOp planOp,
                                             PatternRewriter& rewriter) {
  return applyRowStripBiasAdd(input, planOp.getBias(), rewriter, planOp.getLoc());
}

static FailureOr<Value> lowerRowStripAdd(const RowStripPhysicalValue& lhs,
                                         const RowStripPhysicalValue& rhs,
                                         spatial::SpatAddPlanOp planOp,
                                         PatternRewriter& rewriter) {
  return applyRowStripAdd(lhs, rhs, rewriter, planOp.getLoc());
}

static FailureOr<Value> lowerRowStripConcat(ArrayRef<RowStripPhysicalValue> inputs,
                                            spatial::SpatConcatPlanOp planOp,
                                            PatternRewriter& rewriter) {
  auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
  if (!outputType)
    return failure();
  return applyRowStripConcat(inputs, outputType, rewriter, planOp.getLoc());
}

static FailureOr<Value>
materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
  if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
    return failure();
  return createRowStripAssemblyBlueprint(rowStripValue, rewriter, loc);
}

static FailureOr<Value> materializeDenseToRowStrip(
    Value input, RankedTensorType logicalType, Location loc, PatternRewriter& rewriter) {
  if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() != 4
      || logicalType.getDimSize(0) != 1)
    return failure();
  auto nhwcType = RankedTensorType::get(
    {1, logicalType.getDimSize(2), logicalType.getDimSize(3), logicalType.getDimSize(1)},
    logicalType.getElementType(), logicalType.getEncoding());
  auto rowsType = RankedTensorType::get(
    {logicalType.getDimSize(2) * logicalType.getDimSize(3), logicalType.getDimSize(1)},
    logicalType.getElementType(), logicalType.getEncoding());
  auto rowsCompute = createSpatCompute<1>(
    rewriter, loc, rowsType, {}, input, [&](Value denseInput) {
      Value nhwc = createLinalgTranspose(
        denseInput, nhwcType, {0, 2, 3, 1}, rewriter, loc);
      Value rows = tensor::CollapseShapeOp::create(
        rewriter, loc, rowsType, nhwc,
        SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
      spatial::SpatYieldOp::create(rewriter, loc, rows);
    });
  Value rows = rowsCompute->getResult(0);
  FailureOr<Value> storage = createRowStripStorageFromRows(rows, logicalType, rewriter, loc);
  if (failed(storage))
    return failure();
  return createRowStripStorageBlueprint(*storage, logicalType, rewriter, loc);
}

static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType,
                                               PatternRewriter& rewriter, Location loc) {
  auto producer = input.getDefiningOp<spatial::SpatGraphComputeBatch>();
  auto inputType = dyn_cast<RankedTensorType>(input.getType());
  auto biasType = dyn_cast<RankedTensorType>(bias.getType());
  if (!producer || !inputType || !biasType || !inputType.hasStaticShape() || !biasType.hasStaticShape()
      || !resultType.hasStaticShape() || inputType.getDimSize(0) != producer.getLaneCount()
      || biasType.getDimSize(0) != producer.getLaneCount() || resultType.getDimSize(0) != producer.getLaneCount())
    return failure();
  auto inputFragmentType = spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount());
  auto outputFragmentType = spatial::getGraphBatchFragmentType(resultType, producer.getLaneCount());
  if (failed(inputFragmentType) || failed(outputFragmentType) || inputFragmentType->getRank() != biasType.getRank()
      || inputFragmentType->getDimSize(0) != 1 || inputFragmentType->getShape().drop_front() != biasType.getShape().drop_front()
      || inputFragmentType->getRank() != outputFragmentType->getRank() + 1)
    return failure();
  for (auto [inputDim, outputDim] : llvm::zip(inputFragmentType->getShape().drop_front(), outputFragmentType->getShape()))
    if (outputDim > inputDim)
      return failure();

  auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {resultType}, producer.getLaneCount(), {}, ValueRange {input, bias},
    [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
      FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, *inputFragmentType);
      if (failed(fragment))
        return failure();
      MixedSliceGeometry biasSlice;
      for (int64_t dim : inputFragmentType->getShape()) {
        biasSlice.offsets.push_back(biasSlice.offsets.empty() ? OpFoldResult(args.lane) : rewriter.getIndexAttr(0));
        biasSlice.sizes.push_back(rewriter.getIndexAttr(dim));
        biasSlice.strides.push_back(rewriter.getIndexAttr(1));
      }
      Value biasFragment = extractMixedSliceOrIdentity(rewriter, loc, args.inputs[1], *inputFragmentType, biasSlice);
      if (!biasFragment)
        return failure();
      Value added = spatial::SpatVAddOp::create(rewriter, loc, *inputFragmentType, *fragment, biasFragment);
      MixedSliceGeometry outputSlice;
      outputSlice.offsets.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(0));
      outputSlice.sizes.push_back(rewriter.getIndexAttr(1));
      outputSlice.strides.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(1));
      for (int64_t dim : outputFragmentType->getShape())
        outputSlice.sizes.push_back(rewriter.getIndexAttr(dim));
      Value output = extractMixedSliceOrIdentity(rewriter, loc, added, *outputFragmentType, outputSlice);
      if (!output)
        return failure();
      publishGraphBatchPhysicalFragment(rewriter, loc, output, args.outputs.front(), args.lane);
      return success();
    });
  if (failed(batch))
    return failure();
  return batch->getResult(0);
}

struct LowerDenseReluPlan final : OpRewritePattern<spatial::SpatReluPlanOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    auto selected = spatial::getSelectedPhysicalLayout(planOp.getOperation());
    if (!selected || *selected != spatial::PhysicalLayout::DenseNCHW)
      return failure();

    auto computeOp = createSpatCompute<1>(
      rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) {
        auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), planOp.getOutput().getType(), x);
        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult());
      });
    rewriter.replaceOp(planOp, computeOp.getResults());
    return success();
  }
};

struct LowerDenseSiluPlan final : OpRewritePattern<spatial::SpatSiluPlanOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    auto selected = spatial::getSelectedPhysicalLayout(planOp.getOperation());
    if (!selected || *selected != spatial::PhysicalLayout::DenseNCHW)
      return failure();

    auto computeOp = createSpatCompute<1>(
      rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) {
        Value sigmoid = spatial::SpatSigmoidOp::create(
          rewriter, planOp.getLoc(), planOp.getOutput().getType(), x).getResult();
        Value silu = spatial::SpatVMulOp::create(
          rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, sigmoid).getResult();
        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), silu);
      });
    rewriter.replaceOp(planOp, computeOp.getResults());
    return success();
  }
};

struct LowerDenseResizePlan final : OpRewritePattern<spatial::SpatResizeNearestPlanOp> {
  explicit LowerDenseResizePlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
      : OpRewritePattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}

  LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    if (!isDenseSelected(planOp.getOperation()))
      return failure();
    FailureOr<Value> lowered = lowerSelectedResizeNearestPlan(planOp, std::nullopt, target, rewriter);
    if (failed(lowered))
      return planOp.emitOpError("failed to lower selected dense nearest Resize plan");
    rewriter.replaceOp(planOp, *lowered);
    return success();
  }

  const spatial::SpatialTargetInfo& target;
};

struct LowerDenseBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    if (!isDenseSelected(planOp.getOperation()))
      return failure();
    auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
    if (!resultType)
      return planOp.emitOpError("requires ranked output type");

    FailureOr<Value> denseBias = materializeDenseBiasAddTensor(
      planOp.getBias(), resultType, rewriter, planOp.getLoc());
    if (failed(denseBias))
      return planOp.emitOpError("failed to materialize dense Conv-style bias");
    if (planOp.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
      FailureOr<Value> lowered = lowerDenseBatchBiasAdd(
        planOp.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
      if (succeeded(lowered)) {
        rewriter.replaceOp(planOp, *lowered);
        return success();
      }
    }
    auto computeOp = createSpatCompute<2>(
      rewriter,
      planOp.getLoc(),
      planOp.getOutput().getType(),
      {},
      ValueRange {planOp.getInput(), *denseBias},
      [&](Value x, Value y) {
        auto added = spatial::SpatVAddOp::create(
          rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y);
        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult());
      });
    rewriter.replaceOp(planOp, computeOp.getResults());
    return success();
  }
};

struct LowerDenseAddPlan final : OpRewritePattern<spatial::SpatAddPlanOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    if (!isDenseSelected(planOp.getOperation()))
      return failure();
    auto compute = createSpatCompute<2>(
      rewriter,
      planOp.getLoc(),
      planOp.getOutput().getType(),
      {},
      ValueRange {planOp.getLhs(), planOp.getRhs()},
      [&](Value lhsValue, Value rhsValue) {
        Value added = spatial::SpatVAddOp::create(
          rewriter, planOp.getLoc(), planOp.getOutput().getType(), lhsValue, rhsValue);
        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added);
      });
    rewriter.replaceOp(planOp, compute.getResults());
    return success();
  }
};

struct LowerDenseConcatPlan final : OpRewritePattern<spatial::SpatConcatPlanOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    if (!isDenseSelected(planOp.getOperation()))
      return failure();
    auto compute = createSpatCompute(
      rewriter,
      planOp.getLoc(),
      TypeRange {planOp.getOutput().getType()},
      {},
      planOp.getInputs(),
      [&](ValueRange values) {
        Value concatenated = spatial::SpatConcatOp::create(
          rewriter,
          planOp.getLoc(),
          planOp.getOutput().getType(),
          rewriter.getI64IntegerAttr(planOp.getAxis()),
          values);
        spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), concatenated);
      });
    rewriter.replaceOp(planOp, compute.getResults());
    return success();
  }
};

static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp,
                                  PatternRewriter& rewriter) {
  FailureOr<RowStripPhysicalValue> lhs = getRowStripValue(planOp.getLhs());
  FailureOr<RowStripPhysicalValue> rhs = getRowStripValue(planOp.getRhs());
  if (isRowStripSelected(planOp.getOperation()) && failed(lhs)) {
    if (getKnownPhysicalLayout(planOp.getLhs()) == spatial::PhysicalLayout::NHWCRowStrip)
      return failure();
    return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs");
  }
  if (isRowStripSelected(planOp.getOperation()) && failed(rhs)) {
    if (getKnownPhysicalLayout(planOp.getRhs()) == spatial::PhysicalLayout::NHWCRowStrip)
      return failure();
    return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs");
  }
  if (isRowStripSelected(planOp.getOperation())) {
    rewriter.setInsertionPoint(planOp);
    FailureOr<Value> lowered = lowerRowStripAdd(*lhs, *rhs, planOp, rewriter);
    if (failed(lowered))
      return planOp.emitOpError("failed to lower selected row-strip Spatial add plan");
    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
      return failure();
    return success();
  }
  return planOp.emitOpError("dense Add plan was not lowered by the selected-plan patterns");
}

static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp,
                                     PatternRewriter& rewriter) {
  SmallVector<RowStripPhysicalValue> inputs;
  for (Value input : planOp.getInputs()) {
    FailureOr<RowStripPhysicalValue> physical = getRowStripValue(input);
    if (failed(physical)) {
      inputs.clear();
      break;
    }
    inputs.push_back(*physical);
  }
  if (isRowStripSelected(planOp.getOperation()) && inputs.size() != planOp.getInputs().size()) {
    if (llvm::any_of(planOp.getInputs(), [](Value input) {
          return getKnownPhysicalLayout(input) == spatial::PhysicalLayout::NHWCRowStrip;
        }))
      return failure();
    return planOp.emitOpError("selected row-strip Concat plan requires row-strip inputs");
  }
  if (isRowStripSelected(planOp.getOperation())) {
    rewriter.setInsertionPoint(planOp);
    FailureOr<Value> lowered = lowerRowStripConcat(inputs, planOp, rewriter);
    if (failed(lowered))
      return planOp.emitOpError("failed to lower selected row-strip Spatial concat plan");
    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
      return failure();
    return success();
  }
  return planOp.emitOpError("dense Concat plan was not lowered by the selected-plan patterns");
}

struct LowerSelectedConvPlan final : OpRewritePattern<spatial::SpatConv2DPlanOp> {
  explicit LowerSelectedConvPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
      : OpRewritePattern<spatial::SpatConv2DPlanOp>(ctx), target(target) {}

  LogicalResult matchAndRewrite(spatial::SpatConv2DPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    if (isDenseSelected(planOp.getOperation())) {
      FailureOr<Value> lowered = lowerSelectedConv2DPlan(
        planOp, std::nullopt, /*emitRowStripLayout=*/false, target, rewriter);
      if (failed(lowered))
        return planOp.emitOpError("failed to lower selected dense Spatial Conv plan");
      rewriter.replaceOp(planOp, *lowered);
      return success();
    }
    if (!isRowStripSelected(planOp.getOperation()))
      return failure();

    FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(planOp.getInput());
    if (failed(rowStripInput)
        && getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
      return failure();
    std::optional<Value> physicalInput;
    if (succeeded(rowStripInput))
      physicalInput = rowStripInput->storage;
    FailureOr<Value> lowered = lowerSelectedConv2DPlan(
      planOp, physicalInput, /*emitRowStripLayout=*/true, target, rewriter);
    if (failed(lowered))
      return planOp.emitOpError("failed to lower selected row-strip Spatial Conv plan");
    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
      return failure();
    return success();
  }

  const spatial::SpatialTargetInfo& target;
};

struct LowerRowStripReluPlan final : OpRewritePattern<spatial::SpatReluPlanOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    if (!isRowStripSelected(planOp.getOperation()))
      return failure();
    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
    if (failed(input)) {
      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
        return failure();
      return planOp.emitOpError("selected row-strip ReLU plan requires a row-strip input");
    }
    FailureOr<Value> lowered = lowerRowStripRelu(*input, planOp, rewriter);
    if (failed(lowered))
      return planOp.emitOpError("failed to lower selected row-strip Spatial ReLU plan");
    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
      return failure();
    return success();
  }
};

struct LowerRowStripSiluPlan final : OpRewritePattern<spatial::SpatSiluPlanOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    if (!isRowStripSelected(planOp.getOperation()))
      return failure();
    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
    if (failed(input)) {
      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
        return failure();
      return planOp.emitOpError("selected row-strip SiLU plan requires a row-strip input");
    }
    FailureOr<Value> lowered = lowerRowStripSilu(*input, planOp, rewriter);
    if (failed(lowered))
      return planOp.emitOpError("failed to lower selected row-strip Spatial SiLU plan");
    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
      return failure();
    return success();
  }
};

struct LowerRowStripResizePlan final : OpRewritePattern<spatial::SpatResizeNearestPlanOp> {
  explicit LowerRowStripResizePlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
      : OpRewritePattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}

  LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    if (!isRowStripSelected(planOp.getOperation()))
      return failure();
    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
    if (failed(input)) {
      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
        return failure();
      return planOp.emitOpError("selected row-strip Resize plan requires a row-strip input");
    }
    FailureOr<Value> lowered = lowerSelectedResizeNearestPlan(planOp, input->storage, target, rewriter);
    if (failed(lowered))
      return planOp.emitOpError("failed to lower selected row-strip Resize plan");
    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
      return failure();
    return success();
  }

  const spatial::SpatialTargetInfo& target;
};

struct LowerDenseMaxPoolPlan final : OpRewritePattern<spatial::SpatMaxPool2DPlanOp> {
  explicit LowerDenseMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
      : OpRewritePattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(target) {}

  LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    if (!isDenseSelected(planOp.getOperation()))
      return failure();
    FailureOr<Value> lowered = lowerDenseMaxPool2DPlan(planOp, target, rewriter);
    if (failed(lowered))
      return planOp.emitOpError("failed to lower selected dense Spatial MaxPool plan");
    rewriter.replaceOp(planOp, *lowered);
    return success();
  }

  const spatial::SpatialTargetInfo& target;
};

struct LowerRowStripMaxPoolPlan final : OpRewritePattern<spatial::SpatMaxPool2DPlanOp> {
  explicit LowerRowStripMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
      : OpRewritePattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(target) {}

  LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    if (!isRowStripSelected(planOp.getOperation()))
      return failure();
    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
    if (failed(input)
        && getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
      return failure();
    std::optional<Value> physicalInput;
    if (succeeded(input))
      physicalInput = input->storage;
    FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(planOp, physicalInput, target, rewriter);
    if (failed(lowered))
      return planOp.emitOpError("failed to lower selected row-strip Spatial MaxPool plan");
    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
      return failure();
    return success();
  }

  const spatial::SpatialTargetInfo& target;
};

struct LowerRowStripGlobalAveragePoolPlan
    final : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp> {
  explicit LowerRowStripGlobalAveragePoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
      : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}

  LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    if (!isRowStripSelected(planOp.getOperation()))
      return failure();
    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
    if (failed(input)
        && getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
      return failure();
    std::optional<Value> physicalInput;
    if (succeeded(input))
      physicalInput = input->storage;
    FailureOr<Value> lowered = lowerSelectedGlobalAveragePoolPlan(planOp, physicalInput, target, rewriter);
    if (failed(lowered))
      return planOp.emitOpError("failed to lower selected row-strip Spatial global AveragePool plan");
    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
      return failure();
    return success();
  }

  const spatial::SpatialTargetInfo& target;
};

struct LowerDenseGlobalAveragePoolPlan
    final : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp> {
  explicit LowerDenseGlobalAveragePoolPlan(MLIRContext* ctx,
                                           const spatial::SpatialTargetInfo& target)
      : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}

  LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    if (!isDenseSelected(planOp.getOperation()))
      return failure();
    FailureOr<Value> lowered = lowerDenseGlobalAveragePoolPlan(planOp, target, rewriter);
    if (failed(lowered))
      return planOp.emitOpError("failed to lower selected dense Spatial global AveragePool plan");
    rewriter.replaceOp(planOp, *lowered);
    return success();
  }

  const spatial::SpatialTargetInfo& target;
};

struct LowerRowStripBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    if (!isRowStripSelected(planOp.getOperation()))
      return failure();
    FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
    if (failed(input)) {
      if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
        return failure();
      return planOp.emitOpError("selected row-strip bias_add plan requires a row-strip input");
    }
    FailureOr<Value> lowered = lowerRowStripBiasAdd(*input, planOp, rewriter);
    if (failed(lowered))
      return planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan");
    if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
      return failure();
    return success();
  }
};

struct LowerRowStripAddPlan final : OpRewritePattern<spatial::SpatAddPlanOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    if (!isRowStripSelected(planOp.getOperation()))
      return failure();
    return lowerAddPlan(planOp, rewriter);
  }
};

struct LowerRowStripConcatPlan final : OpRewritePattern<spatial::SpatConcatPlanOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp,
                                PatternRewriter& rewriter) const override {
    if (!isRowStripSelected(planOp.getOperation()))
      return failure();
    return lowerConcatPlan(planOp, rewriter);
  }
};

struct LowerMaterializeLayout final
    : OpRewritePattern<spatial::SpatMaterializeLayoutOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(spatial::SpatMaterializeLayoutOp materializeOp,
                                PatternRewriter& rewriter) const override {
    auto source = materializeOp.getSourcePhysicalLayout();
    auto target = materializeOp.getTargetPhysicalLayout();
    if (source == spatial::PhysicalLayout::DenseNCHW
        && target == spatial::PhysicalLayout::DenseNCHW) {
      rewriter.replaceOp(materializeOp, materializeOp.getInput());
      return success();
    }
    if (source == spatial::PhysicalLayout::DenseNCHW
        && target == spatial::PhysicalLayout::NHWCRowStrip) {
      auto logicalType = dyn_cast<RankedTensorType>(materializeOp.getInput().getType());
      if (!logicalType)
        return materializeOp.emitOpError("requires a ranked dense input"), failure();
      FailureOr<Value> rowStrip = materializeDenseToRowStrip(
        materializeOp.getInput(), logicalType, materializeOp.getLoc(), rewriter);
      if (failed(rowStrip))
        return materializeOp.emitOpError(
          "failed to materialize dense NCHW storage to row-strip layout"), failure();
      rewriter.replaceOp(materializeOp, *rowStrip);
      return success();
    }
    if (source != spatial::PhysicalLayout::NHWCRowStrip
        || target != spatial::PhysicalLayout::DenseNCHW)
      return materializeOp.emitOpError(
          "unsupported Spatial layout materialization direction"), failure();
    auto inputType = dyn_cast<RankedTensorType>(materializeOp.getInput().getType());
    if (!inputType)
      return materializeOp.emitOpError("requires a ranked row-strip input"), failure();
    FailureOr<RowStripPhysicalValue> rowStripValue =
        getRowStripValue(materializeOp.getInput());
    if (failed(rowStripValue))
      return materializeOp.emitOpError(
          "requires an explicitly defining row-strip physical value"), failure();
    FailureOr<Value> dense = materializeRowStripToDense(
        *rowStripValue, materializeOp.getLoc(), rewriter);
    if (failed(dense))
      return materializeOp.emitOpError(
          "failed to materialize row-strip storage to dense NCHW"), failure();
    rewriter.replaceOp(materializeOp, *dense);
    return success();
  }
};

struct LowerRowStripFlatten final
    : OpRewritePattern<spatial::SpatGraphCompute> {
  explicit LowerRowStripFlatten(MLIRContext* context,
                                const spatial::SpatialTargetInfo& target)
      : OpRewritePattern<spatial::SpatGraphCompute>(context), target(target) {}

  LogicalResult matchAndRewrite(spatial::SpatGraphCompute flattenOp,
                                PatternRewriter& rewriter) const override {
    if (flattenOp.getInputs().size() != 1)
      return failure();
    FailureOr<RowStripPhysicalValue> input =
        getRowStripValue(flattenOp.getInputs().front());
    if (failed(input) || failed(canLowerFlattenFromRowStrip(flattenOp, target)))
      return failure();
    if (failed(lowerFlattenFromRowStrip(*input, flattenOp, target, rewriter)))
      return flattenOp.emitOpError(
          "failed to preserve row-strip layout through Flatten"), failure();
    return success();
  }

  const spatial::SpatialTargetInfo& target;
};

struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass)

  StringRef getArgument() const override { return "lower-spatial-plans"; }
  StringRef getDescription() const override { return "Lower selected Spatial planning ops to low-level Spatial IR."; }

  LowerSpatialPlansPass() = default;
  explicit LowerSpatialPlansPass(const spatial::SpatialTargetInfo& target)
      : target(target), hasTarget(true) {}

  void runOnOperation() override {
    ModuleOp moduleOp = getOperation();
    if (!hasTarget) {
      moduleOp.emitError("Spatial plan lowering requires an injected SpatialTargetInfo");
      signalPassFailure();
      return;
    }
    MLIRContext* ctx = moduleOp.getContext();
    auto entryFunc = getPimEntryFunc(moduleOp);
    if (failed(entryFunc)) {
      moduleOp.emitError("failed to locate the PIM entry function during LowerSpatialPlans");
      signalPassFailure();
      return;
    }
    func::FuncOp funcOp = *entryFunc;
    PatternRewriter rewriter(ctx);
    auto verifyLogicalPhase = [&](StringRef stage) -> bool {
      if (succeeded(verifyLogicalSpatialGraphInvariants(*entryFunc)))
        return true;
      moduleOp.emitError() << "logical Spatial graph verification failed " << stage;
      signalPassFailure();
      return false;
    };

    if (!verifyLogicalPhase("at the start of LowerSpatialPlans"))
      return;
    if (failed(verifySelectedLayouts(funcOp, target))) {
      moduleOp.emitError("selected Spatial layout verification failed");
      signalPassFailure();
      return;
    }

    RewritePatternSet selectedPlanPatterns(ctx);
    selectedPlanPatterns.add<LowerDenseReluPlan,
                             LowerRowStripReluPlan,
                             LowerDenseSiluPlan,
                             LowerRowStripSiluPlan,
                             LowerDenseBiasAddPlan,
                             LowerRowStripBiasAddPlan,
                             LowerDenseAddPlan,
                             LowerRowStripAddPlan,
                             LowerDenseConcatPlan,
                             LowerRowStripConcatPlan>(ctx);
    selectedPlanPatterns.add<LowerSelectedConvPlan,
                             LowerDenseResizePlan,
                             LowerRowStripResizePlan,
                             LowerDenseMaxPoolPlan,
                             LowerRowStripMaxPoolPlan,
                             LowerDenseGlobalAveragePoolPlan,
                             LowerRowStripGlobalAveragePoolPlan>(ctx, target);
    if (failed(applyPatternsGreedily(funcOp, std::move(selectedPlanPatterns)))) {
      moduleOp.emitError("failed to lower selected Spatial plans");
      signalPassFailure();
      return;
    }

    RewritePatternSet layoutPatterns(ctx);
    layoutPatterns.add<LowerMaterializeLayout>(ctx);
    layoutPatterns.add<LowerRowStripFlatten>(ctx, target);
    ConversionTarget layoutTarget(*ctx);
    layoutTarget.addLegalDialect<spatial::SpatialDialect,
                                 tensor::TensorDialect,
                                 linalg::LinalgDialect,
                                 affine::AffineDialect,
                                 arith::ArithDialect,
                                 scf::SCFDialect,
                                 func::FuncDialect>();
    layoutTarget.addIllegalDialect<ONNXDialect>();
    layoutTarget.addIllegalOp<spatial::SpatMaterializeLayoutOp>();
    layoutTarget.addDynamicallyLegalOp<spatial::SpatGraphCompute>(
        [&](spatial::SpatGraphCompute computeOp) {
          if (computeOp.getInputs().size() != 1)
            return true;
          FailureOr<RowStripPhysicalValue> input =
              getRowStripValue(computeOp.getInputs().front());
          return failed(input) || failed(canLowerFlattenFromRowStrip(computeOp, target));
        });
    FrozenRewritePatternSet frozenLayoutPatterns(std::move(layoutPatterns));
    if (failed(applyFullConversion(funcOp, layoutTarget,
                                   frozenLayoutPatterns))) {
      moduleOp.emitError("failed to lower explicit Spatial layout materialization");
      signalPassFailure();
      return;
    }

    if (!verifyLogicalPhase("after selected-plan conversion"))
      return;
    SmallVector<spatial::SpatBlueprintOp> deadPhysicalViews;
    funcOp.walk([&](spatial::SpatBlueprintOp blueprint) {
      if (spatial::isPhysicalView(blueprint.getMode()) && blueprint.use_empty())
        deadPhysicalViews.push_back(blueprint);
    });
    for (spatial::SpatBlueprintOp blueprint : deadPhysicalViews)
      rewriter.eraseOp(blueprint);
    bool hasIllegalOps = false;
    moduleOp.walk([&](Operation* op) {
      if (isa<ONNXEntryPointOp>(op))
        return;
      if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
        if (spatial::isFragmentAssembly(blueprint.getMode()))
          return;
        op->emitOpError("planning blueprint must not remain after LowerSpatialPlans");
        hasIllegalOps = true;
      }
      else if (isa<spatial::SpatConv2DPlanOp,
                   spatial::SpatBiasAddPlanOp,
                   spatial::SpatAddPlanOp,
                   spatial::SpatReluPlanOp,
                   spatial::SpatSiluPlanOp,
                   spatial::SpatResizeNearestPlanOp,
                   spatial::SpatMaxPool2DPlanOp,
                   spatial::SpatGlobalAveragePoolPlanOp,
                   spatial::SpatMaterializeLayoutOp>(op)
               || op->getDialect()->getNamespace() == "onnx") {
        op->emitOpError("operation must not remain after LowerSpatialPlans");
        hasIllegalOps = true;
      }
    });

    PassManager canonicalizationPM(ctx);
    canonicalizationPM.addPass(createCanonicalizerPass());
    if (failed(canonicalizationPM.run(moduleOp)))
      moduleOp.emitWarning("failed to run LowerSpatialPlansPass canonicalization; continuing");

    if (hasIllegalOps) {
      signalPassFailure();
    } else {
      dumpModule(moduleOp, "spatial1_graph");
      spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage();
      if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Spatial1)
          && failed(spatial::exportSpatialDataflowCsvGraph(funcOp, "spatial1_graph"))) {
        signalPassFailure();
        return;
      }
    }

    if (!verifyLogicalPhase("at the end of LowerSpatialPlans"))
      return;
  }

  spatial::SpatialTargetInfo target;
  bool hasTarget = false;
};

} // namespace

std::unique_ptr<Pass> createLowerSpatialPlansPass() { return std::make_unique<LowerSpatialPlansPass>(); }

std::unique_ptr<Pass> createLowerSpatialPlansPass(const spatial::SpatialTargetInfo& target) {
  return std::make_unique<LowerSpatialPlansPass>(target);
}

} // namespace onnx_mlir