relu_conv_relu Faster on Arch-A

src/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.cpp

@@ -18,10 +18,10 @@ FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, Ra
   if (!storageType || !storageType.hasStaticShape() || !logicalType || !logicalType.hasStaticShape()
       || storageType.getRank() != 5 || logicalType.getRank() != 4 || logicalType.getDimSize(0) != 1
       || storageType.getElementType() != logicalType.getElementType()
-      || storageType.getDimSize(1) != 1 || storageType.getDimSize(2) <= 0
-      || storageType.getDimSize(3) != 1 || storageType.getDimSize(4) != logicalType.getDimSize(3))
+      || storageType.getDimSize(1) != 1 || storageType.getDimSize(2) != 1
+      || storageType.getDimSize(3) != logicalType.getDimSize(3) || storageType.getDimSize(4) <= 0)
     return failure();
-  const int64_t tilesPerRow = ceilIntegerDivide(logicalType.getDimSize(1), storageType.getDimSize(2));
+  const int64_t tilesPerRow = ceilIntegerDivide(logicalType.getDimSize(1), storageType.getDimSize(4));
   if (storageType.getDimSize(0) != logicalType.getDimSize(2) * tilesPerRow)
     return failure();
   return RowStripPhysicalValue {storage, logicalType,
@@ -30,7 +30,8 @@ FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, Ra
 }
 
 RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
-  return RankedTensorType::get({logicalType.getDimSize(0), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
+  return RankedTensorType::get({logicalType.getDimSize(0), 1, logicalType.getDimSize(3),
+                                logicalType.getDimSize(1)},
                                logicalType.getElementType(),
                                logicalType.getEncoding());
 }
@@ -78,16 +79,18 @@ void insertRowStripFragment(Value fragment,
 FailureOr<Value> createPerChannelConstantFragment(DenseElementsAttr denseAttr,
                                                   RankedTensorType fragmentType,
                                                   PatternRewriter& rewriter) {
-  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, fragmentType);
+  auto logicalType = RankedTensorType::get(
+    {1, fragmentType.getDimSize(3), 1, 1}, fragmentType.getElementType());
+  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, logicalType);
   if (failed(channelValues))
     return failure();
 
   SmallVector<Attribute> values;
   values.reserve(fragmentType.getNumElements());
   for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n)
-    for (int64_t channel = 0; channel < fragmentType.getDimSize(1); ++channel)
-      for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h)
-        for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w)
+    for (int64_t h = 0; h < fragmentType.getDimSize(1); ++h)
+      for (int64_t w = 0; w < fragmentType.getDimSize(2); ++w)
+        for (int64_t channel = 0; channel < fragmentType.getDimSize(3); ++channel)
           values.push_back((*channelValues)[channel]);
 
   auto attr = DenseElementsAttr::get(fragmentType, values);
@@ -117,7 +120,6 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
     return failure();
 
   auto rowSliceType = RankedTensorType::get({width, channels}, logicalType.getElementType(), rowsType.getEncoding());
-  auto channelWidthType = RankedTensorType::get({channels, width}, logicalType.getElementType(), rowsType.getEncoding());
   auto fragmentType = getRowStripFragmentType(logicalType);
   auto storageType = getRowStripStorageType(logicalType);
   auto batchOp = createSpatComputeBatch(
@@ -128,10 +130,8 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
       SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(width), rewriter.getIndexAttr(channels)};
       Value rowSlice = tensor::ExtractSliceOp::create(
         rewriter, loc, rowSliceType, args.inputs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
-      Value channelWidth = ONNXTransposeOp::create(
-        rewriter, loc, channelWidthType, rowSlice, rewriter.getI64ArrayAttr({1, 0})).getResult();
       Value fragment = tensor::ExpandShapeOp::create(
-        rewriter, loc, fragmentType, channelWidth, SmallVector<ReassociationIndices> {{0, 1}, {2, 3}});
+        rewriter, loc, fragmentType, rowSlice, SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
       insertRowStripFragment(fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
       return success();
     });
@@ -143,8 +143,28 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
 FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
                                                  PatternRewriter& rewriter,
                                                  Location loc) {
+  const int64_t laneCount = cast<RankedTensorType>(value.storage.getType()).getDimSize(0);
+  const int64_t tileChannels = value.fragmentType.getDimSize(3);
+  auto nchwFragmentType = RankedTensorType::get(
+    {1, tileChannels, 1, value.logicalType.getDimSize(3)}, value.fragmentType.getElementType(),
+    value.fragmentType.getEncoding());
+  auto nchwStorageType = spatial::getGraphBatchPhysicalResultType(laneCount, nchwFragmentType);
+  auto transposed = createSpatComputeBatch(
+    rewriter, loc, TypeRange {nchwStorageType}, laneCount, {}, ValueRange {value.storage},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
+        rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
+      if (failed(fragment))
+        return failure();
+      Value nchw = ONNXTransposeOp::create(
+        rewriter, loc, nchwFragmentType, *fragment, rewriter.getI64ArrayAttr({0, 3, 1, 2}));
+      publishGraphBatchPhysicalFragment(rewriter, loc, nchw, args.outputs.front(), args.lane);
+      return success();
+    });
+  if (failed(transposed))
+    return failure();
+
   SmallVector<FragmentAssemblyEntry> entries;
-  const int64_t tileChannels = value.fragmentType.getDimSize(1);
   for (int64_t row = 0; row < value.logicalType.getDimSize(2); ++row)
     for (int64_t tile = 0; tile < value.tilesPerRow; ++tile) {
       const int64_t channelOffset = tile * tileChannels;
@@ -152,7 +172,7 @@ FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& va
         {1, std::min(tileChannels, value.logicalType.getDimSize(1) - channelOffset), 1,
          value.logicalType.getDimSize(3)}});
     }
-  return createFragmentAssemblyBlueprint(value.storage, value.logicalType, entries, "nchw_row_strip",
+  return createFragmentAssemblyBlueprint(transposed->getResult(0), value.logicalType, entries, "nhwc_row_strip",
                                          kRowStripIndexMap, rewriter, loc);
 }
 
@@ -193,11 +213,12 @@ FailureOr<Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
   auto biasStorageType = spatial::getGraphBatchPhysicalResultType(value.tilesPerRow, value.fragmentType);
   SmallVector<Attribute> biasValues(
     biasStorageType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(value.fragmentType.getElementType())));
-  const int64_t tileChannels = value.fragmentType.getDimSize(1);
-  const int64_t width = value.fragmentType.getDimSize(3);
+  const int64_t tileChannels = value.fragmentType.getDimSize(3);
+  const int64_t width = value.fragmentType.getDimSize(2);
   for (int64_t channel = 0; channel < value.logicalType.getDimSize(1); ++channel)
     for (int64_t w = 0; w < width; ++w)
-      biasValues[((channel / tileChannels) * tileChannels + channel % tileChannels) * width + w] =
+      biasValues[(channel / tileChannels) * width * tileChannels + w * tileChannels
+                 + channel % tileChannels] =
         (*channelValues)[channel];
   Value biasStorage = getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(),
     DenseElementsAttr::get(biasStorageType, biasValues), biasStorageType);