fix vgg16 diamonds

updat ops validations

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

@@ -9,6 +9,53 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
+FailureOr<Value> createFragmentAssemblyBlueprint(Value physicalBatch,
+                                                 RankedTensorType logicalType,
+                                                 ArrayRef<FragmentAssemblyEntry> entries,
+                                                 StringRef physicalLayout,
+                                                 StringRef indexMap,
+                                                 PatternRewriter& rewriter,
+                                                 Location loc) {
+  auto physicalType = dyn_cast<RankedTensorType>(physicalBatch.getType());
+  if (!physicalType || !physicalType.hasStaticShape() || !logicalType || !logicalType.hasStaticShape()
+      || physicalType.getRank() != logicalType.getRank() + 1 || entries.empty())
+    return emitError(loc, "invalid static physical batch for fragment assembly"), failure();
+
+  const int64_t rank = logicalType.getRank();
+  const int64_t laneCount = physicalType.getDimSize(0);
+  if (laneCount <= 0)
+    return emitError(loc, "fragment assembly requires at least one physical source slot"), failure();
+  const int64_t fragmentElements = physicalType.getNumElements() / laneCount;
+  SmallVector<int64_t> operandIndices(entries.size(), 0), sourceSlots, sourceOffsets, offsets, sizes,
+    strides(entries.size() * rank, 1);
+  for (const FragmentAssemblyEntry& entry : entries) {
+    if (entry.sourceSlot < 0 || entry.sourceSlot >= laneCount || entry.sourceOffset < 0
+        || entry.destinationOffsets.size() != static_cast<size_t>(rank)
+        || entry.sizes.size() != static_cast<size_t>(rank))
+      return emitError(loc, "invalid fragment assembly entry"), failure();
+    int64_t entryElements = 1;
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      if (entry.destinationOffsets[dim] < 0 || entry.sizes[dim] <= 0
+          || entry.destinationOffsets[dim] + entry.sizes[dim] > logicalType.getDimSize(dim))
+        return emitError(loc, "fragment assembly entry exceeds the logical tensor"), failure();
+      entryElements *= entry.sizes[dim];
+    }
+    if (entry.sourceOffset + entryElements > fragmentElements)
+      return emitError(loc, "fragment assembly entry exceeds its physical source slot"), failure();
+    sourceSlots.push_back(entry.sourceSlot);
+    sourceOffsets.push_back(entry.sourceOffset);
+    llvm::append_range(offsets, entry.destinationOffsets);
+    llvm::append_range(sizes, entry.sizes);
+  }
+  return spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, physicalBatch, ValueRange {},
+    rewriter.getStringAttr("nchw"), rewriter.getStringAttr(physicalLayout),
+    rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes),
+    rewriter.getStringAttr(indexMap), rewriter.getStringAttr("fragment_assembly"),
+    rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots),
+    rewriter.getDenseI64ArrayAttr(sourceOffsets), rewriter.getDenseI64ArrayAttr(strides),
+    rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete")).getOutput();
+}
+
 Value sumTensors(ArrayRef<Value> tensors, PatternRewriter& rewriter) {
   if (tensors.size() == 1)
     return tensors[0];

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

@@ -19,6 +19,13 @@
 
 namespace onnx_mlir {
 
+struct FragmentAssemblyEntry {
+  int64_t sourceSlot;
+  int64_t sourceOffset;
+  llvm::SmallVector<int64_t, 4> destinationOffsets;
+  llvm::SmallVector<int64_t, 4> sizes;
+};
+
 namespace detail {
 
 inline mlir::ValueRange getBlockArgs(mlir::Block* block) { return mlir::ValueRange(block->getArguments()); }
@@ -407,4 +414,12 @@ mlir::Value materializeOrComputeUnary(mlir::Value input,
 
 mlir::Value sumTensors(mlir::ArrayRef<mlir::Value> tensors, mlir::PatternRewriter& rewriter);
 
+mlir::FailureOr<mlir::Value> createFragmentAssemblyBlueprint(mlir::Value physicalBatch,
+                                                             mlir::RankedTensorType logicalType,
+                                                             llvm::ArrayRef<FragmentAssemblyEntry> entries,
+                                                             llvm::StringRef physicalLayout,
+                                                             llvm::StringRef indexMap,
+                                                             mlir::PatternRewriter& rewriter,
+                                                             mlir::Location loc);
+
 } // namespace onnx_mlir

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

@@ -4,6 +4,7 @@
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -12,6 +13,22 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
+FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, RankedTensorType logicalType) {
+  auto storageType = dyn_cast<RankedTensorType>(storage.getType());
+  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))
+    return failure();
+  const int64_t tilesPerRow = ceilIntegerDivide(logicalType.getDimSize(1), storageType.getDimSize(2));
+  if (storageType.getDimSize(0) != logicalType.getDimSize(2) * tilesPerRow)
+    return failure();
+  return RowStripPhysicalValue {storage, logicalType,
+    RankedTensorType::get(storageType.getShape().drop_front(), storageType.getElementType(), storageType.getEncoding()),
+    tilesPerRow};
+}
+
 RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
   return RankedTensorType::get({logicalType.getDimSize(0), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
                                logicalType.getElementType(),
@@ -123,42 +140,39 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
   return batchOp->getResult(0);
 }
 
-FailureOr<Value>
-createRowStripAssemblyBlueprint(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
-  auto storageType = dyn_cast<RankedTensorType>(storage.getType());
-  if (!storageType || storageType != getRowStripStorageType(logicalType))
-    return failure();
-
-  auto [offsets, sizes] = buildRowStripMetadata(logicalType);
-  int64_t height = logicalType.getDimSize(2);
-  SmallVector<int64_t> operandIndices(height, 0), sourceSlots, sourceOffsets(height, 0), strides(height * 4, 1);
-  for (int64_t row = 0; row < height; ++row)
-    sourceSlots.push_back(row);
-  return spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, storage, ValueRange {},
-    rewriter.getStringAttr("nchw"), rewriter.getStringAttr("nchw_row_strip"),
-    rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes),
-    rewriter.getStringAttr("nchw_row_strip_fragments"), rewriter.getStringAttr("fragment_assembly"),
-    rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots),
-    rewriter.getDenseI64ArrayAttr(sourceOffsets), rewriter.getDenseI64ArrayAttr(strides),
-    rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete")).getOutput();
+FailureOr<Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
+                                                 PatternRewriter& rewriter,
+                                                 Location loc) {
+  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;
+      entries.push_back({row * value.tilesPerRow + tile, 0, {0, channelOffset, row, 0},
+        {1, std::min(tileChannels, value.logicalType.getDimSize(1) - channelOffset), 1,
+         value.logicalType.getDimSize(3)}});
+    }
+  return createFragmentAssemblyBlueprint(value.storage, value.logicalType, entries, "nchw_row_strip",
+                                         kRowStripIndexMap, rewriter, loc);
 }
 
-FailureOr<Value>
-applyRowStripRelu(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
-  auto fragmentType = getRowStripFragmentType(logicalType);
-  auto storageType = getRowStripStorageType(logicalType);
+FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {
+  auto storageType = cast<RankedTensorType>(value.storage.getType());
+  const int64_t laneCount = storageType.getDimSize(0);
   auto batchOp = createSpatComputeBatch(rewriter,
                                         loc,
                                         TypeRange {storageType},
-                                        logicalType.getDimSize(2),
+                                        laneCount,
                                         {},
-                                        ValueRange {storage},
+                                        ValueRange {value.storage},
                                         [&](detail::SpatComputeBatchBodyArgs args) {
-                                          Value fragment =
-                                            extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
-                                          fragment = spatial::SpatReluOp::create(rewriter, loc, fragmentType, fragment).getResult();
-                                          insertRowStripFragment(
-                                            fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+                                          FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
+                                            rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
+                                          if (failed(fragment)) return failure();
+                                          Value relu = spatial::SpatReluOp::create(
+                                            rewriter, loc, value.fragmentType, *fragment).getResult();
+                                          publishGraphBatchPhysicalFragment(
+                                            rewriter, loc, relu, args.outputs.front(), args.lane);
                                           return success();
                                         });
   if (failed(batchOp))
@@ -166,41 +180,47 @@ applyRowStripRelu(Value storage, RankedTensorType logicalType, PatternRewriter&
   return batchOp->getResult(0);
 }
 
-FailureOr<Value>
-applyRowStripBiasAdd(Value storage, RankedTensorType logicalType, Value bias, PatternRewriter& rewriter, Location loc) {
+FailureOr<Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
+                                     Value bias,
+                                     PatternRewriter& rewriter,
+                                     Location loc) {
   DenseElementsAttr denseAttr;
-  if (!isSupportedBiasAddValue(bias, logicalType, &denseAttr))
+  if (!isSupportedBiasAddValue(bias, value.logicalType, &denseAttr))
     return failure();
-  auto fragmentType = getRowStripFragmentType(logicalType);
-  auto storageType = getRowStripStorageType(logicalType);
+  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, value.logicalType);
+  if (failed(channelValues)) return failure();
+  auto storageType = cast<RankedTensorType>(value.storage.getType());
+  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);
+  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] =
+        (*channelValues)[channel];
+  Value biasStorage = getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(),
+    DenseElementsAttr::get(biasStorageType, biasValues), biasStorageType);
+  const int64_t laneCount = storageType.getDimSize(0);
   auto batchOp = createSpatComputeBatch(rewriter,
                                         loc,
                                         TypeRange {storageType},
-                                        logicalType.getDimSize(2),
+                                        laneCount,
                                         {},
-                                        ValueRange {storage},
+                                        ValueRange {value.storage, biasStorage},
                                         [&](detail::SpatComputeBatchBodyArgs args) {
-                                          Value fragment =
-                                            extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
-                                          Value constant;
-                                          if (denseAttr.isSplat()) {
-                                            constant = getOrCreateConstant(
-                                              rewriter,
-                                              rewriter.getInsertionBlock()->getParentOp(),
-                                              DenseElementsAttr::get(fragmentType, denseAttr.getSplatValue<Attribute>()),
-                                              fragmentType);
-                                          }
-                                          else {
-                                            FailureOr<Value> perChannel =
-                                              createPerChannelConstantFragment(denseAttr, fragmentType, rewriter);
-                                            if (failed(perChannel))
-                                              return failure();
-                                            constant = *perChannel;
-                                          }
-                                          fragment =
-                                            spatial::SpatVAddOp::create(rewriter, loc, fragmentType, fragment, constant).getResult();
-                                          insertRowStripFragment(
-                                            fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+                                          Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+                                          FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
+                                            rewriter, loc, args.inputs[0], args.lane, value.fragmentType);
+                                          Value tile = affineModConst(rewriter, loc, args.lane,
+                                            value.tilesPerRow, anchorOp);
+                                          FailureOr<Value> constant = extractGraphBatchPhysicalFragment(
+                                            rewriter, loc, args.inputs[1], tile, value.fragmentType);
+                                          if (failed(fragment) || failed(constant)) return failure();
+                                          Value added = spatial::SpatVAddOp::create(
+                                            rewriter, loc, value.fragmentType, *fragment, *constant).getResult();
+                                          publishGraphBatchPhysicalFragment(
+                                            rewriter, loc, added, args.outputs.front(), args.lane);
                                           return success();
                                         });
   if (failed(batchOp))

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

@@ -11,10 +11,13 @@ inline constexpr llvm::StringLiteral kRowStripIndexMap = "nchw_row_strip_fragmen
 struct RowStripPhysicalValue {
   mlir::Value storage;
   mlir::RankedTensorType logicalType;
-  llvm::SmallVector<int64_t, 16> fragmentOffsets;
-  llvm::SmallVector<int64_t, 16> fragmentSizes;
+  mlir::RankedTensorType fragmentType;
+  int64_t tilesPerRow;
 };
 
+mlir::FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(mlir::Value storage,
+                                                                     mlir::RankedTensorType logicalType);
+
 std::pair<llvm::SmallVector<int64_t>, llvm::SmallVector<int64_t>>
 buildRowStripMetadata(mlir::RankedTensorType type);
 
@@ -50,18 +53,15 @@ mlir::FailureOr<mlir::Value> createRowStripStorageFromRows(mlir::Value rows,
                                                            mlir::PatternRewriter& rewriter,
                                                            mlir::Location loc);
 
-mlir::FailureOr<mlir::Value> createRowStripAssemblyBlueprint(mlir::Value storage,
-                                                               mlir::RankedTensorType logicalType,
+mlir::FailureOr<mlir::Value> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
                                                                mlir::PatternRewriter& rewriter,
                                                                mlir::Location loc);
 
-mlir::FailureOr<mlir::Value> applyRowStripRelu(mlir::Value storage,
-                                               mlir::RankedTensorType logicalType,
+mlir::FailureOr<mlir::Value> applyRowStripRelu(const RowStripPhysicalValue& value,
                                                mlir::PatternRewriter& rewriter,
                                                mlir::Location loc);
 
-mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(mlir::Value storage,
-                                                  mlir::RankedTensorType logicalType,
+mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
                                                   mlir::Value bias,
                                                   mlir::PatternRewriter& rewriter,
                                                   mlir::Location loc);