meno diamantini

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -6,10 +6,11 @@
 
 #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"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
 
 using namespace mlir;
 
@@ -19,7 +20,6 @@ namespace {
 static constexpr StringLiteral kLogicalLayout = "nchw";
 static constexpr StringLiteral kDenseLayout = "dense_nchw";
 static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";
-static constexpr StringLiteral kRowStripIndexMap = "packed_hwc_rows_to_nchw";
 
 enum class SelectedLayout {
   DenseNchw,
@@ -34,6 +34,8 @@ static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& l
 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;
@@ -49,21 +51,26 @@ static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap<Value, Selecte
   return true;
 }
 
-static std::pair<SmallVector<int64_t>, SmallVector<int64_t>> buildRowStripMetadata(RankedTensorType type) {
-  SmallVector<int64_t> offsets;
-  SmallVector<int64_t> sizes;
-  const int64_t channels = type.getDimSize(1);
-  const int64_t height = type.getDimSize(2);
-  const int64_t width = type.getDimSize(3);
-  offsets.reserve(height * 4);
-  sizes.reserve(height * 4);
-  for (int64_t row = 0; row < height; ++row) {
-    offsets.append({0, 0, row, 0});
-    sizes.append({1, channels, 1, width});
+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);
   }
-  return {offsets, sizes};
+  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());
@@ -76,6 +83,9 @@ 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;
@@ -85,11 +95,27 @@ 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);
@@ -173,6 +199,14 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
           }
           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;
+        }
       }
     }
 
@@ -180,6 +214,8 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
       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