Resnet is fast

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -38,6 +38,8 @@ static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, Selected
     return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::NchwRowStrip;
   if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
     return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::NchwRowStrip;
+  if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
+    return getSelectedLayout(layouts, maxPoolPlan.getResult()) == SelectedLayout::NchwRowStrip;
   return false;
 }
 
@@ -60,6 +62,8 @@ static bool canConsumeRowStripAsUser(Operation* user) {
   }
   if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
     return succeeded(canConsumeAndProduceRowStrip(convPlan));
+  if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(user))
+    return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan));
   return false;
 }
 
@@ -70,7 +74,6 @@ static bool hasRowStripConsumer(Value value) {
   return false;
 }
 
-
 static bool canSelectConvRowStrip(spatial::SpatConv2DPlanOp convPlan,
                                   llvm::DenseMap<Value, SelectedLayout>& layouts) {
   SelectedLayout inputLayout = getSelectedLayout(layouts, convPlan.getInput());
@@ -83,9 +86,6 @@ 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;
@@ -116,6 +116,11 @@ static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan
   return SelectedLayout::NchwRowStrip;
 }
 
+static SelectedLayout chooseMaxPoolLayout(spatial::SpatMaxPool2DPlanOp maxPoolPlan) {
+  return succeeded(canLowerMaxPoolPlanToRowStrip(maxPoolPlan)) ? SelectedLayout::NchwRowStrip
+                                                               : SelectedLayout::DenseNchw;
+}
+
 static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Value value) {
   auto outputType = cast<RankedTensorType>(value.getType());
   auto [offsets, sizes] = buildRowStripMetadata(outputType);
@@ -208,6 +213,14 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
           }
           continue;
         }
+        if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op)) {
+          SelectedLayout selected = chooseMaxPoolLayout(maxPoolPlan);
+          if (layouts[maxPoolPlan.getResult()] != selected) {
+            layouts[maxPoolPlan.getResult()] = selected;
+            changed = true;
+          }
+          continue;
+        }
       }
     }
 
@@ -219,6 +232,8 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
         producedValue = biasAddPlan.getResult();
       else if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op))
         producedValue = reluPlan.getResult();
+      else if (auto maxPoolPlan = dyn_cast<spatial::SpatMaxPool2DPlanOp>(&op))
+        producedValue = maxPoolPlan.getResult();
       else
         continue;