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

#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"

#include "llvm/ADT/DenseMap.h"

#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Accelerators/PIM/Pass/PIMPasses.h"

using namespace mlir;

namespace onnx_mlir {
namespace {

static constexpr StringLiteral kLogicalLayout = "nchw";
static constexpr StringLiteral kDenseLayout = "dense_nchw";
static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";

enum class SelectedLayout {
  DenseNchw,
  NchwRowStrip,
};

static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& layouts, Value value) {
  auto it = layouts.find(value);
  return it == layouts.end() ? SelectedLayout::DenseNchw : it->second;
}

static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, SelectedLayout>& layouts) {
  if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(user))
    return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::NchwRowStrip;
  if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user))
    return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::NchwRowStrip;
  if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
    return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::NchwRowStrip;
  return false;
}

static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap<Value, SelectedLayout>& layouts) {
  for (Operation* user : value.getUsers()) {
    if (usesSelectedRowStrip(user, layouts))
      continue;
    // Dense-only users must be materialized explicitly.
    continue;
  }
  return true;
}

static bool canConsumeRowStripAsUser(Operation* user) {
  if (isa<spatial::SpatReluPlanOp>(user))
    return true;
  if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user)) {
    auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
    return resultType && isSupportedBiasAddValue(biasAddPlan.getBias(), resultType);
  }
  if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
    return succeeded(canConsumeAndProduceRowStrip(convPlan));
  return false;
}

static bool hasRowStripConsumer(Value value) {
  for (Operation* user : value.getUsers())
    if (canConsumeRowStripAsUser(user))
      return true;
  return false;
}


static bool canSelectConvRowStrip(spatial::SpatConv2DPlanOp convPlan,
                                  llvm::DenseMap<Value, SelectedLayout>& layouts) {
  SelectedLayout inputLayout = getSelectedLayout(layouts, convPlan.getInput());
  if (inputLayout == SelectedLayout::NchwRowStrip)
    return succeeded(canConsumeAndProduceRowStrip(convPlan));
  return succeeded(canLowerConvPlanToRowStrip(convPlan));
}

static SelectedLayout chooseConvLayout(spatial::SpatConv2DPlanOp convPlan,
                                       llvm::DenseMap<Value, SelectedLayout>& layouts) {
  if (!canSelectConvRowStrip(convPlan, layouts))
    return SelectedLayout::DenseNchw;
  if (getSelectedLayout(layouts, convPlan.getInput()) != SelectedLayout::NchwRowStrip
      && !hasRowStripConsumer(convPlan.getResult()))
    return SelectedLayout::DenseNchw;
  if (!allUsersCanHandleRowStrip(convPlan.getResult(), layouts))
    return SelectedLayout::DenseNchw;
  return SelectedLayout::NchwRowStrip;
}

static SelectedLayout chooseReluLayout(spatial::SpatReluPlanOp reluPlan,
                                       llvm::DenseMap<Value, SelectedLayout>& layouts) {
  if (getSelectedLayout(layouts, reluPlan.getInput()) != SelectedLayout::NchwRowStrip)
    return SelectedLayout::DenseNchw;
  if (!hasRowStripConsumer(reluPlan.getResult()))
    return SelectedLayout::DenseNchw;
  if (!allUsersCanHandleRowStrip(reluPlan.getResult(), layouts))
    return SelectedLayout::DenseNchw;
  return SelectedLayout::NchwRowStrip;
}

static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan,
                                          llvm::DenseMap<Value, SelectedLayout>& layouts) {
  if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::NchwRowStrip)
    return SelectedLayout::DenseNchw;
  auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
  if (!resultType || !isSupportedBiasAddValue(biasAddPlan.getBias(), resultType))
    return SelectedLayout::DenseNchw;
  if (!hasRowStripConsumer(biasAddPlan.getResult()))
    return SelectedLayout::DenseNchw;
  if (!allUsersCanHandleRowStrip(biasAddPlan.getResult(), layouts))
    return SelectedLayout::DenseNchw;
  return SelectedLayout::NchwRowStrip;
}

static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Value value) {
  auto outputType = cast<RankedTensorType>(value.getType());
  auto [offsets, sizes] = buildRowStripMetadata(outputType);
  return spatial::SpatBlueprintOp::create(rewriter,
                                              value.getLoc(),
                                              outputType,
                                              value,
                                              ValueRange {},
                                              rewriter.getStringAttr(kLogicalLayout),
                                              rewriter.getStringAttr(kRowStripLayout),
                                              rewriter.getDenseI64ArrayAttr(offsets),
                                              rewriter.getDenseI64ArrayAttr(sizes),
                                              rewriter.getStringAttr(kRowStripIndexMap),
                                              nullptr,
                                              nullptr,
                                              nullptr,
                                              nullptr,
                                              nullptr,
                                              nullptr);
}

static void materializeDenseUses(IRRewriter& rewriter,
                                 Value layoutValue,
                                 llvm::DenseMap<Value, SelectedLayout>& layouts) {
  SmallVector<OpOperand*> denseUses;
  for (OpOperand& use : layoutValue.getUses()) {
    if (usesSelectedRowStrip(use.getOwner(), layouts))
      continue;
    denseUses.push_back(&use);
  }

  for (OpOperand* use : denseUses) {
    Operation* owner = use->getOwner();
    rewriter.setInsertionPoint(owner);
    auto materialized = spatial::SpatMaterializeLayoutOp::create(rewriter,
                                                                 owner->getLoc(),
                                                                 use->get().getType(),
                                                                 use->get(),
                                                                 rewriter.getStringAttr(kLogicalLayout),
                                                                 rewriter.getStringAttr(kRowStripLayout),
                                                                 rewriter.getStringAttr(kDenseLayout));
    use->set(materialized.getResult());
  }
}

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

  StringRef getArgument() const override { return "spatial-layout-planning"; }
  StringRef getDescription() const override { return "Select conservative Spatial layouts and insert reconciliation barriers."; }

  void runOnOperation() override {
    auto entryFunc = getPimEntryFunc(getOperation());
    if (failed(entryFunc)) {
      getOperation().emitError("failed to locate the PIM entry function during Spatial layout planning");
      signalPassFailure();
      return;
    }

    func::FuncOp funcOp = *entryFunc;
    IRRewriter rewriter(&getContext());
    llvm::DenseMap<Value, SelectedLayout> layouts;

    bool changed = true;
    while (changed) {
      changed = false;
      for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
        if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(&op)) {
          SelectedLayout selected = chooseConvLayout(convPlan, layouts);
          if (layouts[convPlan.getResult()] != selected) {
            layouts[convPlan.getResult()] = selected;
            changed = true;
          }
          continue;
        }
        if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op)) {
          SelectedLayout selected = chooseReluLayout(reluPlan, layouts);
          if (layouts[reluPlan.getResult()] != selected) {
            layouts[reluPlan.getResult()] = selected;
            changed = true;
          }
          continue;
        }
        if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
          SelectedLayout selected = chooseBiasAddLayout(biasAddPlan, layouts);
          if (layouts[biasAddPlan.getResult()] != selected) {
            layouts[biasAddPlan.getResult()] = selected;
            changed = true;
          }
          continue;
        }
      }
    }

    for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
      Value producedValue;
      if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(&op))
        producedValue = convPlan.getResult();
      else if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op))
        producedValue = biasAddPlan.getResult();
      else if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op))
        producedValue = reluPlan.getResult();
      else
        continue;

      if (getSelectedLayout(layouts, producedValue) != SelectedLayout::NchwRowStrip)
        continue;

      rewriter.setInsertionPointAfter(&op);
      auto blueprint = insertRowStripBlueprint(rewriter, producedValue);
      rewriter.replaceAllUsesExcept(producedValue, blueprint.getResult(), blueprint);
      materializeDenseUses(rewriter, blueprint.getResult(), layouts);
    }
    if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
      getOperation().emitError("logical Spatial graph verification failed after SpatialLayoutPlanning");
      signalPassFailure();
    }
  }
};

} // namespace

std::unique_ptr<Pass> createSpatialLayoutPlanningPass() { return std::make_unique<SpatialLayoutPlanningPass>(); }

} // namespace onnx_mlir