temp commit: i will soft-reset and recommit after next changes

2026-08-02 11:37:22 +02:00
parent f4a3b012cc
commit 893e90feac
43 changed files with 2159 additions and 2048 deletions
@@ -11,6 +11,70 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace spatial {

+bool hasCanonicalContiguousRowMajorFragments(RankedTensorType logicalType,
+                                             ArrayRef<int64_t> offsets,
+                                             ArrayRef<int64_t> sizes,
+                                             ArrayRef<int64_t> strides) {
+  if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() <= 0
+      || logicalType.getDimSize(logicalType.getRank() - 1) <= 0)
+    return false;
+  const int64_t rank = logicalType.getRank();
+  const int64_t rowCount = logicalType.getNumElements() / logicalType.getDimSize(rank - 1);
+  if (offsets.size() != static_cast<size_t>(rowCount * rank) || sizes.size() != offsets.size()
+      || strides.size() != offsets.size())
+    return false;
+  for (int64_t row = 0; row < rowCount; ++row) {
+    int64_t remaining = row;
+    for (int64_t dim = rank - 2; dim >= 0; --dim) {
+      const int64_t index = row * rank + dim;
+      if (offsets[index] != remaining % logicalType.getDimSize(dim) || sizes[index] != 1 || strides[index] != 1)
+        return false;
+      remaining /= logicalType.getDimSize(dim);
+    }
+    const int64_t last = row * rank + rank - 1;
+    if (offsets[last] != 0 || sizes[last] != logicalType.getDimSize(rank - 1) || strides[last] != 1)
+      return false;
+  }
+  return true;
+}
+
+bool isCanonicalContiguousRowMajorFragmentAssembly(SpatBlueprintOp blueprint) {
+  auto logicalType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+  auto physicalType = dyn_cast<RankedTensorType>(blueprint.getInput().getType());
+  auto operandIndices = blueprint.getFragmentOperandIndices();
+  auto sourceSlots = blueprint.getFragmentSourceSlots();
+  auto sourceOffsets = blueprint.getFragmentSourceOffsets();
+  auto fragmentStrides = blueprint.getFragmentStrides();
+  if (!logicalType || !physicalType || !logicalType.hasStaticShape() || !physicalType.hasStaticShape()
+      || logicalType.getRank() < 2 || !blueprint.getFragments().empty()
+      || blueprint.getMode() != "fragment_assembly" || !operandIndices || !sourceSlots || !sourceOffsets
+      || !fragmentStrides)
+    return false;
+
+  ArrayRef<int64_t> offsets = blueprint.getFragmentOffsets();
+  ArrayRef<int64_t> sizes = blueprint.getFragmentSizes();
+  if (!hasCanonicalContiguousRowMajorFragments(logicalType, offsets, sizes, *fragmentStrides)
+      || operandIndices->empty() || operandIndices->size() != sourceSlots->size()
+      || operandIndices->size() != sourceOffsets->size()
+      || operandIndices->size() * static_cast<size_t>(logicalType.getRank()) != offsets.size()
+      || physicalType.getRank() != logicalType.getRank() + 1
+      || physicalType.getDimSize(0) != static_cast<int64_t>(operandIndices->size())
+      || physicalType.getDimSize(0)
+           != logicalType.getNumElements() / logicalType.getDimSize(logicalType.getRank() - 1)
+      || physicalType.getElementType() != logicalType.getElementType()
+      || physicalType.getNumElements() != logicalType.getNumElements()
+      || physicalType.getDimSize(physicalType.getRank() - 1) != logicalType.getDimSize(logicalType.getRank() - 1)
+      || llvm::any_of(physicalType.getShape().slice(1, physicalType.getRank() - 2),
+                     [](int64_t dim) { return dim != 1; }))
+    return false;
+
+  for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices))
+    if (operandIndex != 0 || (*sourceSlots)[fragmentIndex] != static_cast<int64_t>(fragmentIndex)
+        || (*sourceOffsets)[fragmentIndex] != 0)
+      return false;
+  return true;
+}
+
 RankedTensorType getGraphBatchPhysicalResultType(int64_t laneCount, RankedTensorType fragmentType) {
  SmallVector<int64_t> shape {laneCount};
  llvm::append_range(shape, fragmentType.getShape());