E' ancora tutto rotto

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -383,7 +383,7 @@ LogicalResult SpatConcatOp::verify() {
 static bool isKnownLogicalLayout(StringRef layout) { return layout == "nchw"; }
 
 static bool isKnownPhysicalLayout(StringRef layout) {
-  return layout == "dense_nchw" || layout == "nchw_row_strip";
+  return layout == "dense_nchw" || layout == "nchw_row_strip" || layout == "fragmented";
 }
 
 static LogicalResult verifyPlanTensorTypes(Operation* op, Value input, Value output, StringRef kind) {
@@ -437,7 +437,9 @@ LogicalResult SpatReluPlanOp::verify() {
 }
 
 LogicalResult SpatReconciliatorOp::verify() {
-  if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.reconciliator")))
+  auto modeAttr = getModeAttr();
+  bool isFragmentAssembly = modeAttr && modeAttr.getValue() == "fragment_assembly";
+  if (!isFragmentAssembly && failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.reconciliator")))
     return failure();
   if (!isKnownLogicalLayout(getLogicalLayout()))
     return emitError("requires a known logical layout");
@@ -452,23 +454,154 @@ LogicalResult SpatReconciliatorOp::verify() {
   auto sizes = getFragmentSizes();
   if (offsets.size() != sizes.size())
     return emitError("fragment offset and size arrays must have the same length");
+  int64_t rank = logicalType.getRank();
   if (offsets.empty())
     return success();
-
-  int64_t rank = logicalType.getRank();
   if (rank <= 0 || offsets.size() % rank != 0)
     return emitError("fragment metadata must be a whole number of rank-sized fragments");
 
-  ArrayRef<int64_t> shape = logicalType.getShape();
-  for (int64_t index = 0; index < static_cast<int64_t>(offsets.size()); ++index) {
-    int64_t dim = index % rank;
-    int64_t offset = offsets[index];
-    int64_t size = sizes[index];
-    if (offset < 0 || size < 0)
-      return emitError("fragment offsets and sizes must be non-negative");
-    int64_t logicalDim = shape[dim];
-    if (!ShapedType::isDynamic(logicalDim) && offset + size > logicalDim)
-      return emitError("fragment bounds must stay within the logical tensor shape");
+  auto verifyBoundsOnly = [&](ArrayRef<int64_t> strideValues) -> LogicalResult {
+    ArrayRef<int64_t> shape = logicalType.getShape();
+    for (int64_t index = 0; index < static_cast<int64_t>(offsets.size()); ++index) {
+      int64_t dim = index % rank;
+      int64_t offset = offsets[index];
+      int64_t size = sizes[index];
+      int64_t stride = strideValues.empty() ? 1 : strideValues[index];
+      if (offset < 0 || size < 0 || stride < 0)
+        return emitError("fragment offsets, sizes, and strides must be non-negative");
+      int64_t logicalDim = shape[dim];
+      if (!ShapedType::isDynamic(logicalDim) && offset + size > logicalDim)
+        return emitError("fragment bounds must stay within the logical tensor shape");
+      if (stride != 1)
+        return emitError("fragment assembly currently requires unit strides");
+    }
+    return success();
+  };
+
+  if (!isFragmentAssembly) {
+    if (failed(verifyBoundsOnly({})))
+      return failure();
+    if (!getFragments().empty())
+      return emitError("legacy reconciliator does not accept extra fragment operands");
+    if (getFragmentStridesAttr() || getConflictPolicyAttr() || getCoveragePolicyAttr())
+      return emitError("legacy reconciliator does not accept fragment assembly attributes");
+    return success();
+  }
+
+  auto stridesAttr = getFragmentStridesAttr();
+  auto operandIndicesAttr = getFragmentOperandIndicesAttr();
+  if (!operandIndicesAttr)
+    return emitError("fragment assembly reconciliator requires fragment operand indices");
+  if (!stridesAttr)
+    return emitError("fragment assembly reconciliator requires fragment strides");
+  ArrayRef<int64_t> operandIndices = operandIndicesAttr.asArrayRef();
+  ArrayRef<int64_t> strides = stridesAttr.asArrayRef();
+  if (strides.size() != offsets.size())
+    return emitError("fragment stride and offset arrays must have the same length");
+  if (!getConflictPolicyAttr() || !getCoveragePolicyAttr())
+    return emitError("fragment assembly reconciliator requires conflict and coverage policies");
+  if (getConflictPolicy() != "disjoint")
+    return emitError("fragment assembly reconciliator currently supports only conflict_policy=\"disjoint\"");
+  if (getCoveragePolicy() != "complete" && getCoveragePolicy() != "partial")
+    return emitError("fragment assembly reconciliator coverage_policy must be \"complete\" or \"partial\"");
+
+  SmallVector<Value> operands;
+  operands.push_back(getInput());
+  llvm::append_range(operands, getFragments());
+  int64_t operandCount = static_cast<int64_t>(operands.size());
+  int64_t fragmentCount = static_cast<int64_t>(operandIndices.size());
+  if (operandCount == 0)
+    return emitError("fragment assembly reconciliator requires at least one operand");
+  if (static_cast<int64_t>(offsets.size()) != fragmentCount * rank)
+    return emitError("fragment assembly metadata count must match operand count * result rank");
+  if (failed(verifyBoundsOnly(strides)))
+    return failure();
+
+  SmallVector<std::pair<SmallVector<int64_t, 4>, SmallVector<int64_t, 4>>, 8> slices;
+  slices.reserve(static_cast<size_t>(fragmentCount));
+  SmallVector<SmallVector<SmallVector<int64_t, 4>, 4>, 8> sizesByOperand(static_cast<size_t>(operandCount));
+  for (int64_t fragmentIndex = 0; fragmentIndex < fragmentCount; ++fragmentIndex) {
+    int64_t operandIndex = operandIndices[fragmentIndex];
+    if (operandIndex < 0 || operandIndex >= operandCount)
+      return emitError("fragment assembly operand index is out of range");
+
+    auto operandType = dyn_cast<RankedTensorType>(operands[operandIndex].getType());
+    if (!operandType || !operandType.hasStaticShape())
+      return emitError("fragment assembly reconciliator requires static ranked tensor operands");
+    if (operandType.getRank() != rank)
+      return emitError("fragment assembly reconciliator requires operand/result rank match");
+
+    SmallVector<int64_t, 4> fragmentOffsets;
+    SmallVector<int64_t, 4> fragmentSizes;
+    fragmentOffsets.reserve(rank);
+    fragmentSizes.reserve(rank);
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      int64_t flatIndex = fragmentIndex * rank + dim;
+      fragmentOffsets.push_back(offsets[flatIndex]);
+      fragmentSizes.push_back(sizes[flatIndex]);
+    }
+
+    sizesByOperand[static_cast<size_t>(operandIndex)].push_back(fragmentSizes);
+
+    for (const auto& [existingOffsets, existingSizes] : slices) {
+      bool overlaps = true;
+      for (int64_t dim = 0; dim < rank; ++dim) {
+        int64_t begin = fragmentOffsets[dim];
+        int64_t end = begin + fragmentSizes[dim];
+        int64_t existingBegin = existingOffsets[dim];
+        int64_t existingEnd = existingBegin + existingSizes[dim];
+        if (end <= existingBegin || existingEnd <= begin) {
+          overlaps = false;
+          break;
+        }
+      }
+      if (overlaps)
+        return emitError("fragment assembly reconciliator requires disjoint static slices");
+    }
+    slices.push_back({std::move(fragmentOffsets), std::move(fragmentSizes)});
+  }
+
+  for (int64_t operandIndex = 0; operandIndex < operandCount; ++operandIndex) {
+    if (sizesByOperand[static_cast<size_t>(operandIndex)].empty())
+      return emitError("fragment assembly reconciliator requires every operand to contribute at least one fragment");
+
+    auto operandType = cast<RankedTensorType>(operands[operandIndex].getType());
+    ArrayRef<int64_t> operandShape = operandType.getShape();
+    auto& fragmentShapes = sizesByOperand[static_cast<size_t>(operandIndex)];
+    if (fragmentShapes.size() == 1) {
+      if (!llvm::equal(operandShape, fragmentShapes.front()))
+        return emitError("single-fragment reconciliator operand shape must match declared fragment size");
+      continue;
+    }
+
+    ArrayRef<int64_t> fragmentShape = fragmentShapes.front();
+    for (ArrayRef<int64_t> otherShape : fragmentShapes)
+      if (!llvm::equal(fragmentShape, otherShape))
+        return emitError("packed reconciliator operand requires equal fragment sizes per operand");
+    if (llvm::equal(operandShape, fragmentShape))
+      continue;
+    if (!llvm::equal(operandShape.drop_front(), fragmentShape.drop_front()))
+      return emitError("packed reconciliator operand must match fragment shape on non-packed dimensions");
+    if (operandShape.front() != static_cast<int64_t>(fragmentShapes.size()) * fragmentShape.front())
+      return emitError("packed reconciliator operand first dimension must equal fragment_count * fragment_size");
+  }
+
+  if (getCoveragePolicy() == "complete") {
+    int64_t covered = 0;
+    int64_t logicalElements = 1;
+    for (int64_t dimSize : logicalType.getShape()) {
+      if (ShapedType::isDynamic(dimSize))
+        return emitError("fragment assembly complete coverage requires static result shape");
+      logicalElements *= dimSize;
+    }
+    for (const auto& [ignoredOffsets, fragmentSizes] : slices) {
+      int64_t fragmentElements = 1;
+      for (int64_t dimSize : fragmentSizes)
+        fragmentElements *= dimSize;
+      covered += fragmentElements;
+    }
+    if (covered != logicalElements)
+      return emitError("fragment assembly complete coverage must cover the whole result exactly");
   }
 
   return success();