finally fast googlenet with correct latency artifacts for fair comparison

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -64,6 +64,22 @@ static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
   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())
@@ -122,6 +138,99 @@ static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTe
   return batch->getResult(0);
 }
 
+static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp,
+                                  llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
+                                  llvm::SmallPtrSetImpl<Operation*>& eraseAfterLowering,
+                                  PatternRewriter& rewriter) {
+  FailureOr<RowStripPhysicalValue> lhs = getRowStripValue(rowStripValues, planOp.getLhs());
+  FailureOr<RowStripPhysicalValue> rhs = getRowStripValue(rowStripValues, planOp.getRhs());
+  if (succeeded(lhs) && succeeded(rhs)) {
+    auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
+      auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
+      return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
+    });
+    if (outputBlueprint == planOp.getResult().getUsers().end())
+      return planOp.emitOpError("row-strip add plan requires a row-strip blueprint result");
+
+    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");
+    auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
+    FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
+    if (failed(output))
+      return failure();
+    rowStripValues[blueprint.getResult()] = *output;
+    eraseAfterLowering.insert(planOp);
+    eraseAfterLowering.insert(blueprint);
+    return success();
+  }
+
+  rewriter.setInsertionPoint(planOp);
+  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();
+}
+
+static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp,
+                                     llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
+                                     llvm::SmallPtrSetImpl<Operation*>& eraseAfterLowering,
+                                     PatternRewriter& rewriter) {
+  SmallVector<RowStripPhysicalValue> inputs;
+  for (Value input : planOp.getInputs()) {
+    FailureOr<RowStripPhysicalValue> physical = getRowStripValue(rowStripValues, input);
+    if (failed(physical)) {
+      inputs.clear();
+      break;
+    }
+    inputs.push_back(*physical);
+  }
+  if (inputs.size() == planOp.getInputs().size()) {
+    auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
+      auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
+      return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
+    });
+    if (outputBlueprint == planOp.getResult().getUsers().end())
+      return planOp.emitOpError("row-strip concat plan requires a row-strip blueprint result");
+    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");
+    auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
+    FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
+    if (failed(output))
+      return failure();
+    rowStripValues[blueprint.getResult()] = *output;
+    eraseAfterLowering.insert(planOp);
+    eraseAfterLowering.insert(blueprint);
+    return success();
+  }
+
+  rewriter.setInsertionPoint(planOp);
+  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();
+}
+
 struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass)
 
@@ -274,6 +383,40 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         eraseAfterLowering.insert(blueprint);
         continue;
       }
+      if (auto planOp = dyn_cast<spatial::SpatGlobalAveragePoolPlanOp>(&op)) {
+        auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
+          auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
+          return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
+        });
+        if (outputBlueprint == planOp.getResult().getUsers().end()) {
+          planOp.emitOpError("selected global AveragePool plan requires a row-strip blueprint result");
+          signalPassFailure();
+          return;
+        }
+
+        FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
+        rewriter.setInsertionPoint(planOp);
+        std::optional<Value> physicalInput;
+        if (succeeded(input))
+          physicalInput = input->storage;
+        FailureOr<Value> lowered =
+          lowerSelectedGlobalAveragePoolPlan(planOp, physicalInput, rewriter);
+        if (failed(lowered)) {
+          planOp.emitOpError("failed to lower selected row-strip Spatial global AveragePool plan");
+          signalPassFailure();
+          return;
+        }
+        auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
+        FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
+        if (failed(output)) {
+          signalPassFailure();
+          return;
+        }
+        rowStripValues[blueprint.getResult()] = *output;
+        eraseAfterLowering.insert(planOp);
+        eraseAfterLowering.insert(blueprint);
+        continue;
+      }
       if (auto planOp = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
         if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
           auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
@@ -339,6 +482,20 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         rewriter.replaceOp(planOp, computeOp.getResults());
         continue;
       }
+      if (auto planOp = dyn_cast<spatial::SpatAddPlanOp>(&op)) {
+        if (failed(lowerAddPlan(planOp, rowStripValues, eraseAfterLowering, rewriter))) {
+          signalPassFailure();
+          return;
+        }
+        continue;
+      }
+      if (auto planOp = dyn_cast<spatial::SpatConcatPlanOp>(&op)) {
+        if (failed(lowerConcatPlan(planOp, rowStripValues, eraseAfterLowering, rewriter))) {
+          signalPassFailure();
+          return;
+        }
+        continue;
+      }
       if (auto flattenOp = dyn_cast<spatial::SpatGraphCompute>(&op)) {
         if (flattenOp.getInputs().size() == 1) {
           FailureOr<RowStripPhysicalValue> input =
@@ -488,12 +645,15 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
           return;
         op->emitOpError("planning blueprint must not remain after LowerSpatialPlans");
         hasIllegalOps = true;
-      } else if (isa<spatial::SpatConv2DPlanOp,
-                     spatial::SpatBiasAddPlanOp,
-                     spatial::SpatReluPlanOp,
-                     spatial::SpatMaxPool2DPlanOp,
-                     spatial::SpatMaterializeLayoutOp>(op)
-          || op->getDialect()->getNamespace() == "onnx") {
+      }
+      else if (isa<spatial::SpatConv2DPlanOp,
+                   spatial::SpatBiasAddPlanOp,
+                   spatial::SpatAddPlanOp,
+                   spatial::SpatReluPlanOp,
+                   spatial::SpatMaxPool2DPlanOp,
+                   spatial::SpatGlobalAveragePoolPlanOp,
+                   spatial::SpatMaterializeLayoutOp>(op)
+               || op->getDialect()->getNamespace() == "onnx") {
         op->emitOpError("operation must not remain after LowerSpatialPlans");
         hasIllegalOps = true;
       }