finally fast googlenet with correct latency artifacts for fair comparison

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

@@ -19,9 +19,11 @@ FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, Ra
       || storageType.getRank() != 5 || logicalType.getRank() != 4 || logicalType.getDimSize(0) != 1
       || storageType.getElementType() != logicalType.getElementType()
       || storageType.getDimSize(1) != 1 || storageType.getDimSize(2) != 1
-      || storageType.getDimSize(3) != logicalType.getDimSize(3) || storageType.getDimSize(4) <= 0)
+      || storageType.getDimSize(3) != logicalType.getDimSize(3)
+      || storageType.getDimSize(4) <= 0)
     return failure();
-  const int64_t tilesPerRow = ceilIntegerDivide(logicalType.getDimSize(1), storageType.getDimSize(4));
+  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,
@@ -249,4 +251,111 @@ FailureOr<Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
   return batchOp->getResult(0);
 }
 
+FailureOr<Value> applyRowStripAdd(const RowStripPhysicalValue& lhs,
+                                  const RowStripPhysicalValue& rhs,
+                                  PatternRewriter& rewriter,
+                                  Location loc) {
+  if (lhs.logicalType != rhs.logicalType || lhs.fragmentType != rhs.fragmentType
+      || lhs.storage.getType() != rhs.storage.getType() || lhs.tilesPerRow != rhs.tilesPerRow)
+    return failure();
+  auto storageType = cast<RankedTensorType>(lhs.storage.getType());
+  const int64_t laneCount = storageType.getDimSize(0);
+  auto batch = createSpatComputeBatch(
+    rewriter,
+    loc,
+    TypeRange {storageType},
+    laneCount,
+    {},
+    ValueRange {lhs.storage, rhs.storage},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      FailureOr<Value> lhsFragment =
+        extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, lhs.fragmentType);
+      FailureOr<Value> rhsFragment =
+        extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[1], args.lane, rhs.fragmentType);
+      if (failed(lhsFragment) || failed(rhsFragment))
+        return failure();
+      Value added = spatial::SpatVAddOp::create(rewriter, loc, lhs.fragmentType, *lhsFragment, *rhsFragment);
+      publishGraphBatchPhysicalFragment(rewriter, loc, added, args.outputs.front(), args.lane);
+      return success();
+    });
+  if (failed(batch))
+    return failure();
+  return batch->getResult(0);
+}
+
+FailureOr<Value> applyRowStripConcat(ArrayRef<RowStripPhysicalValue> inputs,
+                                     RankedTensorType outputType,
+                                     PatternRewriter& rewriter,
+                                     Location loc) {
+  if (inputs.empty() || !outputType || !outputType.hasStaticShape() || outputType.getRank() != 4
+      || outputType.getDimSize(0) != 1)
+    return failure();
+  int64_t channels = 0;
+  for (const RowStripPhysicalValue& input : inputs) {
+    if (input.logicalType.getElementType() != outputType.getElementType()
+        || input.logicalType.getDimSize(0) != outputType.getDimSize(0)
+        || input.logicalType.getDimSize(2) != outputType.getDimSize(2)
+        || input.logicalType.getDimSize(3) != outputType.getDimSize(3))
+      return failure();
+    channels += input.logicalType.getDimSize(1);
+  }
+  if (channels != outputType.getDimSize(1))
+    return failure();
+
+  SmallVector<Value> storages;
+  llvm::transform(
+    inputs, std::back_inserter(storages), [](const RowStripPhysicalValue& input) { return input.storage; });
+  const int64_t tileWidth = outputType.getDimSize(3);
+  auto fragmentType = getRowStripFragmentType(outputType);
+  auto storageType = getRowStripStorageType(outputType);
+  auto batch = createSpatComputeBatch(
+    rewriter,
+    loc,
+    TypeRange {storageType},
+    outputType.getDimSize(2),
+    {},
+    storages,
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      Operation* anchor = rewriter.getInsertionBlock()->getParentOp();
+      SmallVector<Value> fragments;
+      for (auto [inputIndex, input] : llvm::enumerate(inputs)) {
+        Value tileStart = affineMulConst(
+          rewriter, loc, args.lane, input.tilesPerRow, anchor);
+        for (int64_t tile = 0; tile < input.tilesPerRow; ++tile) {
+          Value slot =
+            affineAddConst(rewriter, loc, tileStart, tile, anchor);
+          FailureOr<Value> fragment =
+            extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[inputIndex], slot, input.fragmentType);
+          if (failed(fragment))
+            return failure();
+          int64_t channelOffset = tile * input.fragmentType.getDimSize(3);
+          int64_t validChannels =
+            std::min(input.fragmentType.getDimSize(3), input.logicalType.getDimSize(1) - channelOffset);
+          auto validType =
+            RankedTensorType::get(
+              {1, 1, tileWidth, validChannels},
+              outputType.getElementType());
+          MixedSliceGeometry slice;
+          slice.offsets.assign(4, rewriter.getIndexAttr(0));
+          slice.sizes = {rewriter.getIndexAttr(1),
+                         rewriter.getIndexAttr(1),
+                         rewriter.getIndexAttr(tileWidth),
+                         rewriter.getIndexAttr(validChannels)};
+          slice.strides.assign(4, rewriter.getIndexAttr(1));
+          Value valid = extractMixedSliceOrIdentity(rewriter, loc, *fragment, validType, slice);
+          if (!valid)
+            return failure();
+          fragments.push_back(valid);
+        }
+      }
+      Value concatenated =
+        spatial::SpatConcatOp::create(rewriter, loc, fragmentType, rewriter.getI64IntegerAttr(3), fragments);
+      publishGraphBatchPhysicalFragment(rewriter, loc, concatenated, args.outputs.front(), args.lane);
+      return success();
+    });
+  if (failed(batch))
+    return failure();
+  return batch->getResult(0);
+}
+
 } // namespace onnx_mlir

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

@@ -66,4 +66,14 @@ mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(const RowStripPhysicalValue& v
                                                   mlir::PatternRewriter& rewriter,
                                                   mlir::Location loc);
 
+mlir::FailureOr<mlir::Value> applyRowStripAdd(const RowStripPhysicalValue& lhs,
+                                              const RowStripPhysicalValue& rhs,
+                                              mlir::PatternRewriter& rewriter,
+                                              mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> applyRowStripConcat(llvm::ArrayRef<RowStripPhysicalValue> inputs,
+                                                 mlir::RankedTensorType outputType,
+                                                 mlir::PatternRewriter& rewriter,
+                                                 mlir::Location loc);
+
 } // namespace onnx_mlir