big refactor

2026-08-04 11:28:05 +02:00
parent f4a3b012cc
commit 10b6ee6c32
150 changed files with 6737 additions and 4816 deletions
@@ -129,6 +129,32 @@ LogicalResult validateFragmentAssemblyMetadata(spatial::SpatBlueprintOp blueprin
  return success();
 }

+FailureOr<mlir::Value> reshapeContiguousRowMajorFragments(RewriterBase& rewriter,
+                                                          Location loc,
+                                                          mlir::Value source,
+                                                          RankedTensorType resultType) {
+  auto sourceType = dyn_cast<RankedTensorType>(source.getType());
+  if (!sourceType || !sourceType.hasStaticShape() || !resultType.hasStaticShape() || resultType.getRank() < 2
+      || sourceType.getRank() != resultType.getRank() + 1 || sourceType.getElementType() != resultType.getElementType()
+      || sourceType.getNumElements() != resultType.getNumElements()
+      || sourceType.getDimSize(0) != getStaticShapeElementCount(resultType.getShape().drop_back())
+      || sourceType.getDimSize(sourceType.getRank() - 1) != resultType.getDimSize(resultType.getRank() - 1)
+      || llvm::any_of(sourceType.getShape().slice(1, sourceType.getRank() - 2), [](int64_t dim) { return dim != 1; }))
+    return failure();
+
+  SmallVector<ReassociationIndices> collapse {{}, {sourceType.getRank() - 1}};
+  for (int64_t dim = 0; dim < sourceType.getRank() - 1; ++dim)
+    collapse.front().push_back(dim);
+  auto flatType = RankedTensorType::get(
+    {sourceType.getDimSize(0), sourceType.getDimSize(sourceType.getRank() - 1)}, resultType.getElementType());
+  mlir::Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, source, collapse);
+
+  SmallVector<ReassociationIndices> expand {{}, {resultType.getRank() - 1}};
+  for (int64_t dim = 0; dim < resultType.getRank() - 1; ++dim)
+    expand.front().push_back(dim);
+  return tensor::ExpandShapeOp::create(rewriter, loc, resultType, flat, expand).getResult();
+}
+
 static SmallVector<int64_t, 4> expandFlatElementIndex(int64_t flatIndex, ArrayRef<int64_t> shape) {
  SmallVector<int64_t, 4> indices(shape.size(), 0);
  for (int64_t dim = static_cast<int64_t>(shape.size()) - 1; dim >= 0; --dim) {