Raptor sync wait

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